feat: coverage for constructors (#7661)

* wip

* better naming

crates/common/src/contracts.rs

@@ -1,34 +1,48 @@
 //! Commonly used contract types and functions.
 
 use alloy_json_abi::{Event, Function, JsonAbi};
-use alloy_primitives::{hex, Address, Selector, B256};
+use alloy_primitives::{Address, Bytes, Selector, B256};
 use eyre::Result;
 use foundry_compilers::{
     artifacts::{CompactContractBytecode, ContractBytecodeSome},
     ArtifactId,
 };
 use std::{
     collections::BTreeMap,
-    fmt,
     ops::{Deref, DerefMut},
 };
 
-type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a (JsonAbi, Vec<u8>));
+/// Container for commonly used contract data.
+#[derive(Debug, Clone)]
+pub struct ContractData {
+    /// Contract name.
+    pub name: String,
+    /// Contract ABI.
+    pub abi: JsonAbi,
+    /// Contract creation code.
+    pub bytecode: Bytes,
+    /// Contract runtime code.
+    pub deployed_bytecode: Bytes,
+}
 
-/// Wrapper type that maps an artifact to a contract ABI and bytecode.
-#[derive(Clone, Default)]
-pub struct ContractsByArtifact(pub BTreeMap<ArtifactId, (JsonAbi, Vec<u8>)>);
+type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
 
-impl fmt::Debug for ContractsByArtifact {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_map().entries(self.iter().map(|(k, (v1, v2))| (k, (v1, hex::encode(v2))))).finish()
-    }
-}
+/// Wrapper type that maps an artifact to a contract ABI and bytecode.
+#[derive(Clone, Default, Debug)]
+pub struct ContractsByArtifact(pub BTreeMap<ArtifactId, ContractData>);
 
 impl ContractsByArtifact {
     /// Finds a contract which has a similar bytecode as `code`.
-    pub fn find_by_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
-        self.iter().find(|(_, (_, known_code))| bytecode_diff_score(known_code, code) <= 0.1)
+    pub fn find_by_creation_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
+        self.iter()
+            .find(|(_, contract)| bytecode_diff_score(contract.bytecode.as_ref(), code) <= 0.1)
+    }
+
+    /// Finds a contract which has a similar deployed bytecode as `code`.
+    pub fn find_by_deployed_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
+        self.iter().find(|(_, contract)| {
+            bytecode_diff_score(contract.deployed_bytecode.as_ref(), code) <= 0.1
+        })
     }
 
     /// Finds a contract which has the same contract name or identifier as `id`. If more than one is
