Compute 4 Byte selector from SolError and SolFunction

crates/ast/src/util.rs

@@ -1,4 +1,7 @@
-use alloy_primitives::U256;
+use crate::ast::{SolError, SolFunction};
+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 +56,79 @@ pub fn u256_as_push(value: U256) -> Opcode {
         _ => unreachable!(),
     }
 }
+
+type Selector = (FixedBytes<4>, FixedBytes<4>);
+
+pub fn compute_selector(
+    func: SolFunction,
+    err: SolError,
+) -> Option<(FixedBytes<4>, FixedBytes<4>)> {
+    let build_signature = |name: &Spanned<&str>, args: &Box<[Spanned<DynSolType>]>| -> Vec<u8> {
+        let arg_types: Vec<_> = args
+            .iter()
+            .map(|arg| match arg.0 {
+                DynSolType::Address => "address",
+                DynSolType::Uint(_) => "uint256",
+                DynSolType::String => "string",
+                _ => panic!("Unsupported type: {:?}", arg.0),
+            })
+            .collect();
+
+        let mut signature = String::new();
+        signature.push_str(name.0);
+        signature.push('(');
+        signature.push_str(&arg_types.join(","));
+        signature.push(')');
+        signature.into_bytes()
+    };
+
+    let build_selector =
+        |name: &Spanned<&str>, args: &Box<[Spanned<DynSolType>]>| -> FixedBytes<4> {
+            let signature = build_signature(name, args);
+            let hash = keccak256(signature);
+            FixedBytes::<4>::from_slice(&hash[..4])
+        };
+
+    let func_selector = build_selector(&func.name, &func.args);
+    let err_selector = build_selector(&err.name, &err.args);
+
+    Some((func_selector, err_selector))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ast::*;
+    use alloy_dyn_abi::DynSolType;
+    use alloy_primitives::keccak256;
+    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 selectors = compute_selector(func.clone(), err.clone()).unwrap();
+        let expected_func_hash = keccak256(selectors.0.clone());
+        let expected_err_hash = keccak256(selectors.1.clone());
+        //println!("{}", &selectors);
+        let expected_func_selector: Selector = (
+            FixedBytes::from_slice(&expected_func_hash[..4]),
+            FixedBytes::from_slice(&expected_err_hash[..4]),
+        );
+        assert_eq!(selectors.0, expected_func_selector.0);
+        assert_eq!(selectors.1, expected_func_selector.1);
+    }
+}