Skip to content

Commit 1043551

Browse files
DaniPopesDaniPopes
committed
refactor: optimize trace identifiers
1 parent efedf1f commit 1043551

File tree

11 files changed

+275
-232
lines changed

11 files changed

+275
-232
lines changed

‎crates/chisel/src/dispatcher.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -958,13 +958,13 @@ impl ChiselDispatcher {
958958
session_config.evm_opts.get_remote_chain_id(),
959959
)?;
960960

961-
let mut decoder =
962-
CallTraceDecoderBuilder::new().with_labels(result.labeled_addresses.clone()).build();
963-
964-
decoder.add_signature_identifier(SignaturesIdentifier::new(
965-
Config::foundry_cache_dir(),
966-
session_config.foundry_config.offline,
967-
)?);
961+
let mut decoder = CallTraceDecoderBuilder::new()
962+
.with_labels(result.labeled_addresses.iter().map(|(a, s)| (*a, s.clone())))
963+
.with_signature_identifier(SignaturesIdentifier::new(
964+
Config::foundry_cache_dir(),
965+
session_config.foundry_config.offline,
966+
)?)
967+
.build();
968968

969969
for (_, trace) in &mut result.traces {
970970
// decoder.identify(trace, &mut local_identifier);

‎crates/cli/src/utils/cmd.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -378,12 +378,13 @@ pub async fn handle_traces(
378378
None
379379
});
380380

381-
let mut decoder = CallTraceDecoderBuilder::new().with_labels(labeled_addresses).build();
382-
383-
decoder.add_signature_identifier(SignaturesIdentifier::new(
384-
Config::foundry_cache_dir(),
385-
config.offline,
386-
)?);
381+
let mut decoder = CallTraceDecoderBuilder::new()
382+
.with_labels(labeled_addresses)
383+
.with_signature_identifier(SignaturesIdentifier::new(
384+
Config::foundry_cache_dir(),
385+
config.offline,
386+
)?)
387+
.build();
387388

