From 679a43605a4d9bc8f8efbd137a2ccf7641518103 Mon Sep 17 00:00:00 2001
From: William Conyea <wconyea@crestviewglobal.com>
Date: Mon, 18 Nov 2024 08:15:30 -0500
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A7=20Add=20compute=20selector=20helpe?=
 =?UTF-8?q?r?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* wip: poc compute selectors

* refactor: compute selectors  #6

* fix

* remove

* rm unused

* rm unused

* fix: redudant closure / put test back

* clippy

* accidentally broke the test, whoops
---
 crates/ast/src/lib.rs  |  2 +-
 crates/ast/src/util.rs | 55 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/crates/ast/src/lib.rs b/crates/ast/src/lib.rs
index 2696f3e..b7273c7 100644
--- a/crates/ast/src/lib.rs
+++ b/crates/ast/src/lib.rs
@@ -5,4 +5,4 @@ mod util;
 
 pub use ast::*;
 pub use parser::parse;
-pub use util::u256_as_push;
+pub use util::*;
diff --git a/crates/ast/src/util.rs b/crates/ast/src/util.rs
index 126d98c..38d91f9 100644
--- a/crates/ast/src/util.rs
+++ b/crates/ast/src/util.rs
@@ -1,4 +1,6 @@
-use alloy_primitives::U256;
+use crate::Spanned;
+use alloy_dyn_abi::DynSolType;
+use alloy_primitives::{keccak256, FixedBytes, U256};
 use evm_glue::opcodes::Opcode;
 
 pub(crate) fn u256_as_push_data<const N: usize>(value: U256) -> Result<[u8; N], String> {
@@ -53,3 +55,54 @@ pub fn u256_as_push(value: U256) -> Opcode {
         _ => unreachable!(),
     }
 }
+
+pub fn compute_selector(name: &Spanned<&str>, args: Box<[&Spanned<DynSolType>]>) -> FixedBytes<4> {
+    let arg_types: Vec<String> = args.iter().map(|arg| arg.0.to_string()).collect();
+
+    let signature = format!("{}({})", name.0, arg_types.join(","));
+    let hash = keccak256(signature.as_bytes());
+    FixedBytes::<4>::from_slice(&hash[..4])
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ast::{SolError, SolFunction};
+    use chumsky::span::Span;
+
+    #[test]
+
+    fn test_compute_selector() {
+        let func = SolFunction {
+            name: Spanned::new("transfer", 0..8),
+            args: Box::new([
+                Spanned::new(DynSolType::Address, 9..17),
+                Spanned::new(DynSolType::Uint(256), 18..26),
+            ]),
+            rets: Box::new([]),
+        };
+
+        let err = SolError {
+            name: Spanned::new("TransferFailed", 0..15),
+            args: Box::new([
+                Spanned::new(DynSolType::String, 16..21),
+                Spanned::new(DynSolType::Uint(256), 22..30),
+            ]),
+        };
+
+        let func_selector = compute_selector(&func.name, func.args.iter().collect::<Box<[_]>>());
+        let err_selector = compute_selector(&err.name, err.args.iter().collect::<Box<[_]>>());
+
+        let expected_func_signature = "transfer(address,uint256)";
+        let expected_err_signature = "TransferFailed(string,uint256)";
+
+        let expected_func_hash = keccak256(expected_func_signature.as_bytes());
+        let expected_err_hash = keccak256(expected_err_signature.as_bytes());
+
+        let expected_func_selector = FixedBytes::<4>::from_slice(&expected_func_hash[..4]);
+        let expected_err_selector = FixedBytes::<4>::from_slice(&expected_err_hash[..4]);
+
+        assert_eq!(func_selector, expected_func_selector);
+        assert_eq!(err_selector, expected_err_selector);
+    }
+}