refactor sources_by_name to ids_by_name

foundry-rs · Feb 9, 2024 · 1ab8d4d · 1ab8d4d
1 parent e7515d8
commit 1ab8d4d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 24 deletions.
diff --git a/crates/common/src/compile.rs b/crates/common/src/compile.rs
@@ -275,7 +275,7 @@ impl ProjectCompiler {
 #[derive(Clone, Debug, Default)]
 pub struct ContractSources {
     /// Map over artifacts contract sources name -> file_id -> (source, contract)
-    pub sources_by_name: HashMap<String, HashMap<u32, (String, ContractBytecodeSome)>>,
+    pub ids_by_name: HashMap<String, Vec<u32>>,
     /// Map over artifacts contract sources file_id -> (source, contract)
     pub sources_by_id: HashMap<u32, (String, ContractBytecodeSome)>,
 }
@@ -289,22 +289,32 @@ impl ContractSources {
         source: String,
         bytecode: ContractBytecodeSome,
     ) {
-        self.sources_by_name
-            .entry(artifact_id.name.clone())
-            .or_default()
-            .insert(file_id, (source.clone(), bytecode.clone()));
+        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 sources for contracts by name.
-    pub fn get(&self, name: &str) -> Option<&HashMap<u32, (String, ContractBytecodeSome)>> {
-        self.sources_by_name.get(name)
-    }
-
     /// Returns the source for a contract by file ID.
-    pub fn get_by_id(&self, id: u32) -> Option<&(String, ContractBytecodeSome)> {
+    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<Vec<(u32, &(String, ContractBytecodeSome))>> {
+        self.ids_by_name.get(name).map(|ids| {
+            ids.iter().filter_map(|id| Some((*id, self.sources_by_id.get(id)?))).collect()
+        })
+    }
+
+    /// Returns all (name, source) pairs.
+    pub fn entries(&self) -> Vec<(String, &(String, ContractBytecodeSome))> {
+        self.ids_by_name
+            .iter()
+            .flat_map(|(name, ids)| {
+                ids.iter()
+                    .filter_map(move |id| self.sources_by_id.get(id).map(|s| (name.clone(), s)))
+            })
+            .collect()
+    }
 }
 
 // https://eips.ethereum.org/EIPS/eip-170

diff --git a/crates/debugger/src/tui/draw.rs b/crates/debugger/src/tui/draw.rs
@@ -334,8 +334,7 @@ impl DebuggerContext<'_> {
             return Err(format!("Unknown contract at address {address}"));
         };
 
-        let Some(files_source_code) =
-            self.debugger.contracts_sources.sources_by_name.get(contract_name)
+        let Some(files_source_code) = self.debugger.contracts_sources.get_sources(contract_name)
         else {
             return Err(format!("No source map index for contract {contract_name}"));
         };

diff --git a/crates/debugger/src/tui/mod.rs b/crates/debugger/src/tui/mod.rs
@@ -66,18 +66,16 @@ impl Debugger {
         breakpoints: Breakpoints,
     ) -> Self {
         let pc_ic_maps = contracts_sources
-            .sources_by_name
+            .entries()
             .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()?),
-                        ),
-                    ))
-                })
+            .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 }