From 6ce627030e412e43f49622a5100a12d4fcf30c7b Mon Sep 17 00:00:00 2001
From: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
Date: Tue, 13 Feb 2024 17:58:46 +0100
Subject: [PATCH] chore: better ContractsByArtifact flatten

---
 crates/common/src/contracts.rs   | 58 +++++++++++++++-----------------
 crates/forge/src/multi_runner.rs |  4 +--
 2 files changed, 30 insertions(+), 32 deletions(-)

diff --git a/crates/common/src/contracts.rs b/crates/common/src/contracts.rs
index b47e928b11eb..650280cc5375 100644
--- a/crates/common/src/contracts.rs
+++ b/crates/common/src/contracts.rs
@@ -1,7 +1,7 @@
 //! Commonly used contract types and functions.
 
 use alloy_json_abi::{Event, Function, JsonAbi};
-use alloy_primitives::{hex, Address, B256};
+use alloy_primitives::{hex, Address, Selector, B256};
 use foundry_compilers::{
     artifacts::{CompactContractBytecode, ContractBytecodeSome},
     ArtifactId, ProjectPathsConfig,
@@ -43,36 +43,34 @@ impl ContractsByArtifact {
         Ok(contracts.first().cloned())
     }
 
-    /// Flattens a group of contracts into maps of all events and functions
-    pub fn flatten(&self) -> (BTreeMap<[u8; 4], Function>, BTreeMap<B256, Event>, JsonAbi) {
-        let flattened_funcs: BTreeMap<[u8; 4], Function> = self
-            .iter()
-            .flat_map(|(_name, (abi, _code))| {
-                abi.functions()
-                    .map(|func| (func.selector().into(), func.clone()))
-                    .collect::<BTreeMap<[u8; 4], Function>>()
-            })
-            .collect();
+    /// Flattens the contracts into functions, events and errors.
+    pub fn flatten(&self) -> (BTreeMap<Selector, Function>, BTreeMap<B256, Event>, JsonAbi) {
+        let mut funcs = BTreeMap::new();
+        let mut events = BTreeMap::new();
+        let mut errors_abi = JsonAbi::new();
+        for (_name, (abi, _code)) in self.iter() {
+            for func in abi.functions() {
+                funcs.insert(func.selector(), func.clone());
+            }
+            for event in abi.events() {
+                events.insert(event.selector(), event.clone());
+            }
+            for error in abi.errors() {
+                errors_abi.errors.entry(error.name.clone()).or_default().push(error.clone());
+            }
+        }
+        (funcs, events, errors_abi)
+    }
 
-        let flattened_events: BTreeMap<B256, Event> = self
-            .iter()
-            .flat_map(|(_name, (abi, _code))| {
-                abi.events()
-                    .map(|event| (event.selector(), event.clone()))
-                    .collect::<BTreeMap<B256, Event>>()
-            })
-            .collect();
-
-        // We need this for better revert decoding, and want it in abi form
-        let mut errors_abi = JsonAbi::default();
-        self.iter().for_each(|(_name, (abi, _code))| {
-            abi.errors().for_each(|error| {
-                let entry =
-                    errors_abi.errors.entry(error.name.clone()).or_insert_with(Default::default);
-                entry.push(error.clone());
-            });
-        });
-        (flattened_funcs, flattened_events, errors_abi)
+    /// Flattens the errors into a single JsonAbi.
+    pub fn flatten_errors(&self) -> JsonAbi {
+        let mut errors_abi = JsonAbi::new();
+        for (_name, (abi, _code)) in self.iter() {
+            for error in abi.errors() {
+                errors_abi.errors.entry(error.name.clone()).or_default().push(error.clone());
+            }
+        }
+        errors_abi
     }
 }
 
diff --git a/crates/forge/src/multi_runner.rs b/crates/forge/src/multi_runner.rs
index 5cfd59d2752a..0c876a640f29 100644
--- a/crates/forge/src/multi_runner.rs
+++ b/crates/forge/src/multi_runner.rs
@@ -316,7 +316,7 @@ impl MultiContractRunnerBuilder {
             );
         }
 
-        let execution_info = known_contracts.flatten();
+        let errors = known_contracts.flatten_errors();
         Ok(MultiContractRunner {
             contracts: deployable_contracts,
             known_contracts,
@@ -324,7 +324,7 @@ impl MultiContractRunnerBuilder {
             env,
             evm_spec: self.evm_spec.unwrap_or(SpecId::MERGE),
             sender: self.sender,
-            errors: Some(execution_info.2),
+            errors: Some(errors),
             source_paths,
             fork: self.fork,
             cheats_config: self.cheats_config.unwrap_or_default().into(),