@@ -51,14 +65,14 @@ impl ContractsByArtifact {
         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() {
+        for (_name, contract) in self.iter() {
+            for func in contract.abi.functions() {
                 funcs.insert(func.selector(), func.clone());
             }
-            for event in abi.events() {
+            for event in contract.abi.events() {
                 events.insert(event.selector(), event.clone());
             }
-            for error in abi.errors() {
+            for error in contract.abi.errors() {
                 errors_abi.errors.entry(error.name.clone()).or_default().push(error.clone());
             }
         }
@@ -67,7 +81,7 @@ impl ContractsByArtifact {
 }
 
 impl Deref for ContractsByArtifact {
-    type Target = BTreeMap<ArtifactId, (JsonAbi, Vec<u8>)>;
+    type Target = BTreeMap<ArtifactId, ContractData>;
 
     fn deref(&self) -> &Self::Target {
         &self.0

crates/evm/coverage/src/analysis.rs

@@ -45,13 +45,10 @@ impl<'a> ContractVisitor<'a> {
         let name: String =
             node.attribute("name").ok_or_else(|| eyre::eyre!("Function has no name"))?;
 
-        // TODO(onbjerg): Re-enable constructor parsing when we walk both the deployment and runtime
-        // sourcemaps. Currently this fails because we are trying to look for anchors in the runtime
-        // sourcemap.
         // TODO(onbjerg): Figure out why we cannot find anchors for the receive function
         let kind: String =
             node.attribute("kind").ok_or_else(|| eyre::eyre!("Function has no kind"))?;
-        if kind == "constructor" || kind == "receive" {
+        if kind == "receive" {
             return Ok(())
         }
 

crates/evm/coverage/src/lib.rs

@@ -37,7 +37,7 @@ pub struct CoverageReport {
     /// All coverage items for the codebase, keyed by the compiler version.
     pub items: HashMap<Version, Vec<CoverageItem>>,
     /// All item anchors for the codebase, keyed by their contract ID.
-    pub anchors: HashMap<ContractId, Vec<ItemAnchor>>,
+    pub anchors: HashMap<ContractId, (Vec<ItemAnchor>, Vec<ItemAnchor>)>,
     /// All the bytecode hits for the codebase
     pub bytecode_hits: HashMap<ContractId, HitMap>,
     /// The bytecode -> source mappings
@@ -70,7 +70,10 @@ impl CoverageReport {
     }
 
     /// Add anchors to this report
-    pub fn add_anchors(&mut self, anchors: HashMap<ContractId, Vec<ItemAnchor>>) {
+    pub fn add_anchors(
+        &mut self,
+        anchors: HashMap<ContractId, (Vec<ItemAnchor>, Vec<ItemAnchor>)>,
+    ) {
         self.anchors.extend(anchors);
     }
 
@@ -124,7 +127,12 @@ impl CoverageReport {
     ///
     /// This function should only be called *after* all the relevant sources have been processed and
     /// added to the map (see [add_source]).
-    pub fn add_hit_map(&mut self, contract_id: &ContractId, hit_map: &HitMap) -> Result<()> {
+    pub fn add_hit_map(
+        &mut self,
+        contract_id: &ContractId,
+        hit_map: &HitMap,
+        is_deployed_code: bool,
+    ) -> Result<()> {
         // Add bytecode level hits
         let e = self
             .bytecode_hits
@@ -139,6 +147,7 @@ impl CoverageReport {
 
         // Add source level hits
         if let Some(anchors) = self.anchors.get(contract_id) {
+            let anchors = if is_deployed_code { &anchors.1 } else { &anchors.0 };
             for anchor in anchors {
                 if let Some(hits) = hit_map.hits.get(&anchor.instruction) {
                     self.items

crates/evm/evm/src/executors/invariant/mod.rs

@@ -435,8 +435,9 @@ impl<'a> InvariantExecutor<'a> {
         }
 
         // Exclude any artifact without mutable functions.
-        for (artifact, (abi, _)) in self.project_contracts.iter() {
-            if abi
+        for (artifact, contract) in self.project_contracts.iter() {
+            if contract
+                .abi
                 .functions()
                 .filter(|func| {
                     !matches!(
@@ -474,12 +475,14 @@ impl<'a> InvariantExecutor<'a> {
         contract: String,
         selectors: &[FixedBytes<4>],
     ) -> eyre::Result<String> {
-        if let Some((artifact, (abi, _))) =
+        if let Some((artifact, contract_data)) =
             self.project_contracts.find_by_name_or_identifier(&contract)?
         {
             // Check that the selectors really exist for this contract.
             for selector in selectors {
-                abi.functions()
+                contract_data
+                    .abi
+                    .functions()
                     .find(|func| func.selector().as_slice() == selector.as_slice())
                     .wrap_err(format!("{contract} does not have the selector {selector:?}"))?;
             }
@@ -563,7 +566,7 @@ impl<'a> InvariantExecutor<'a> {
             // Identifiers are specified as an array, so we loop through them.
             for identifier in artifacts {
                 // Try to find the contract by name or identifier in the project's contracts.
-                if let Some((_, (abi, _))) =
+                if let Some((_, contract)) =
                     self.project_contracts.find_by_name_or_identifier(identifier)?
                 {
                     combined
@@ -574,10 +577,10 @@ impl<'a> InvariantExecutor<'a> {
                             let (_, contract_abi, _) = entry;
 
                             // Extend the ABI's function list with the new functions.
-                            contract_abi.functions.extend(abi.functions.clone());
+                            contract_abi.functions.extend(contract.abi.functions.clone());
                         })
                         // Otherwise insert it into the map.
-                        .or_insert_with(|| (identifier.to_string(), abi.clone(), vec![]));
+                        .or_insert_with(|| (identifier.to_string(), contract.abi.clone(), vec![]));
                 }
             }
         }

crates/evm/fuzz/src/strategies/state.rs

@@ -316,15 +316,17 @@ pub fn collect_created_contracts(
         if !setup_contracts.contains_key(address) {
             if let (true, Some(code)) = (&account.is_touched(), &account.info.code) {
                 if !code.is_empty() {
-                    if let Some((artifact, (abi, _))) =
-                        project_contracts.find_by_code(&code.original_bytes())
+                    if let Some((artifact, contract)) =
+                        project_contracts.find_by_deployed_code(&code.original_bytes())
                     {
                         if let Some(functions) =
-                            artifact_filters.get_targeted_functions(artifact, abi)?
+                            artifact_filters.get_targeted_functions(artifact, &contract.abi)?
                         {
                             created_contracts.push(*address);
-                            writable_targeted
-                                .insert(*address, (artifact.name.clone(), abi.clone(), functions));
+                            writable_targeted.insert(
+                                *address,
+                                (artifact.name.clone(), contract.abi.clone(), functions),
+                            );
                         }
                     }
                 }

crates/evm/traces/src/decoder/mod.rs

@@ -56,8 +56,8 @@ impl CallTraceDecoderBuilder {
     #[inline]
     pub fn with_known_contracts(mut self, contracts: &ContractsByArtifact) -> Self {
         trace!(target: "evm::traces", len=contracts.len(), "collecting known contract ABIs");
-        for (abi, _) in contracts.values() {
-            self.decoder.collect_abi(abi, None);
+        for contract in contracts.values() {
+            self.decoder.collect_abi(&contract.abi, None);
         }
         self
     }

crates/evm/traces/src/identifier/local.rs

@@ -9,16 +9,18 @@ use std::borrow::Cow;
 pub struct LocalTraceIdentifier<'a> {
     /// Known contracts to search through.
     known_contracts: &'a ContractsByArtifact,
-    /// Vector of pairs of artifact ID and the code length of the given artifact.
+    /// Vector of pairs of artifact ID and the runtime code length of the given artifact.
     ordered_ids: Vec<(&'a ArtifactId, usize)>,
 }
 
 impl<'a> LocalTraceIdentifier<'a> {
     /// Creates a new local trace identifier.
     #[inline]
     pub fn new(known_contracts: &'a ContractsByArtifact) -> Self {
-        let mut ordered_ids =
-            known_contracts.iter().map(|(id, contract)| (id, contract.1.len())).collect::<Vec<_>>();
+        let mut ordered_ids = known_contracts
+            .iter()
+            .map(|(id, contract)| (id, contract.deployed_bytecode.len()))
+            .collect::<Vec<_>>();
         ordered_ids.sort_by_key(|(_, len)| *len);
         Self { known_contracts, ordered_ids }
     }
@@ -37,15 +39,15 @@ impl<'a> LocalTraceIdentifier<'a> {
         let mut min_score_id = None;
 
         let mut check = |id| {
-            let (abi, known_code) = self.known_contracts.get(id)?;
-            let score = bytecode_diff_score(known_code, code);
+            let contract = self.known_contracts.get(id)?;
+            let score = bytecode_diff_score(&contract.deployed_bytecode, code);
             if score == 0.0 {
                 trace!(target: "evm::traces", "found exact match");
-                return Some((id, abi));
+                return Some((id, &contract.abi));
             }
             if score < min_score {
                 min_score = score;
-                min_score_id = Some((id, abi));
+                min_score_id = Some((id, &contract.abi));
             }
             None
         };

crates/evm/traces/src/lib.rs

@@ -308,8 +308,8 @@ pub fn load_contracts(
         .contracts
         .iter()
         .filter_map(|(addr, name)| {
-            if let Ok(Some((_, (abi, _)))) = contracts.find_by_name_or_identifier(name) {
-                return Some((*addr, (name.clone(), abi.clone())));
+            if let Ok(Some((_, contract))) = contracts.find_by_name_or_identifier(name) {
+                return Some((*addr, (name.clone(), contract.abi.clone())));
             }
             None
         })

crates/forge/bin/cmd/coverage.rs

@@ -5,7 +5,7 @@ use eyre::{Context, Result};
 use forge::{
     coverage::{
         analysis::SourceAnalyzer, anchors::find_anchors, BytecodeReporter, ContractId,
-        CoverageReport, CoverageReporter, DebugReporter, ItemAnchor, LcovReporter, SummaryReporter,
+        CoverageReport, CoverageReporter, DebugReporter, LcovReporter, SummaryReporter,
     },
     inspectors::CheatsConfig,
     opts::EvmOpts,
@@ -271,21 +271,26 @@ impl CoverageArgs {
                 items_by_source_id.entry(item.loc.source_id).or_default().push(item_id);
             }
 
-            let anchors: HashMap<ContractId, Vec<ItemAnchor>> = source_maps
+            let anchors = source_maps
                 .iter()
                 .filter(|(contract_id, _)| contract_id.version == version)
-                .filter_map(|(contract_id, (_, deployed_source_map))| {
+                .filter_map(|(contract_id, (creation_source_map, deployed_source_map))| {
+                    let creation_code_anchors = find_anchors(
+                        &bytecodes.get(contract_id)?.0,
+                        creation_source_map,
+                        &ic_pc_maps.get(contract_id)?.0,
+                        &source_analysis.items,
+                        &items_by_source_id,
+                    );
+                    let deployed_code_anchors = find_anchors(
+                        &bytecodes.get(contract_id)?.1,
+                        deployed_source_map,
+                        &ic_pc_maps.get(contract_id)?.1,
+                        &source_analysis.items,
+                        &items_by_source_id,
+                    );
                     // TODO: Creation source map/bytecode as well
-                    Some((
-                        contract_id.clone(),
-                        find_anchors(
-                            &bytecodes.get(contract_id)?.1,
-                            deployed_source_map,
-                            &ic_pc_maps.get(contract_id)?.1,
-                            &source_analysis.items,
-                            &items_by_source_id,
-                        ),
-                    ))
+                    Some((contract_id.clone(), (creation_code_anchors, deployed_code_anchors)))
                 })
                 .collect();
             report.add_items(version, source_analysis.items);
@@ -345,24 +350,34 @@ impl CoverageArgs {
             .filter_map(|mut result| result.coverage.take())
             .flat_map(|hit_maps| {
                 hit_maps.0.into_values().filter_map(|map| {
-                    Some((known_contracts.find_by_code(map.bytecode.as_ref())?.0, map))
+                    if let Some((id, _)) =
+                        known_contracts.find_by_deployed_code(map.bytecode.as_ref())
+                    {
+                        Some((id, map, true))
+                    } else if let Some((id, _)) =
+                        known_contracts.find_by_creation_code(map.bytecode.as_ref())
+                    {
+                        Some((id, map, false))
+                    } else {
+                        None
+                    }
                 })
             });
-        for (artifact_id, hits) in data {
+        for (artifact_id, hits, is_deployed_code) in data {
             // TODO: Note down failing tests
             if let Some(source_id) = report.get_source_id(
                 artifact_id.version.clone(),
                 artifact_id.source.to_string_lossy().to_string(),
             ) {
                 let source_id = *source_id;
-                // TODO: Distinguish between creation/runtime in a smart way
                 report.add_hit_map(
                     &ContractId {
                         version: artifact_id.version.clone(),
                         source_id,
                         contract_name: artifact_id.name.clone(),
                     },
                     &hits,
+                    is_deployed_code,
                 )?;
             }
         }