From 0f7960f5571d5f266fb8165635052f5c9c9d1c61 Mon Sep 17 00:00:00 2001
From: Daniel Perez <daniel@perez.sh>
Date: Fri, 20 Sep 2024 01:07:50 +0100
Subject: [PATCH] Add abi.decodeData

---
 CHANGELOG.md                    |  1 +
 docs/src/builtin_values.md      |  9 +++++++
 src/interpreter/builtins/abi.rs | 48 ++++++++++++++++++++++++++-------
 src/interpreter/builtins/mod.rs |  1 +
 4 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7226332..22722f4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@
 * Add `repl.fork` to start run and use an Anvil instance as a fork of the current URL
 * Add `repl.startPrank` / `repl.stopPrank` to start/stop impersonating an address
 * Add `FUNC.trace_call` method to contract functions
+* Add `abi.decodeData` to decode function calldata and errors from any known ABI
 
 ### Other changes
 
diff --git a/docs/src/builtin_values.md b/docs/src/builtin_values.md
index 31dea6c..7ddce72 100644
--- a/docs/src/builtin_values.md
+++ b/docs/src/builtin_values.md
@@ -256,6 +256,15 @@ Behaves like the regular Solidity `abi.decode` function.
 (1, 0x789f8F7B547183Ab8E99A5e0E6D567E90e0EB03B)
 ```
 
+### `abi.decodeData(bytes data) -> any`
+
+Decodes the data (either function calldata or error data) using any registered ABI.
+
+```javascript
+>> abi.decodeData(0xa9059cbb000000000000000000000000789f8f7b547183ab8e99a5e0e6d567e90e0eb03b0000000000000000000000000000000000000000000000000de0b6b3a7640000)
+("transfer(address,uint256)", (0x789f8F7B547183Ab8E99A5e0E6D567E90e0EB03B, 1000000000000000000))
+```
+
 ## `block` functions
 
 ### `block.number -> uint256`
diff --git a/src/interpreter/builtins/abi.rs b/src/interpreter/builtins/abi.rs
index 9c692a2..d8a2e8c 100644
--- a/src/interpreter/builtins/abi.rs
+++ b/src/interpreter/builtins/abi.rs
@@ -36,6 +36,17 @@ impl Decodable for json_abi::Error {
     }
 }
 
+fn _run_decode(signature: String, decoded: Vec<DynSolValue>) -> Result<Value> {
+    let values = decoded
+        .into_iter()
+        .map(Value::try_from)
+        .collect::<Result<Vec<_>>>()?;
+    Ok(Value::Tuple(vec![
+        Value::Str(signature),
+        Value::Tuple(values),
+    ]))
+}
+
 fn _generic_abi_decode<D: Decodable, F>(
     receiver: &Value,
     args: &[Value],
@@ -55,7 +66,7 @@ where
     };
     let selector = alloy::primitives::FixedBytes::<4>::from_slice(&data[..4]);
     let options = get_options(abi);
-    let error = options
+    let decodable = options
         .iter()
         .find(|f| f.selector() == selector)
         .ok_or(anyhow!(
@@ -64,15 +75,27 @@ where
             selector,
             name
         ))?;
-    let decoded = error.abi_decode_input(&data[4..], true)?;
-    let values = decoded
-        .into_iter()
-        .map(Value::try_from)
-        .collect::<Result<Vec<_>>>()?;
-    Ok(Value::Tuple(vec![
-        Value::Str(error.signature()),
-        Value::Tuple(values),
-    ]))
+    let decoded = decodable.abi_decode_input(&data[4..], true)?;
+    _run_decode(decodable.signature(), decoded)
+}
+
+fn abi_decode_data(env: &mut Env, _receiver: &Value, args: &[Value]) -> Result<Value> {
+    let data = match args.first() {
+        Some(Value::Bytes(bytes)) => bytes,
+        _ => bail!("abi.decodeData expects bytes as argument"),
+    };
+    if data.len() < 4 {
+        bail!("abi.decodeData expects at least 4 bytes");
+    }
+    let selector = alloy::primitives::FixedBytes::<4>::from_slice(&data[..4]);
+    let (signature, decoded) = if let Some(func) = env.get_function(&selector) {
+        (func.signature(), func.abi_decode_input(&data[4..], true)?)
+    } else if let Some(error) = env.get_error(&selector) {
+        (error.signature(), error.abi_decode_input(&data[4..], true)?)
+    } else {
+        bail!("function or error with selector {} not found", selector);
+    };
+    _run_decode(signature, decoded)
 }
 
 fn abi_decode_calldata(_env: &mut Env, receiver: &Value, args: &[Value]) -> Result<Value> {
@@ -142,6 +165,11 @@ lazy_static! {
         SyncMethod::arc("encodePacked", abi_encode_packed_, vec![]);
     pub static ref ABI_DECODE: Arc<dyn FunctionDef> =
         SyncMethod::arc("decode", abi_decode_, vec![]);
+    pub static ref ABI_DECODE_DATA: Arc<dyn FunctionDef> = SyncMethod::arc(
+        "decodeData",
+        abi_decode_data,
+        vec![vec![FunctionParam::new("data", Type::Bytes)]]
+    );
     pub static ref ABI_DECODE_CALLDATA: Arc<dyn FunctionDef> = SyncMethod::arc(
         "decode",
         abi_decode_calldata,
diff --git a/src/interpreter/builtins/mod.rs b/src/interpreter/builtins/mod.rs
index c365732..36d7222 100644
--- a/src/interpreter/builtins/mod.rs
+++ b/src/interpreter/builtins/mod.rs
@@ -124,6 +124,7 @@ lazy_static! {
         abi_methods.insert("encode".to_string(), abi::ABI_ENCODE.clone());
         abi_methods.insert("encodePacked".to_string(), abi::ABI_ENCODE_PACKED.clone());
         abi_methods.insert("decode".to_string(), abi::ABI_DECODE.clone());
+        abi_methods.insert("decodeData".to_string(), abi::ABI_DECODE_DATA.clone());
         m.insert(NonParametricType::Abi, abi_methods);
 
         let mut block_methods = HashMap::new();