fix(forge): debugger should load sourcemaps by file_id (#7058)

* initial fix for sourcemaps in debugger

* refactor sources_by_name to ids_by_name

* bring up nested or_else

* return iterators and update docstring

* cargo fmt

* then_some -> then

* remove whitespace

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<String, HashMap<u32, (String, ContractBytecodeSome)>>);
+pub struct ContractSources {
+    /// Map over artifacts' contract names -> vector of file IDs
+    pub ids_by_name: HashMap<String, Vec<u32>>,
+    /// Map over file_id -> (source code, contract)
+    pub sources_by_id: HashMap<u32, (String, ContractBytecodeSome)>,
+}
+
+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<impl Iterator<Item = (u32, &(String, ContractBytecodeSome))>> {
+        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<Item = (String, &(String, ContractBytecodeSome))> {
+        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;

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}"));

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 }

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)

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");
                 }

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);
                 }
             }