Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: coverage for libraries #7510

Merged
merged 10 commits into from
Mar 29, 2024
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 33 additions & 169 deletions crates/evm/coverage/src/analysis.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -23,36 +23,14 @@ pub struct ContractVisitor<'a> {

/// Coverage items
pub items: Vec<CoverageItem>,

/// Node IDs of this contract's base contracts, as well as IDs for referenced contracts such as
/// libraries
pub base_contract_node_ids: HashSet<usize>,
}

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<Self> {
let linearized_base_contracts: Vec<usize> =
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 {
Expand Down Expand Up @@ -318,25 +296,13 @@ impl<'a> ContractVisitor<'a> {
});

let expr: Option<Node> = 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<String> = 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<String> = 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(())
Expand Down Expand Up @@ -435,9 +401,6 @@ impl<'a> ContractVisitor<'a> {
pub struct SourceAnalysis {
/// A collection of coverage items.
pub items: Vec<CoverageItem>,
/// A mapping of contract IDs to item IDs relevant to the contract (including items in base
/// contracts).
pub contract_items: HashMap<ContractId, Vec<usize>>,
}

/// Analyzes a set of sources to find coverage items.
Expand All @@ -451,10 +414,6 @@ pub struct SourceAnalyzer {
contracts: HashMap<ContractId, Node>,
/// A collection of coverage items.
items: Vec<CoverageItem>,
/// A map of contract IDs to item IDs.
contract_items: HashMap<ContractId, Vec<usize>>,
/// A map of contracts to their base contracts
contract_bases: HashMap<ContractId, Vec<ContractId>>,
}

impl SourceAnalyzer {
Expand Down Expand Up @@ -497,11 +456,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
Expand All @@ -511,127 +466,36 @@ 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<SourceAnalysis> {
// Analyze the contracts
self.analyze_contracts()?;

// Flatten the data
let mut flattened: HashMap<ContractId, Vec<usize>> = HashMap::new();
for contract_id in self.contract_items.keys() {
let mut item_ids: Vec<usize> = 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());
let ContractVisitor { items, .. } = ContractVisitor::new(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

w/o the if !self.contract_items.contains_key(contract_id) { check, will this not lead to re-analyzing contracts? and possibly duping coverage items?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We're walking only ContractDefinition nodes from self.contracts, so it should be fine as long as there is no duplicating ContractId keys

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think thats a fair assumption

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)
mattsse marked this conversation as resolved.
Show resolved Hide resolved
}),
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 {
let item_ids: Vec<usize> =
(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 })
}
}
17 changes: 13 additions & 4 deletions crates/evm/coverage/src/anchors.rs
Original file line number Diff line number Diff line change
@@ -1,24 +1,33 @@
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.
pub fn find_anchors(
bytecode: &Bytes,
source_map: &SourceMap,
ic_pc_map: &IcPcMap,
item_ids: &[usize],
items: &[CoverageItem],
items_by_source_id: &HashMap<usize, Vec<usize>>,
) -> Vec<ItemAnchor> {
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::<HashSet<_>>();

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, .. } => {
Expand Down
24 changes: 19 additions & 5 deletions crates/forge/bin/cmd/coverage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -257,19 +257,33 @@ impl CoverageArgs {
})?,
)?
.analyze()?;
let anchors: HashMap<ContractId, Vec<ItemAnchor>> = source_analysis
.contract_items

// Build helper mapping used by `find_anchors`
let items_by_source_id = source_analysis
.items
.iter()
.enumerate()
.map(|(item_id, item)| (item.loc.source_id, item_id))
.fold(
HashMap::new(),
|mut map: HashMap<usize, Vec<usize>>, (source_id, item_id)| {
map.entry(source_id).or_default().push(item_id);
map
},
);
mattsse marked this conversation as resolved.
Show resolved Hide resolved

let anchors: HashMap<ContractId, Vec<ItemAnchor>> = 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,
),
))
})
Expand Down
Loading