388389
for (_, trace) in &mut result.traces {
389390
decoder.identify(trace, &mut etherscan_identifier);

‎crates/evm/src/trace/decoder.rs‎

Lines changed: 46 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use super::{
2-
identifier::{SingleSignaturesIdentifier, TraceIdentifier},
2+
identifier::{AddressIdentity, SingleSignaturesIdentifier, TraceIdentifier},
33
CallTraceArena, RawOrDecodedCall, RawOrDecodedLog, RawOrDecodedReturnData,
44
};
55
use crate::{
@@ -20,22 +20,27 @@ use std::collections::{BTreeMap, HashMap};
2020

2121
/// Build a new [CallTraceDecoder].
2222
#[derive(Default)]
23+
#[must_use = "builders do nothing unless you call `build` on them"]
2324
pub struct CallTraceDecoderBuilder {
2425
decoder: CallTraceDecoder,
2526
}
2627

2728
impl CallTraceDecoderBuilder {
29+
/// Create a new builder.
30+
#[inline]
2831
pub fn new() -> Self {
2932
Self { decoder: CallTraceDecoder::new().clone() }
3033
}
3134

3235
/// Add known labels to the decoder.
36+
#[inline]
3337
pub fn with_labels(mut self, labels: impl IntoIterator<Item = (Address, String)>) -> Self {
3438
self.decoder.labels.extend(labels);
3539
self
3640
}
3741

3842
/// Add known events to the decoder.
43+
#[inline]
3944
pub fn with_events(mut self, events: impl IntoIterator<Item = Event>) -> Self {
4045
for event in events {
4146
self.decoder
@@ -48,12 +53,21 @@ impl CallTraceDecoderBuilder {
4853
}
4954

5055
/// Sets the verbosity level of the decoder.
56+
#[inline]
5157
pub fn with_verbosity(mut self, level: u8) -> Self {
5258
self.decoder.verbosity = level;
5359
self
5460
}
5561

62+
/// Sets the signature identifier for events and functions.
63+
#[inline]
64+
pub fn with_signature_identifier(mut self, identifier: SingleSignaturesIdentifier) -> Self {
65+
self.decoder.signature_identifier = Some(identifier);
66+
self
67+
}
68+
5669
/// Build the decoder.
70+
#[inline]
5771
pub fn build(self) -> CallTraceDecoder {
5872
self.decoder
5973
}
@@ -168,23 +182,26 @@ impl CallTraceDecoder {
168182
}
169183
}
170184

171-
pub fn add_signature_identifier(&mut self, identifier: SingleSignaturesIdentifier) {
172-
self.signature_identifier = Some(identifier);
173-
}
174-
175185
/// Identify unknown addresses in the specified call trace using the specified identifier.
176186
///
177187
/// Unknown contracts are contracts that either lack a label or an ABI.
188+
#[inline]
178189
pub fn identify(&mut self, trace: &CallTraceArena, identifier: &mut impl TraceIdentifier) {
179-
let unidentified_addresses = trace
180-
.addresses()
181-
.into_iter()
182-
.filter(|(address, _)| {
183-
!self.labels.contains_key(address) || !self.contracts.contains_key(address)
184-
})
185-
.collect();
186-
187-
identifier.identify_addresses(unidentified_addresses).iter().for_each(|identity| {
190+
self.collect_identities(identifier.identify_addresses(self.addresses(trace)));
191+
}
192+
193+
#[inline(always)]
194+
fn addresses<'a>(
195+
&'a self,
196+
trace: &'a CallTraceArena,
197+
) -> impl Iterator<Item = (&'a Address, Option<&'a [u8]>)> + 'a {
198+
trace.addresses().into_iter().filter(|&(address, _)| {
199+
!self.labels.contains_key(address) || !self.contracts.contains_key(address)
200+
})
201+
}
202+
203+
fn collect_identities(&mut self, identities: Vec<AddressIdentity>) {
204+
for identity in identities {
188205
let address = identity.address;
189206

190207
if let Some(contract) = &identity.contract {
@@ -197,30 +214,31 @@ impl CallTraceDecoder {
197214

198215
if let Some(abi) = &identity.abi {
199216
// Store known functions for the address
200-
abi.functions()
201-
.map(|func| (func.short_signature(), func.clone()))
202-
.for_each(|(sig, func)| self.functions.entry(sig).or_default().push(func));
217+
for function in abi.functions() {
218+
self.functions
219+
.entry(function.short_signature())
220+
.or_default()
221+
.push(function.clone())
222+
}
203223

204224
// Flatten events from all ABIs
205-
abi.events()
206-
.map(|event| ((event.signature(), indexed_inputs(event)), event.clone()))
207-
.for_each(|(sig, event)| {
208-
self.events.entry(sig).or_default().push(event);
209-
});
225+
for event in abi.events() {
226+
let sig = (event.signature(), indexed_inputs(event));
227+
self.events.entry(sig).or_default().push(event.clone());
228+
}
210229

211230
// Flatten errors from all ABIs
212-
abi.errors().for_each(|error| {
213-
let entry = self.errors.errors.entry(error.name.clone()).or_default();
214-
entry.push(error.clone());
215-
});
231+
for error in abi.errors() {
232+
self.errors.errors.entry(error.name.clone()).or_default().push(error.clone());
233+
}
216234

217235
self.receive_contracts.entry(address).or_insert(abi.receive);
218236
}
219-
});
237+
}
220238
}
221239

222240
pub async fn decode(&self, traces: &mut CallTraceArena) {
223-
for node in traces.arena.iter_mut() {
241+
for node in &mut traces.arena {
224242
// Set contract name
225243
if let Some(contract) = self.contracts.get(&node.trace.address).cloned() {
226244
node.trace.contract = Some(contract);

‎crates/evm/src/trace/identifier/etherscan.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -100,11 +100,11 @@ impl EtherscanIdentifier {
100100
}
101101

102102
impl TraceIdentifier for EtherscanIdentifier {
103-
fn identify_addresses(
104-
&mut self,
105-
addresses: Vec<(&Address, Option<&[u8]>)>,
106-
) -> Vec<AddressIdentity> {
107-
trace!(target: "etherscanidentifier", "identify {} addresses", addresses.len());
103+
fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity>
104+
where
105+
A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>,
106+
{
107+
trace!(target: "etherscanidentifier", "identify {:?} addresses", addresses.size_hint().1);
108108

109109
let Some(client) = self.client.clone() else {
110110
// no client was configured

‎crates/evm/src/trace/identifier/local.rs‎

Lines changed: 19 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,56 +1,45 @@
11
use super::{AddressIdentity, TraceIdentifier};
2-
use ethers::{
3-
abi::{Abi, Address, Event},
4-
prelude::ArtifactId,
5-
};
2+
use ethers::abi::{Address, Event};
63
use foundry_common::contracts::{diff_score, ContractsByArtifact};
7-
use itertools::Itertools;
84
use ordered_float::OrderedFloat;
9-
use std::{borrow::Cow, collections::BTreeMap};
5+
use std::borrow::Cow;
106

117
/// A trace identifier that tries to identify addresses using local contracts.
12-
pub struct LocalTraceIdentifier {
13-
local_contracts: BTreeMap<Vec<u8>, (ArtifactId, Abi)>,
8+
pub struct LocalTraceIdentifier<'a> {
9+
known_contracts: &'a ContractsByArtifact,
1410
}
1511

16-
impl LocalTraceIdentifier {
17-
pub fn new(known_contracts: &ContractsByArtifact) -> Self {
18-
Self {
19-
local_contracts: known_contracts
20-
.iter()
21-
.map(|(id, (abi, runtime_code))| (runtime_code.clone(), (id.clone(), abi.clone())))
22-
.collect(),
23-
}
12+
impl<'a> LocalTraceIdentifier<'a> {
13+
pub fn new(known_contracts: &'a ContractsByArtifact) -> Self {
14+
Self { known_contracts }
2415
}
2516

2617
/// Get all the events of the local contracts.
2718
pub fn events(&self) -> impl Iterator<Item = &Event> {
28-
self.local_contracts.iter().flat_map(|(_, (_, abi))| abi.events())
19+
self.known_contracts.iter().flat_map(|(_, (abi, _))| abi.events())
2920
}
3021
}
3122

32-
impl TraceIdentifier for LocalTraceIdentifier {
33-
fn identify_addresses(
34-
&mut self,
35-
addresses: Vec<(&Address, Option<&[u8]>)>,
36-
) -> Vec<AddressIdentity> {
23+
impl TraceIdentifier for LocalTraceIdentifier<'_> {
24+
fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity>
25+
where
26+
A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>,
27+
{
3728
addresses
38-
.into_iter()
3929
.filter_map(|(address, code)| {
4030
let code = code?;
41-
let (_, (_, (id, abi))) = self
42-
.local_contracts
31+
let (_, id, abi) = self
32+
.known_contracts
4333
.iter()
44-
.filter_map(|entry| {
45-
let score = diff_score(entry.0, code);
34+
.filter_map(|(id, (abi, known_code))| {
35+
let score = diff_score(known_code, code);
4636
if score < 0.1 {
47-
Some((OrderedFloat(score), entry))
37+
Some((OrderedFloat(score), id, abi))
4838
} else {
4939
None
5040
}
5141
})
52-
.sorted_by_key(|(score, _)| *score)
53-
.next()?;
42+
.min_by_key(|(score, _, _)| *score)?;
5443

5544
Some(AddressIdentity {
5645
address: *address,

‎crates/evm/src/trace/identifier/mod.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,7 @@ pub struct AddressIdentity<'a> {
3333
pub trait TraceIdentifier {
3434
// TODO: Update docs
3535
/// Attempts to identify an address in one or more call traces.
36-
#[allow(clippy::type_complexity)]
37-
fn identify_addresses(
38-
&mut self,
39-
addresses: Vec<(&Address, Option<&[u8]>)>,
40-
) -> Vec<AddressIdentity>;
36+
fn identify_addresses<'a, A>(&mut self, addresses: A) -> Vec<AddressIdentity>
37+
where
38+
A: Iterator<Item = (&'a Address, Option<&'a [u8]>)>;
4139
}

‎crates/evm/src/trace/mod.rs‎

Lines changed: 43 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -561,13 +561,39 @@ impl fmt::Display for CallTrace {
561561
}
562562

563563
/// Specifies the kind of trace.
564-
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
564+
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
565565
pub enum TraceKind {
566566
Deployment,
567567
Setup,
568568
Execution,
569569
}
570570

571+
impl TraceKind {
572+
/// Returns `true` if the trace kind is [`Deployment`].
573+
///
574+
/// [`Deployment`]: TraceKind::Deployment
575+
#[must_use]
576+
pub fn is_deployment(self) -> bool {
577+
matches!(self, Self::Deployment)
578+
}
579+
580+
/// Returns `true` if the trace kind is [`Setup`].
581+
///
582+
/// [`Setup`]: TraceKind::Setup
583+
#[must_use]
584+
pub fn is_setup(self) -> bool {
585+
matches!(self, Self::Setup)
586+
}
587+
588+
/// Returns `true` if the trace kind is [`Execution`].
589+
///
590+
/// [`Execution`]: TraceKind::Execution
591+
#[must_use]
592+
pub fn is_execution(self) -> bool {
593+
matches!(self, Self::Execution)
594+
}
595+
}
596+
571597
/// Chooses the color of the trace depending on the destination address and status of the call.
572598
fn trace_color(trace: &CallTrace) -> Color {
573599
if trace.address == CHEATCODE_ADDRESS {
@@ -584,26 +610,23 @@ pub fn load_contracts(
584610
traces: Traces,
585611
known_contracts: Option<&ContractsByArtifact>,
586612
) -> ContractsByAddress {
587-
if let Some(contracts) = known_contracts {
588-
let mut local_identifier = LocalTraceIdentifier::new(contracts);
589-
let mut decoder = CallTraceDecoderBuilder::new().build();
590-
for (_, trace) in &traces {
591-
decoder.identify(trace, &mut local_identifier);
592-
}
593-
594-
decoder
595-
.contracts
596-
.iter()
597-
.filter_map(|(addr, name)| {
598-
if let Ok(Some((_, (abi, _)))) = contracts.find_by_name_or_identifier(name) {
599-
return Some((*addr, (name.clone(), abi.clone())))
600-
}
601-
None
602-
})
603-
.collect()
604-
} else {
605-
BTreeMap::new()
613+
let Some(contracts) = known_contracts else { return BTreeMap::new() };
614+
let mut local_identifier = LocalTraceIdentifier::new(contracts);
615+
let mut decoder = CallTraceDecoderBuilder::new().build();
616+
for (_, trace) in &traces {
617+
decoder.identify(trace, &mut local_identifier);
606618
}
619+
620+
decoder
621+
.contracts
622+
.iter()
623+
.filter_map(|(addr, name)| {
624+
if let Ok(Some((_, (abi, _)))) = contracts.find_by_name_or_identifier(name) {
625+
return Some((*addr, (name.clone(), abi.clone())))
626+
}
627+
None
628+
})
629+
.collect()
607630
}
608631

609632
/// creates the memory data in 32byte chunks

‎crates/evm/src/trace/utils.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ pub(crate) fn decode_cheatcode_inputs(
7979
"serializeBytes32" |
8080
"serializeString" |
8181
"serializeBytes" => {
82-
if verbosity == 5 {
82+
if verbosity >= 5 {
8383
None
8484
} else {
8585
let mut decoded = func.decode_input(&data[SELECTOR_LEN..]).ok()?;
@@ -111,10 +111,10 @@ pub(crate) fn decode_cheatcode_outputs(
111111
// redacts derived private key
112112
return Some("<pk>".to_string())
113113
}
114-
if func.name == "parseJson" && verbosity != 5 {
114+
if func.name == "parseJson" && verbosity < 5 {
115115
return Some("<encoded JSON value>".to_string())
116116
}
117-
if func.name == "readFile" && verbosity != 5 {
117+
if func.name == "readFile" && verbosity < 5 {
118118
return Some("<file>".to_string())
119119
}
120120
None

0 commit comments

Comments
 (0)