fix: debugger breaks when source file has multiple contract definitio…

…ns (#7550) * fix contracts sources Co-authored-by: lazymio <mio@lazym.io> * fix doc * fmt --------- Co-authored-by: lazymio <mio@lazym.io>
foundry-rs · Apr 3, 2024 · f6208d8 · f6208d8
1 parent dbc48ea
commit f6208d8
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 21 deletions.
diff --git a/crates/common/src/compile.rs b/crates/common/src/compile.rs
@@ -288,8 +288,10 @@ impl ProjectCompiler {
 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: FxHashMap<u32, (String, ContractBytecodeSome)>,
+    /// Map over file_id -> source code
+    pub sources_by_id: FxHashMap<u32, String>,
+    /// Map over file_id -> contract name -> bytecode
+    pub artifacts_by_id: FxHashMap<u32, HashMap<String, ContractBytecodeSome>>,
 }
 
 impl ContractSources {
@@ -327,30 +329,45 @@ impl ContractSources {
         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));
+        self.sources_by_id.insert(file_id, source);
+        self.artifacts_by_id.entry(file_id).or_default().insert(artifact_id.name.clone(), bytecode);
     }
 
     /// Returns the source for a contract by file ID.
-    pub fn get(&self, id: u32) -> Option<&(String, ContractBytecodeSome)> {
+    pub fn get(&self, id: u32) -> Option<&String> {
         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)))
+    pub fn get_sources<'a>(
+        &'a self,
+        name: &'a str,
+    ) -> Option<impl Iterator<Item = (u32, &'_ str, &'_ ContractBytecodeSome)>> {
+        self.ids_by_name.get(name).map(|ids| {
+            ids.iter().filter_map(|id| {
+                Some((
+                    *id,
+                    self.sources_by_id.get(id)?.as_ref(),
+                    self.artifacts_by_id.get(id)?.get(name)?,
+                ))
+            })
         })
     }
+
+    /// Returns all (name, source, bytecode) sets.
+    pub fn entries(&self) -> impl Iterator<Item = (&str, &str, &ContractBytecodeSome)> {
+        self.artifacts_by_id
+            .iter()
+            .filter_map(|(id, artifacts)| {
+                let source = self.sources_by_id.get(id)?;
+                Some(
+                    artifacts
+                        .iter()
+                        .map(move |(name, bytecode)| (name.as_ref(), source.as_ref(), bytecode)),
+                )
+            })
+            .flatten()
+    }
 }
 
 // https://eips.ethereum.org/EIPS/eip-170

diff --git a/crates/debugger/src/tui/draw.rs b/crates/debugger/src/tui/draw.rs
@@ -346,7 +346,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.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 {
@@ -369,7 +369,7 @@ impl DebuggerContext<'_> {
                             .contracts_sources
                             .sources_by_id
                             .get(&(source_element.index?))
-                            .map(|(source_code, _)| (source_element.clone(), source_code))
+                            .map(|source_code| (source_element.clone(), source_code.as_ref()))
                     })
             })
         else {

diff --git a/crates/debugger/src/tui/mod.rs b/crates/debugger/src/tui/mod.rs
@@ -67,9 +67,9 @@ impl Debugger {
     ) -> Self {
         let pc_ic_maps = contracts_sources
             .entries()
-            .filter_map(|(contract_name, (_, contract))| {
+            .filter_map(|(contract_name, _, contract)| {
                 Some((
-                    contract_name.clone(),
+                    contract_name.to_owned(),
                     (
                         PcIcMap::new(SpecId::LATEST, contract.bytecode.bytes()?),
                         PcIcMap::new(SpecId::LATEST, contract.deployed_bytecode.bytes()?),