chore(cheatcodes): clean up cheatcodes util.rs with new wallet and parsing

crates/evm/src/executor/inspector/cheatcodes/ext.rs

@@ -1,12 +1,13 @@
-use super::{bail, ensure, fmt_err, Cheatcodes, Result};
-use crate::{abi::HEVMCalls, executor::inspector::cheatcodes::util};
+use super::{bail, ensure, fmt_err, util::MAGIC_SKIP_BYTES, Cheatcodes, Error, Result};
+use crate::{abi::HEVMCalls, executor::inspector::cheatcodes::parse};
 use ethers::{
     abi::{self, AbiEncode, JsonAbi, ParamType, Token},
     prelude::artifacts::CompactContractBytecode,
     types::*,
 };
 use foundry_common::{fmt::*, fs, get_artifact_path};
 use foundry_config::fs_permissions::FsAccessKind;
+use revm::{Database, EVMData};
 use serde::Deserialize;
 use serde_json::Value;
 use std::{collections::BTreeMap, env, path::Path, process::Command};
@@ -190,9 +191,9 @@ fn get_env(key: &str, ty: ParamType, delim: Option<&str>, default: Option<String
         })
     })?;
     if let Some(d) = delim {
-        util::parse_array(val.split(d).map(str::trim), &ty)
+        parse::parse_array(val.split(d).map(str::trim), &ty)
     } else {
-        util::parse(&val, &ty)
+        parse::parse(&val, &ty)
     }
 }
 
@@ -342,9 +343,9 @@ fn parse_json(json_str: &str, key: &str, coerce: Option<ParamType>) -> Result {
                 };
                 trace!(target : "forge::evm", ?values, "parsign values");
                 return if let Some(array) = values[0].as_array() {
-                    util::parse_array(array.iter().map(to_string), &coercion_type)
+                    parse::parse_array(array.iter().map(to_string), &coercion_type)
                 } else {
-                    util::parse(&to_string(values[0]), &coercion_type)
+                    parse::parse(&to_string(values[0]), &coercion_type)
                 }
             }
 
@@ -482,7 +483,7 @@ fn key_exists(json_str: &str, key: &str) -> Result {
     let json: Value =
         serde_json::from_str(json_str).map_err(|e| format!("Could not convert to JSON: {e}"))?;
     let values = jsonpath_lib::select(&json, &canonicalize_json_key(key))?;
-    let exists = util::parse(&(!values.is_empty()).to_string(), &ParamType::Bool)?;
+    let exists = parse::parse(&(!values.is_empty()).to_string(), &ParamType::Bool)?;
     Ok(exists)
 }
 
@@ -494,8 +495,28 @@ fn sleep(milliseconds: &U256) -> Result {
     Ok(Default::default())
 }
 
+/// Skip the current test, by returning a magic value that will be checked by the test runner.
+pub fn skip(state: &mut Cheatcodes, depth: u64, skip: bool) -> Result {
+    if !skip {
+        return Ok(b"".into())
+    }
+
+    // Skip should not work if called deeper than at test level.
+    // As we're not returning the magic skip bytes, this will cause a test failure.
+    if depth > 1 {
+        return Err(Error::custom("The skip cheatcode can only be used at test level"))
+    }
+
+    state.skip = true;
+    Err(Error::custom_bytes(MAGIC_SKIP_BYTES))
+}
+
 #[instrument(level = "error", name = "ext", target = "evm::cheatcodes", skip_all)]
