Stream traces

.gitignore

@@ -10,3 +10,4 @@ on_chain_bytecode.txt
 .env
 tests/hardhat/package-lock.json
 .dv_config.json
+/.claude

docs/config.md

@@ -23,7 +23,7 @@ When running the `dvf` command, the default configuration file is expected at `$
 | - - `api_url` | Chain-specific Blockscout API URL |
 | - - `api_key` | Chain-specific Blockscout API Key |
 | `max_blocks_per_event_query` | Number of blocks that can be queried at once in `getLogs`, optional, defaults to 9999 |
-| `web3_timeout` | Timeout in seconds for web3 RPC queries, optional, defaults to 5000 |
+| `web3_timeout` | Timeout in milliseconds for web3 RPC queries, optional, defaults to 5000 |
 | `signer` | Configuration on how to sign, optional |
 | - `wallet_address` | Address which is used to sign |
 | - `wallet_type` | Can have different structure |

lib/dvf/discovery.rs

@@ -15,14 +15,13 @@ use crate::dvf::config::DVFConfig;
 use crate::dvf::parse::{self, ValidationError};
 use crate::dvf::registry;
 use crate::state::contract_state::ContractState;
+use crate::state::contract_state::MappingUsages;
 use crate::state::forge_inspect;
 use crate::utils::pretty::PrettyPrinter;
 use crate::utils::progress::{print_progress, ProgressMode};
 use crate::utils::read_write_file::get_project_paths;
 use crate::web3;
 use crate::web3::stop_anvil_instance;
-use crate::web3::TraceWithAddress;
-use alloy_node_bindings::AnvilInstance;
 
 pub struct DiscoveryParams<'a> {
     pub config: &'a DVFConfig,
@@ -46,9 +45,8 @@ pub struct DiscoveryParams<'a> {
     pub progress_mode: &'a ProgressMode,
     pub use_storage_range: bool,
     pub tx_hashes: Option<Vec<String>>,
-    // Optional cache (used by inspect-tx): reuse an already computed trace and config
-    pub cached_traces: Option<Vec<TraceWithAddress>>,
-    pub cached_anvil_config: Option<&'a DVFConfig>,
+    // Optional cache (used by inspect-tx): reuse pre-computed mapping usages
+    pub cached_mapping_usages: Option<MappingUsages>,
 }
 
 pub struct DiscoveryResult {
@@ -211,47 +209,48 @@ pub fn discover_storage_and_events(
     )?;
 
     print_progress("Getting relevant traces.", params.pc, params.progress_mode);
-    let mut seen_transactions = HashSet::new();
     let mut missing_traces = false;
 
-    for (index, tx_hash) in tx_hashes.iter().enumerate() {
-        if seen_transactions.contains(tx_hash) {
-            continue;
-        }
-        seen_transactions.insert(tx_hash);
-
-        info!("Getting trace for {}", tx_hash);
-        // Use cached trace if provided (inspect-tx), otherwise fetch
-        let fetched = if let Some(ref cached) = params.cached_traces {
-            debug!("Using cached trace at index {} of {}", index, cached.len());
-            Ok((
-                cached[index].clone(),
-                None::<DVFConfig>,
-                None::<AnvilInstance>,
-            ))
-        } else {
-            web3::get_eth_debug_trace_sim(params.config, tx_hash)
-        };
-        match fetched {
-            Ok((trace, anvil_config, anvil_instance)) => {
-                let record_traces_config: &DVFConfig = if params.cached_traces.is_some() {
-                    params.cached_anvil_config.unwrap_or(params.config)
-                } else {
-                    anvil_config.as_ref().unwrap_or(params.config)
-                };
-                if let Err(err) = contract_state.record_traces(record_traces_config, vec![trace]) {
+    if let Some(cached_usages) = params.cached_mapping_usages {
+        // Inject pre-computed mapping usages directly (from inspect-tx)
+        debug!("Using cached mapping usages");
+        contract_state.inject_mapping_usages(cached_usages);
+    } else {
+        // Stream traces to collect mapping usages
+        let mut seen_transactions = HashSet::new();
+        for tx_hash in tx_hashes.iter() {
+            if seen_transactions.contains(tx_hash) {
+                continue;
+            }
+            seen_transactions.insert(tx_hash);
+
+            info!("Getting trace for {}", tx_hash);
+            let trace_address = match web3::get_receipt_address(params.config, tx_hash) {
+                Ok(addr) => addr,
+                Err(err) => {
                     missing_traces = true;
                     info!("Warning. The trace for {tx_hash} cannot be obtained. Some mapping slots might not be decodable. You can try to increase the timeout in the config. Error: {}", err);
+                    continue;
                 }
-                if params.cached_traces.is_none() {
+            };
+            let mut processor = crate::state::contract_state::MappingUsageProcessor::new(
+                *params.address,
+                trace_address,
+                params.config,
+                tx_hash.clone(),
+                true,
+            );
+            match web3::stream_eth_debug_trace_sim(params.config, tx_hash, &mut processor) {
+                Ok((_metadata, _address, _anvil_config, anvil_instance)) => {
+                    contract_state.inject_mapping_usages(processor.mapping_usages);
                     if let Some(anvil_instance) = anvil_instance {
                         stop_anvil_instance(anvil_instance);
                     }
                 }
-            }
-            Err(err) => {
-                missing_traces = true;
-                info!("Warning. The trace for {tx_hash} cannot be obtained. Some mapping slots might not be decodable. You can try to increase the timeout in the config. Error: {}", err);
+                Err(err) => {
+                    missing_traces = true;
+                    info!("Warning. The trace for {tx_hash} cannot be obtained. Some mapping slots might not be decodable. You can try to increase the timeout in the config. Error: {}", err);
+                }
             }
         }
     }
@@ -546,8 +545,7 @@ pub fn create_discovery_params_for_init<'a>(
         progress_mode,
         use_storage_range: true,
         tx_hashes: None,
-        cached_traces: None,
-        cached_anvil_config: None,
+        cached_mapping_usages: None,
     }
 }
 
@@ -593,7 +591,6 @@ pub fn create_discovery_params_for_update<'a>(
         progress_mode,
         use_storage_range: false, // cannot use storage range here as we are only trying to get a subset of the state
         tx_hashes: None,
-        cached_traces: None,
-        cached_anvil_config: None,
+        cached_mapping_usages: None,
     }
 }

lib/state/contract_state.rs

@@ -17,7 +17,9 @@ use crate::state::forge_inspect::{
     ForgeInspectIrOptimized, ForgeInspectLayoutStorage, StateVariable, TypeDescription,
 };
 use crate::utils::pretty::PrettyPrinter;
-use crate::web3::{get_internal_create_addresses, StorageSnapshot, TraceWithAddress};
+use crate::web3::{
+    get_internal_create_addresses, StorageSnapshot, StructLogProcessor, TraceWithAddress,
+};
 
 fn hash_u256(u: &U256) -> B256 {
     keccak256(u.to_be_bytes::<32>())
@@ -388,6 +390,16 @@ impl<'a> ContractState<'a> {
         Ok(())
     }
 
+    /// Inject pre-computed mapping usages directly, skipping trace processing.
+    pub fn inject_mapping_usages(&mut self, usages: HashMap<U256, HashSet<(String, U256)>>) {
+        for (index, entries) in usages {
+            self.mapping_usages
+                .entry(index)
+                .or_default()
+                .extend(entries);
+        }
+    }
+
     fn add_to_table(storage_entry: &parse::DVFStorageEntry, table: &mut Table) {
         PrettyPrinter::add_formatted_to_table(
             &storage_entry.var_name,
@@ -931,3 +943,225 @@ impl<'a> ContractState<'a> {
         var_type.starts_with("t_userDefinedValueType")
     }
 }
+
+/// Type alias for mapping usages: storage index -> set of (key, derived_slot).
+pub type MappingUsages = HashMap<U256, HashSet<(String, U256)>>;
+
+/// Helper: extract mapping key from memory at a SHA3/KECCAK256 opcode.
+/// Returns `Some((key_hex, storage_index))` if the input looks like a mapping hash.
+fn extract_mapping_key_from_sha3(stack: &[U256], memory: &[String]) -> Option<(String, U256)> {
+    let length_in_bytes = stack[stack.len() - 2];
+    if length_in_bytes < U256::from(32_u64) || length_in_bytes >= U256::from(usize::MAX / 2) {
+        return None;
+    }
+    let mem_str: String = memory.iter().cloned().collect();
+    let start_idx = stack[stack.len() - 1].to::<usize>() * 2;
+    let length = length_in_bytes.to::<usize>() * 2;
+    let sha3_input = format!("0x{}", &mem_str[start_idx..(start_idx + length)]);
+
+    let usize_str_length = length_in_bytes.to::<usize>() * 2 + 2;
+    assert!(sha3_input.len() == usize_str_length);
+    let key = sha3_input[2..usize_str_length - 64].to_string();
+    let index = U256::from_str_radix(&sha3_input[usize_str_length - 64..], 16).ok()?;
+    Some((key, index))
+}
+
+/// Processor that collects mapping usages for a single address from trace logs.
+/// Extracted from `ContractState::record_traces()`.
+pub struct MappingUsageProcessor<'a> {
+    address: Address,
+    config: &'a DVFConfig,
+    tx_id: String,
+    is_first_trace: bool,
+    depth_to_address: HashMap<u64, Address>,
+    create_addresses: Option<Vec<Address>>,
+    key: Option<String>,
+    index: U256,
+    pub mapping_usages: MappingUsages,
+    failed: bool,
+}
+
+impl<'a> MappingUsageProcessor<'a> {
+    pub fn new(
+        address: Address,
+        trace_address: Address,
+        config: &'a DVFConfig,
+        tx_id: String,
+        is_first_trace: bool,
+    ) -> Self {
+        let mut depth_to_address: HashMap<u64, Address> = HashMap::new();
+        depth_to_address.insert(1, trace_address);
+        MappingUsageProcessor {
+            address,
+            config,
+            tx_id,
+            is_first_trace,
+            depth_to_address,
+            create_addresses: None,
+            key: None,
+            index: U256::from(1),
+            mapping_usages: HashMap::new(),
+            failed: false,
+        }
+    }
+}
+
+impl<'a> StructLogProcessor for MappingUsageProcessor<'a> {
+    fn process_log(
+        &mut self,
+        log: alloy_rpc_types_trace::geth::StructLog,
+    ) -> Result<(), crate::dvf::parse::ValidationError> {
+        if log.stack.is_none() {
+            return Ok(());
+        }
+        let stack = log.stack.unwrap();
+
+        if log.op == "CREATE" || log.op == "CREATE2" {
+            if self.is_first_trace {
+                if self.create_addresses.is_none() {
+                    self.create_addresses =
+                        Some(get_internal_create_addresses(self.config, &self.tx_id)?);
+                }
+                if let Some(ref mut create_ref) = self.create_addresses {
+                    self.depth_to_address
+                        .insert(log.depth + 1, create_ref.remove(0));
+                }
+            } else {
+                self.depth_to_address
+                    .insert(log.depth + 1, Address::from([0; 20]));
+            }
+        }
+
+        if log.op == "CALL" || log.op == "STATICCALL" {
+            let address_bytes = stack[stack.len() - 2].to_be_bytes::<32>();
+            let a = Address::from_slice(&address_bytes[12..]);
+            self.depth_to_address.insert(log.depth + 1, a);
+        }
+
+        if log.op == "DELEGATECALL" || log.op == "CALLCODE" {
+            self.depth_to_address
+                .insert(log.depth + 1, self.depth_to_address[&log.depth]);
+        }
+
+        if self.depth_to_address[&log.depth] == self.address {
+            if let Some(key_in) = self.key.take() {
+                let target_slot = &stack[stack.len() - 1];
+                self.mapping_usages
+                    .entry(self.index)
+                    .or_default()
+                    .insert((key_in, *target_slot));
+            }
+
+            if log.op == "KECCAK256" || log.op == "SHA3" {
+                if let Some(memory) = &log.memory {
+                    if let Some((key, index)) = extract_mapping_key_from_sha3(&stack, memory) {
+                        debug!("Found key {} for index {}.", key, index);
+                        self.key = Some(key);
+                        self.index = index;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn trace_failed(&mut self) {
+        self.failed = true;
+        self.mapping_usages.clear();
+    }
+}
+
+/// Processor that collects mapping usages for ALL addresses encountered in a trace.
+/// Used by `inspect-tx` to process the trace once for all contracts.
+pub struct MultiAddressMappingProcessor<'a> {
+    config: &'a DVFConfig,
+    tx_id: String,
+    depth_to_address: HashMap<u64, Address>,
+    create_addresses: Option<Vec<Address>>,
+    key: Option<String>,
+    index: U256,
+    /// Per-address mapping usages.
+    pub all_mapping_usages: HashMap<Address, MappingUsages>,
+    failed: bool,
+}
+
+impl<'a> MultiAddressMappingProcessor<'a> {
+    pub fn new(trace_address: Address, config: &'a DVFConfig, tx_id: String) -> Self {
+        let mut depth_to_address: HashMap<u64, Address> = HashMap::new();
+        depth_to_address.insert(1, trace_address);
+        MultiAddressMappingProcessor {
+            config,
+            tx_id,
+            depth_to_address,
+            create_addresses: None,
+            key: None,
+            index: U256::from(1),
+            all_mapping_usages: HashMap::new(),
+            failed: false,
+        }
+    }
+}
+
+impl<'a> StructLogProcessor for MultiAddressMappingProcessor<'a> {
+    fn process_log(
+        &mut self,
+        log: alloy_rpc_types_trace::geth::StructLog,
+    ) -> Result<(), crate::dvf::parse::ValidationError> {
+        if log.stack.is_none() {
+            return Ok(());
+        }
+        let stack = log.stack.unwrap();
+
+        if log.op == "CREATE" || log.op == "CREATE2" {
+            if self.create_addresses.is_none() {
+                self.create_addresses =
+                    Some(get_internal_create_addresses(self.config, &self.tx_id)?);
+            }
+            if let Some(ref mut create_ref) = self.create_addresses {
+                self.depth_to_address
+                    .insert(log.depth + 1, create_ref.remove(0));
+            }
+        }
+
+        if log.op == "CALL" || log.op == "STATICCALL" {
+            let address_bytes = stack[stack.len() - 2].to_be_bytes::<32>();
+            let a = Address::from_slice(&address_bytes[12..]);
+            self.depth_to_address.insert(log.depth + 1, a);
+        }
+
+        if log.op == "DELEGATECALL" || log.op == "CALLCODE" {
+            self.depth_to_address
+                .insert(log.depth + 1, self.depth_to_address[&log.depth]);
+        }
+
+        let current_address = self.depth_to_address[&log.depth];
+
+        if let Some(key_in) = self.key.take() {
+            let target_slot = &stack[stack.len() - 1];
+            self.all_mapping_usages
+                .entry(current_address)
+                .or_default()
+                .entry(self.index)
+                .or_default()
+                .insert((key_in, *target_slot));
+        }
+
+        if log.op == "KECCAK256" || log.op == "SHA3" {
+            if let Some(memory) = &log.memory {
+                if let Some((key, index)) = extract_mapping_key_from_sha3(&stack, memory) {
+                    debug!("Found key {} for index {}.", key, index);
+                    self.key = Some(key);
+                    self.index = index;
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn trace_failed(&mut self) {
+        self.failed = true;
+        self.all_mapping_usages.clear();
+    }
+}