diff --git a/crates/evm/coverage/src/analysis.rs b/crates/evm/coverage/src/analysis.rs index 1926d2b4ed20..a136f2d79cd2 100644 --- a/crates/evm/coverage/src/analysis.rs +++ b/crates/evm/coverage/src/analysis.rs @@ -3,7 +3,7 @@ use foundry_common::TestFunctionExt; use foundry_compilers::artifacts::ast::{self, Ast, Node, NodeType}; use rustc_hash::FxHashMap; use semver::Version; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; /// A visitor that walks the AST of a single contract and finds coverage items. #[derive(Clone, Debug)] @@ -23,36 +23,14 @@ pub struct ContractVisitor<'a> { /// Coverage items pub items: Vec, - - /// Node IDs of this contract's base contracts, as well as IDs for referenced contracts such as - /// libraries - pub base_contract_node_ids: HashSet, } impl<'a> ContractVisitor<'a> { pub fn new(source_id: usize, source: &'a str, contract_name: String) -> Self { - Self { - source_id, - source, - contract_name, - branch_id: 0, - last_line: 0, - items: Vec::new(), - base_contract_node_ids: HashSet::new(), - } + Self { source_id, source, contract_name, branch_id: 0, last_line: 0, items: Vec::new() } } pub fn visit(mut self, contract_ast: Node) -> eyre::Result { - let linearized_base_contracts: Vec = - contract_ast.attribute("linearizedBaseContracts").ok_or_else(|| { - eyre::eyre!( - "The contract's AST node is missing a list of linearized base contracts" - ) - })?; - - // We skip the first ID because that's the ID of the contract itself - self.base_contract_node_ids.extend(&linearized_base_contracts[1..]); - // Find all functions and walk their AST for node in contract_ast.nodes { if node.node_type == NodeType::FunctionDefinition { @@ -318,25 +296,13 @@ impl<'a> ContractVisitor<'a> { }); let expr: Option = node.attribute("expression"); - match expr.as_ref().map(|expr| &expr.node_type) { + if let Some(NodeType::Identifier) = expr.as_ref().map(|expr| &expr.node_type) { // Might be a require/assert call - Some(NodeType::Identifier) => { - let name: Option = expr.and_then(|expr| expr.attribute("name")); - if let Some("assert" | "require") = name.as_deref() { - self.push_branches(&node.src, self.branch_id); - self.branch_id += 1; - } - } - // Might be a call to a library - Some(NodeType::MemberAccess) => { - let referenced_declaration_id = expr - .and_then(|expr| expr.attribute("expression")) - .and_then(|subexpr: Node| subexpr.attribute("referencedDeclaration")); - if let Some(id) = referenced_declaration_id { - self.base_contract_node_ids.insert(id); - } + let name: Option = expr.and_then(|expr| expr.attribute("name")); + if let Some("assert" | "require") = name.as_deref() { + self.push_branches(&node.src, self.branch_id); + self.branch_id += 1; } - _ => (), } Ok(()) @@ -435,9 +401,6 @@ impl<'a> ContractVisitor<'a> { pub struct SourceAnalysis { /// A collection of coverage items. pub items: Vec, - /// A mapping of contract IDs to item IDs relevant to the contract (including items in base - /// contracts). - pub contract_items: HashMap>, } /// Analyzes a set of sources to find coverage items. @@ -445,16 +408,10 @@ pub struct SourceAnalysis { pub struct SourceAnalyzer { /// A map of source IDs to their source code sources: FxHashMap, - /// A map of AST node IDs of contracts to their contract IDs. - contract_ids: FxHashMap, /// A map of contract IDs to their AST nodes. contracts: HashMap, /// A collection of coverage items. items: Vec, - /// A map of contract IDs to item IDs. - contract_items: HashMap>, - /// A map of contracts to their base contracts - contract_bases: HashMap>, } impl SourceAnalyzer { @@ -476,8 +433,6 @@ impl SourceAnalyzer { continue } - let node_id = - child.id.ok_or_else(|| eyre::eyre!("The contract's AST node has no ID"))?; let contract_id = ContractId { version: version.clone(), source_id, @@ -485,7 +440,6 @@ impl SourceAnalyzer { .attribute("name") .ok_or_else(|| eyre::eyre!("Contract has no name"))?, }; - analyzer.contract_ids.insert(node_id, contract_id.clone()); analyzer.contracts.insert(contract_id, child); } } @@ -497,11 +451,7 @@ impl SourceAnalyzer { /// /// Coverage items are found by: /// - Walking the AST of each contract (except interfaces) - /// - Recording the items and base contracts of each contract - /// - /// Finally, the item IDs of each contract and its base contracts are flattened, and the return - /// value represents all coverage items in the project, along with a mapping of contract IDs to - /// item IDs. + /// - Recording the items of each contract /// /// Each coverage item contains relevant information to find opcodes corresponding to them: the /// source ID the item is in, the source code range of the item, and the contract name the item @@ -511,127 +461,32 @@ impl SourceAnalyzer { /// two different solc versions will produce overlapping source IDs if the compiler version is /// not taken into account. pub fn analyze(mut self) -> eyre::Result { - // Analyze the contracts - self.analyze_contracts()?; - - // Flatten the data - let mut flattened: HashMap> = HashMap::new(); - for contract_id in self.contract_items.keys() { - let mut item_ids: Vec = Vec::new(); - - // - // for a specific contract (id == contract_id): - // - // self.contract_bases.get(contract_id) includes the following contracts: - // 1. all the ancestors of this contract (including parent, grandparent, ... - // contracts) - // 2. the libraries **directly** used by this contract - // - // The missing contracts are: - // 1. libraries used in ancestors of this contracts - // 2. libraries used in libraries (i.e libs indirectly used by this contract) - // - // We want to find out all the above contracts and libraries related to this contract. - - for contract_or_lib in { - // A set of contracts and libraries related to this contract (will include "this" - // contract itself) - let mut contracts_libraries: HashSet<&ContractId> = HashSet::new(); - - // we use a stack for depth-first search. - let mut stack: Vec<&ContractId> = Vec::new(); - - // push "this" contract onto the stack - stack.push(contract_id); - - while let Some(contract_or_lib) = stack.pop() { - // whenever a contract_or_lib is removed from the stack, it is added to the set - contracts_libraries.insert(contract_or_lib); - - // push all ancestors of contract_or_lib and libraries used by contract_or_lib - // onto the stack - if let Some(bases) = self.contract_bases.get(contract_or_lib) { - stack.extend( - bases.iter().filter(|base| !contracts_libraries.contains(base)), - ); - } - } - - contracts_libraries - } { - // get items of each contract or library - if let Some(items) = self.contract_items.get(contract_or_lib) { - item_ids.extend(items.iter()); - } - } - - // If there are no items for this contract, then it was most likely filtered - if !item_ids.is_empty() { - flattened.insert(contract_id.clone(), item_ids); - } - } - - Ok(SourceAnalysis { items: self.items.clone(), contract_items: flattened }) - } - - fn analyze_contracts(&mut self) -> eyre::Result<()> { - for contract_id in self.contracts.keys() { - // Find this contract's coverage items if we haven't already - if !self.contract_items.contains_key(contract_id) { - let ContractVisitor { items, base_contract_node_ids, .. } = ContractVisitor::new( - contract_id.source_id, - self.sources.get(&contract_id.source_id).unwrap_or_else(|| { - panic!( - "We should have the source code for source ID {}", - contract_id.source_id - ) - }), - contract_id.contract_name.clone(), - ) - .visit( - self.contracts - .get(contract_id) - .unwrap_or_else(|| { - panic!("We should have the AST of contract: {contract_id:?}") - }) - .clone(), - )?; - - let is_test = items.iter().any(|item| { - if let CoverageItemKind::Function { name } = &item.kind { - name.is_test() - } else { - false - } - }); - - // Record this contract's base contracts - // We don't do this for test contracts because we don't care about them - if !is_test { - self.contract_bases.insert( - contract_id.clone(), - base_contract_node_ids - .iter() - .filter_map(|base_contract_node_id| { - self.contract_ids.get(base_contract_node_id).cloned() - }) - .collect(), - ); - } - - // For tests and contracts with no items we still record an empty Vec so we don't - // end up here again - if items.is_empty() || is_test { - self.contract_items.insert(contract_id.clone(), Vec::new()); + for (contract_id, ast) in self.contracts { + let ContractVisitor { items, .. } = ContractVisitor::new( + contract_id.source_id, + self.sources.get(&contract_id.source_id).ok_or_else(|| { + eyre::eyre!( + "We should have the source code for source ID {}", + contract_id.source_id + ) + })?, + contract_id.contract_name.clone(), + ) + .visit(ast)?; + + let is_test = items.iter().any(|item| { + if let CoverageItemKind::Function { name } = &item.kind { + name.is_test() } else { - let item_ids: Vec = - (self.items.len()..self.items.len() + items.len()).collect(); - self.items.extend(items); - self.contract_items.insert(contract_id.clone(), item_ids.clone()); + false } + }); + + if !is_test { + self.items.extend(items); } } - Ok(()) + Ok(SourceAnalysis { items: self.items }) } } diff --git a/crates/evm/coverage/src/anchors.rs b/crates/evm/coverage/src/anchors.rs index 237d941d1a3d..d45fdae57dca 100644 --- a/crates/evm/coverage/src/anchors.rs +++ b/crates/evm/coverage/src/anchors.rs @@ -1,10 +1,12 @@ +use std::collections::HashMap; + use super::{CoverageItem, CoverageItemKind, ItemAnchor, SourceLocation}; use alloy_primitives::Bytes; use foundry_compilers::sourcemap::{SourceElement, SourceMap}; use foundry_evm_core::utils::IcPcMap; use revm::{ interpreter::opcode::{self, spec_opcode_gas}, - primitives::SpecId, + primitives::{HashSet, SpecId}, }; /// Attempts to find anchors for the given items using the given source map and bytecode. @@ -12,13 +14,20 @@ pub fn find_anchors( bytecode: &Bytes, source_map: &SourceMap, ic_pc_map: &IcPcMap, - item_ids: &[usize], items: &[CoverageItem], + items_by_source_id: &HashMap>, ) -> Vec { - item_ids + // Prepare coverage items from all sources referenced in the source map + let potential_item_ids = source_map .iter() + .filter_map(|element| items_by_source_id.get(&(element.index? as usize))) + .flatten() + .collect::>(); + + potential_item_ids + .into_iter() .filter_map(|item_id| { - let item = items.get(*item_id)?; + let item = &items[*item_id]; match item.kind { CoverageItemKind::Branch { path_id, .. } => { diff --git a/crates/forge/bin/cmd/coverage.rs b/crates/forge/bin/cmd/coverage.rs index e140b8436aaf..c59a54ade72f 100644 --- a/crates/forge/bin/cmd/coverage.rs +++ b/crates/forge/bin/cmd/coverage.rs @@ -257,19 +257,27 @@ impl CoverageArgs { })?, )? .analyze()?; - let anchors: HashMap> = source_analysis - .contract_items + + // Build helper mapping used by `find_anchors` + let mut items_by_source_id: HashMap<_, Vec<_>> = + HashMap::with_capacity(source_analysis.items.len()); + + for (item_id, item) in source_analysis.items.iter().enumerate() { + items_by_source_id.entry(item.loc.source_id).or_default().push(item_id); + } + + let anchors: HashMap> = source_maps .iter() - .filter_map(|(contract_id, item_ids)| { + .filter_map(|(contract_id, (_, deployed_source_map))| { // TODO: Creation source map/bytecode as well Some(( contract_id.clone(), find_anchors( &bytecodes.get(contract_id)?.1, - &source_maps.get(contract_id)?.1, + deployed_source_map, &ic_pc_maps.get(contract_id)?.1, - item_ids, &source_analysis.items, + &items_by_source_id, ), )) })