diff --git a/Cargo.lock b/Cargo.lock index ec2e562b8f2f..21251c8ef069 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3299,7 +3299,6 @@ dependencies = [ "hashbrown 0.14.3", "itertools 0.11.0", "once_cell", - "ordered-float", "revm-inspectors", "serde", "tempfile", @@ -5135,15 +5134,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" -[[package]] -name = "ordered-float" -version = "4.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a76df7075c7d4d01fdcb46c912dd17fba5b60c78ea480b475f2b6ab6f666584e" -dependencies = [ - "num-traits", -] - [[package]] name = "overload" version = "0.1.1" diff --git a/crates/evm/traces/Cargo.toml b/crates/evm/traces/Cargo.toml index 28d81e2fee7f..94dd36b63d28 100644 --- a/crates/evm/traces/Cargo.toml +++ b/crates/evm/traces/Cargo.toml @@ -29,7 +29,6 @@ hashbrown = "0.14" hex.workspace = true itertools.workspace = true once_cell = "1" -ordered-float = "4" serde = "1" tokio = { version = "1", features = ["time", "macros"] } tracing = "0.1" diff --git a/crates/evm/traces/src/identifier/local.rs b/crates/evm/traces/src/identifier/local.rs index d37552a78426..13f0c8d7ba96 100644 --- a/crates/evm/traces/src/identifier/local.rs +++ b/crates/evm/traces/src/identifier/local.rs @@ -3,19 +3,25 @@ use alloy_json_abi::JsonAbi; use alloy_primitives::Address; use foundry_common::contracts::{bytecode_diff_score, ContractsByArtifact}; use foundry_compilers::ArtifactId; -use ordered_float::OrderedFloat; use std::borrow::Cow; /// A trace identifier that tries to identify addresses using local contracts. pub struct LocalTraceIdentifier<'a> { known_contracts: &'a ContractsByArtifact, + + // Vector of pairs of artifact id and the code length of the given artifact. + // Stored in descending order of code length to optimize the search. + ordered_ids: Vec<(&'a ArtifactId, usize)>, } impl<'a> LocalTraceIdentifier<'a> { /// Creates a new local trace identifier. #[inline] pub fn new(known_contracts: &'a ContractsByArtifact) -> Self { - Self { known_contracts } + let mut ordered_ids = + known_contracts.iter().map(|(id, contract)| (id, contract.1.len())).collect::>(); + ordered_ids.sort_by_key(|(_, len)| *len); + Self { known_contracts, ordered_ids } } /// Returns the known contracts. @@ -24,21 +30,40 @@ impl<'a> LocalTraceIdentifier<'a> { self.known_contracts } - fn find_contract_from_bytecode( - &mut self, - code: &[u8], - ) -> Option<(&'a ArtifactId, &'a JsonAbi)> { - self.known_contracts - .iter() - .filter_map(|(id, (abi, known_code))| { - // Note: the diff score can be inaccurate for small contracts so we're using - // a relatively high threshold here to avoid filtering out too many - // contracts. - let score = bytecode_diff_score(known_code, code); - (score < 0.85).then_some((score, id, abi)) - }) - .min_by_key(|(score, _, _)| OrderedFloat(*score)) - .map(|(_, id, abi)| (id, abi)) + /// Iterates over artifacts with code length less than or equal to the given code and tries to + /// find a match. + /// + /// We do not consider artifacts with code length greater than the given code length as it is + /// considered that after compilation code can only be extended by additional parameters + /// (immutables) and cannot be shortened. + pub fn identify_code(&self, code: &[u8]) -> Option<(&'a ArtifactId, &'a JsonAbi)> { + let mut min_score = f64::MAX; + let mut min_score_id = None; + + let ids_start = match self + .ordered_ids + .binary_search_by(|(_, known_code_len)| known_code_len.cmp(&code.len())) + { + // Exact match. + Ok(i) => i, + // Not found, start searching from the previous index. + Err(i) => i.saturating_sub(1), + }; + for &(id, _) in &self.ordered_ids[ids_start..] { + let (abi, known_code) = self.known_contracts.get(id)?; + let score = bytecode_diff_score(known_code, code); + trace!(%score, abi=?abi.functions().collect::>()); + if score < 0.1 { + return Some((id, abi)); + } + if score < min_score { + min_score = score; + min_score_id = Some((id, abi)); + } + } + + trace!(%min_score, "no close-enough match found"); + min_score_id } } @@ -47,9 +72,16 @@ impl TraceIdentifier for LocalTraceIdentifier<'_> { where A: Iterator)>, { + trace!("identify {:?} addresses", addresses.size_hint().1); + addresses .filter_map(|(address, code)| { - let (id, abi) = self.find_contract_from_bytecode(code?)?; + let _span = trace_span!("identify", %address).entered(); + + trace!("identifying"); + let (id, abi) = self.identify_code(code?)?; + trace!(id=%id.identifier(), "identified"); + Some(AddressIdentity { address: *address, contract: Some(id.identifier()),