-pub fn apply(state: &mut Cheatcodes, call: &HEVMCalls) -> Option<Result> {
+pub fn apply<DB: Database>(
+    state: &mut Cheatcodes,
+    data: &mut EVMData<'_, DB>,
+    call: &HEVMCalls,
+) -> Option<Result> {
     Some(match call {
         HEVMCalls::Ffi(inner) => {
             if state.config.ffi {
@@ -680,6 +701,7 @@ pub fn apply(state: &mut Cheatcodes, call: &HEVMCalls) -> Option<Result> {
         HEVMCalls::WriteJson0(inner) => write_json(state, &inner.0, &inner.1, None),
         HEVMCalls::WriteJson1(inner) => write_json(state, &inner.0, &inner.1, Some(&inner.2)),
         HEVMCalls::KeyExists(inner) => key_exists(&inner.0, &inner.1),
+        HEVMCalls::Skip(inner) => skip(state, data.journaled_state.depth(), inner.0),
         _ => return None,
     })
 }

crates/evm/src/executor/inspector/cheatcodes/mod.rs

@@ -57,10 +57,15 @@ mod fs;
 mod fuzz;
 /// Mapping related cheatcodes
 mod mapping;
+/// Parsing related cheatcodes.
+/// Does not include JSON-related cheatcodes to cut complexity.
+mod parse;
 /// Snapshot related cheatcodes
 mod snapshot;
-/// Utility cheatcodes (`sign` etc.)
+/// Utility functions and constants.
 pub mod util;
+/// Wallet / key management related cheatcodes
+mod wallet;
 pub use util::{BroadcastableTransaction, DEFAULT_CREATE2_DEPLOYER};
 
 mod config;
@@ -219,10 +224,11 @@ impl Cheatcodes {
 
         let opt = env::apply(self, data, caller, &decoded)
             .transpose()
-            .or_else(|| util::apply(self, data, &decoded))
+            .or_else(|| wallet::apply(self, data, &decoded))
+            .or_else(|| parse::apply(self, data, &decoded))
             .or_else(|| expect::apply(self, data, &decoded))
             .or_else(|| fuzz::apply(&decoded))
-            .or_else(|| ext::apply(self, &decoded))
+            .or_else(|| ext::apply(self, data, &decoded))
             .or_else(|| fs::apply(self, &decoded))
             .or_else(|| snapshot::apply(data, &decoded))
             .or_else(|| fork::apply(self, data, &decoded));

crates/evm/src/executor/inspector/cheatcodes/parse.rs

@@ -0,0 +1,151 @@
+use super::{fmt_err, Cheatcodes, Result};
+use crate::abi::HEVMCalls;
+use ethers::{
+    abi::{ParamType, Token},
+    prelude::*,
+};
+use foundry_macros::UIfmt;
+use revm::{Database, EVMData};
+
+pub fn parse(s: &str, ty: &ParamType) -> Result {
+    parse_token(s, ty)
+        .map(|token| abi::encode(&[token]).into())
+        .map_err(|e| fmt_err!("Failed to parse `{s}` as type `{ty}`: {e}"))
+}
+
+pub fn parse_array<I, T>(values: I, ty: &ParamType) -> Result
+where
+    I: IntoIterator<Item = T>,
+    T: AsRef<str>,
+{
+    let mut values = values.into_iter();
+    match values.next() {
+        Some(first) if !first.as_ref().is_empty() => {
+            let tokens = std::iter::once(first)
+                .chain(values)
+                .map(|v| parse_token(v.as_ref(), ty))
+                .collect::<Result<Vec<_>, _>>()?;
+            Ok(abi::encode(&[Token::Array(tokens)]).into())
+        }
+        // return the empty encoded Bytes when values is empty or the first element is empty
+        _ => Ok(abi::encode(&[Token::String(String::new())]).into()),
+    }
+}
+
+fn parse_token(s: &str, ty: &ParamType) -> Result<Token, String> {
+    match ty {
+        ParamType::Bool => {
+            s.to_ascii_lowercase().parse().map(Token::Bool).map_err(|e| e.to_string())
+        }
+        ParamType::Uint(256) => parse_uint(s).map(Token::Uint),
+        ParamType::Int(256) => parse_int(s).map(Token::Int),
+        ParamType::Address => s.parse().map(Token::Address).map_err(|e| e.to_string()),
+        ParamType::FixedBytes(32) => parse_bytes(s).map(Token::FixedBytes),
+        ParamType::Bytes => parse_bytes(s).map(Token::Bytes),
+        ParamType::String => Ok(Token::String(s.to_string())),
+        _ => Err("unsupported type".into()),
+    }
+}
+
+fn parse_int(s: &str) -> Result<U256, String> {
+    // hex string may start with "0x", "+0x", or "-0x" which needs to be stripped for
+    // `I256::from_hex_str`
+    if s.starts_with("0x") || s.starts_with("+0x") || s.starts_with("-0x") {
+        s.replacen("0x", "", 1).parse::<I256>().map_err(|err| err.to_string())
+    } else {
+        match I256::from_dec_str(s) {
+            Ok(val) => Ok(val),
+            Err(dec_err) => s.parse::<I256>().map_err(|hex_err| {
+                format!("could not parse value as decimal or hex: {dec_err}, {hex_err}")
+            }),
+        }
+    }
+    .map(|v| v.into_raw())
+}
+
+fn parse_uint(s: &str) -> Result<U256, String> {
+    if s.starts_with("0x") {
+        s.parse::<U256>().map_err(|err| err.to_string())
+    } else {
+        match U256::from_dec_str(s) {
+            Ok(val) => Ok(val),
+            Err(dec_err) => s.parse::<U256>().map_err(|hex_err| {
+                format!("could not parse value as decimal or hex: {dec_err}, {hex_err}")
+            }),
+        }
+    }
+}
+
+fn parse_bytes(s: &str) -> Result<Vec<u8>, String> {
+    hex::decode(s).map_err(|e| e.to_string())
+}
+
+#[instrument(level = "error", name = "util", target = "evm::cheatcodes", skip_all)]
+pub fn apply<DB: Database>(
+    _state: &mut Cheatcodes,
+    _data: &mut EVMData<'_, DB>,
+    call: &HEVMCalls,
+) -> Option<Result> {
+    Some(match call {
+        HEVMCalls::ToString0(inner) => {
+            Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into())
+        }
+        HEVMCalls::ToString1(inner) => {
+            Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into())
+        }
+        HEVMCalls::ToString2(inner) => {
+            Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into())
+        }
+        HEVMCalls::ToString3(inner) => {
+            Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into())
+        }
+        HEVMCalls::ToString4(inner) => {
+            Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into())
+        }
+        HEVMCalls::ToString5(inner) => {
+            Ok(ethers::abi::encode(&[Token::String(inner.0.pretty())]).into())
+        }
+        HEVMCalls::ParseBytes(inner) => parse(&inner.0, &ParamType::Bytes),
+        HEVMCalls::ParseAddress(inner) => parse(&inner.0, &ParamType::Address),
+        HEVMCalls::ParseUint(inner) => parse(&inner.0, &ParamType::Uint(256)),
+        HEVMCalls::ParseInt(inner) => parse(&inner.0, &ParamType::Int(256)),
+        HEVMCalls::ParseBytes32(inner) => parse(&inner.0, &ParamType::FixedBytes(32)),
+        HEVMCalls::ParseBool(inner) => parse(&inner.0, &ParamType::Bool),
+        _ => return None,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ethers::abi::AbiDecode;
+
+    #[test]
+    fn test_uint_env() {
+        let pk = "0x10532cc9d0d992825c3f709c62c969748e317a549634fb2a9fa949326022e81f";
+        let val: U256 = pk.parse().unwrap();
+        let parsed = parse(pk, &ParamType::Uint(256)).unwrap();
+        let decoded = U256::decode(&parsed).unwrap();
+        assert_eq!(val, decoded);
+
+        let parsed = parse(pk.strip_prefix("0x").unwrap(), &ParamType::Uint(256)).unwrap();
+        let decoded = U256::decode(&parsed).unwrap();
+        assert_eq!(val, decoded);
+
+        let parsed = parse("1337", &ParamType::Uint(256)).unwrap();
+        let decoded = U256::decode(&parsed).unwrap();
+        assert_eq!(U256::from(1337u64), decoded);
+    }
+
+    #[test]
+    fn test_int_env() {
+        let val = U256::from(100u64);
+        let parsed = parse(&val.to_string(), &ParamType::Int(256)).unwrap();
+        let decoded = I256::decode(parsed).unwrap();
+        assert_eq!(val, decoded.try_into().unwrap());
+
+        let parsed = parse("100", &ParamType::Int(256)).unwrap();
+        let decoded = I256::decode(parsed).unwrap();
+        assert_eq!(U256::from(100u64), decoded.try_into().unwrap());
+    }
+}