diff --git a/crates/common/src/compile.rs b/crates/common/src/compile.rs index eccea208cd04..15debf8a539b 100644 --- a/crates/common/src/compile.rs +++ b/crates/common/src/compile.rs @@ -271,9 +271,50 @@ impl ProjectCompiler { } } -/// Map over artifacts contract sources name -> file_id -> (source, contract) +/// Contract source code and bytecode. #[derive(Clone, Debug, Default)] -pub struct ContractSources(pub HashMap>); +pub struct ContractSources { + /// Map over artifacts' contract names -> vector of file IDs + pub ids_by_name: HashMap>, + /// Map over file_id -> (source code, contract) + pub sources_by_id: HashMap, +} + +impl ContractSources { + /// Inserts a contract into the sources. + pub fn insert( + &mut self, + artifact_id: &ArtifactId, + file_id: u32, + source: String, + bytecode: ContractBytecodeSome, + ) { + self.ids_by_name.entry(artifact_id.name.clone()).or_default().push(file_id); + self.sources_by_id.insert(file_id, (source, bytecode)); + } + + /// Returns the source for a contract by file ID. + pub fn get(&self, id: u32) -> Option<&(String, ContractBytecodeSome)> { + self.sources_by_id.get(&id) + } + + /// Returns all sources for a contract by name. + pub fn get_sources( + &self, + name: &str, + ) -> Option> { + self.ids_by_name + .get(name) + .map(|ids| ids.iter().filter_map(|id| Some((*id, self.sources_by_id.get(id)?)))) + } + + /// Returns all (name, source) pairs. + pub fn entries(&self) -> impl Iterator { + self.ids_by_name.iter().flat_map(|(name, ids)| { + ids.iter().filter_map(|id| self.sources_by_id.get(id).map(|s| (name.clone(), s))) + }) + } +} // https://eips.ethereum.org/EIPS/eip-170 const CONTRACT_SIZE_LIMIT: usize = 24576; diff --git a/crates/debugger/src/tui/draw.rs b/crates/debugger/src/tui/draw.rs index e38ac170b910..7a834392986d 100644 --- a/crates/debugger/src/tui/draw.rs +++ b/crates/debugger/src/tui/draw.rs @@ -334,7 +334,9 @@ impl DebuggerContext<'_> { return Err(format!("Unknown contract at address {address}")); }; - let Some(files_source_code) = self.debugger.contracts_sources.0.get(contract_name) else { + let Some(mut files_source_code) = + self.debugger.contracts_sources.get_sources(contract_name) + else { return Err(format!("No source map index for contract {contract_name}")); }; @@ -345,7 +347,7 @@ impl DebuggerContext<'_> { let is_create = matches!(self.call_kind(), CallKind::Create | CallKind::Create2); let pc = self.current_step().pc; let Some((source_element, source_code)) = - files_source_code.iter().find_map(|(file_id, (source_code, contract_source))| { + files_source_code.find_map(|(file_id, (source_code, contract_source))| { let bytecode = if is_create { &contract_source.bytecode } else { @@ -356,7 +358,20 @@ impl DebuggerContext<'_> { let pc_ic_map = if is_create { create_map } else { rt_map }; let ic = pc_ic_map.get(pc)?; let source_element = source_map.swap_remove(ic); - (*file_id == source_element.index?).then_some((source_element, source_code)) + // if the source element has an index, find the sourcemap for that index + source_element + .index + .and_then(|index| + // if index matches current file_id, return current source code + (index == file_id).then(|| (source_element.clone(), source_code))) + .or_else(|| { + // otherwise find the source code for the element's index + self.debugger + .contracts_sources + .sources_by_id + .get(&(source_element.index?)) + .map(|(source_code, _)| (source_element.clone(), source_code)) + }) }) else { return Err(format!("No source map for contract {contract_name}")); diff --git a/crates/debugger/src/tui/mod.rs b/crates/debugger/src/tui/mod.rs index 523e7faf15a0..d9543fc58a8d 100644 --- a/crates/debugger/src/tui/mod.rs +++ b/crates/debugger/src/tui/mod.rs @@ -66,18 +66,15 @@ impl Debugger { breakpoints: Breakpoints, ) -> Self { let pc_ic_maps = contracts_sources - .0 - .iter() - .flat_map(|(contract_name, files_sources)| { - files_sources.iter().filter_map(|(_, (_, contract))| { - Some(( - contract_name.clone(), - ( - PcIcMap::new(SpecId::LATEST, contract.bytecode.bytes()?), - PcIcMap::new(SpecId::LATEST, contract.deployed_bytecode.bytes()?), - ), - )) - }) + .entries() + .filter_map(|(contract_name, (_, contract))| { + Some(( + contract_name.clone(), + ( + PcIcMap::new(SpecId::LATEST, contract.bytecode.bytes()?), + PcIcMap::new(SpecId::LATEST, contract.deployed_bytecode.bytes()?), + ), + )) }) .collect(); Self { debug_arena, identified_contracts, contracts_sources, pc_ic_maps, breakpoints } diff --git a/crates/evm/traces/src/identifier/etherscan.rs b/crates/evm/traces/src/identifier/etherscan.rs index cb6952ab50dd..1ddbf2457d6f 100644 --- a/crates/evm/traces/src/identifier/etherscan.rs +++ b/crates/evm/traces/src/identifier/etherscan.rs @@ -87,11 +87,7 @@ impl EtherscanIdentifier { for (results, (_, metadata)) in artifacts.into_iter().zip(contracts_iter) { // get the inner type let (artifact_id, file_id, bytecode) = results?; - sources - .0 - .entry(artifact_id.clone().name) - .or_default() - .insert(file_id, (metadata.source_code(), bytecode)); + sources.insert(&artifact_id, file_id, metadata.source_code(), bytecode); } Ok(sources) diff --git a/crates/forge/bin/cmd/script/build.rs b/crates/forge/bin/cmd/script/build.rs index 27d691eeb745..979ac5ee5a0f 100644 --- a/crates/forge/bin/cmd/script/build.rs +++ b/crates/forge/bin/cmd/script/build.rs @@ -48,11 +48,7 @@ impl ScriptArgs { })?; let contract = artifact.clone().into_contract_bytecode(); let source_contract = compact_to_contract(contract)?; - sources - .0 - .entry(id.clone().name) - .or_default() - .insert(source.id, (source_code, source_contract)); + sources.insert(&id, source.id, source_code, source_contract); } else { warn!(?id, "source not found"); } diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index 8fe5966faa47..1a8cf9eed418 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -272,11 +272,7 @@ impl TestArgs { let source_code = fs::read_to_string(abs_path)?; let contract = artifact.clone().into_contract_bytecode(); let source_contract = compact_to_contract(contract)?; - sources - .0 - .entry(id.name.clone()) - .or_default() - .insert(source.id, (source_code, source_contract)); + sources.insert(&id, source.id, source_code, source_contract); } }