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(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]>) -> FixedBytes<4> { + let arg_types: Vec = 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::>()); + let err_selector = compute_selector(&err.name, err.args.iter().collect::>()); + + 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); + } +}