1use super::{CoverageItem, CoverageItemKind, SourceLocation};
2use alloy_primitives::map::HashMap;
3use foundry_common::TestFunctionExt;
4use foundry_compilers::artifacts::{
5 Source,
6 ast::{self, Ast, Node, NodeType},
7};
8use rayon::prelude::*;
9use std::sync::Arc;
10
11#[derive(Clone, Debug)]
13struct ContractVisitor<'a> {
14 source_id: u32,
16 source: &'a str,
18
19 contract_name: &'a Arc<str>,
21
22 branch_id: u32,
24 last_line: u32,
26
27 items: Vec<CoverageItem>,
29}
30
31impl<'a> ContractVisitor<'a> {
32 fn new(source_id: usize, source: &'a str, contract_name: &'a Arc<str>) -> Self {
33 Self {
34 source_id: source_id.try_into().expect("too many sources"),
35 source,
36 contract_name,
37 branch_id: 0,
38 last_line: 0,
39 items: Vec::new(),
40 }
41 }
42
43 fn clear_if_test(&mut self) {
45 let has_tests = self.items.iter().any(|item| {
46 if let CoverageItemKind::Function { name } = &item.kind {
47 name.is_any_test()
48 } else {
49 false
50 }
51 });
52 if has_tests {
53 self.items = Vec::new();
54 }
55 }
56
57 fn disambiguate_functions(&mut self) {
59 if self.items.is_empty() {
60 return;
61 }
62
63 let mut dups = HashMap::<_, Vec<usize>>::default();
64 for (i, item) in self.items.iter().enumerate() {
65 if let CoverageItemKind::Function { name } = &item.kind {
66 dups.entry(name.clone()).or_default().push(i);
67 }
68 }
69 for dups in dups.values() {
70 if dups.len() > 1 {
71 for (i, &dup) in dups.iter().enumerate() {
72 let item = &mut self.items[dup];
73 if let CoverageItemKind::Function { name } = &item.kind {
74 item.kind =
75 CoverageItemKind::Function { name: format!("{name}.{i}").into() };
76 }
77 }
78 }
79 }
80 }
81
82 fn visit_contract(&mut self, node: &Node) -> eyre::Result<()> {
83 for node in &node.nodes {
85 match node.node_type {
86 NodeType::FunctionDefinition => {
87 self.visit_function_definition(node)?;
88 }
89 NodeType::ModifierDefinition => {
90 self.visit_modifier_or_yul_fn_definition(node)?;
91 }
92 _ => {}
93 }
94 }
95 Ok(())
96 }
97
98 fn visit_function_definition(&mut self, node: &Node) -> eyre::Result<()> {
99 let Some(body) = &node.body else { return Ok(()) };
100
101 let name: Box<str> =
102 node.attribute("name").ok_or_else(|| eyre::eyre!("Function has no name"))?;
103 let kind: Box<str> =
104 node.attribute("kind").ok_or_else(|| eyre::eyre!("Function has no kind"))?;
105
106 if &*kind != "function" && !has_statements(body) {
109 return Ok(());
110 }
111
112 let name = if name.is_empty() { kind } else { name };
115
116 self.push_item_kind(CoverageItemKind::Function { name }, &node.src);
117 self.visit_block(body)
118 }
119
120 fn visit_modifier_or_yul_fn_definition(&mut self, node: &Node) -> eyre::Result<()> {
121 let Some(body) = &node.body else { return Ok(()) };
122
123 let name: Box<str> =
124 node.attribute("name").ok_or_else(|| eyre::eyre!("Modifier has no name"))?;
125 self.push_item_kind(CoverageItemKind::Function { name }, &node.src);
126 self.visit_block(body)
127 }
128
129 fn visit_block(&mut self, node: &Node) -> eyre::Result<()> {
130 let statements: Vec<Node> = node.attribute("statements").unwrap_or_default();
131
132 for statement in &statements {
133 self.visit_statement(statement)?;
134 }
135
136 Ok(())
137 }
138
139 fn visit_statement(&mut self, node: &Node) -> eyre::Result<()> {
140 match node.node_type {
141 NodeType::Block | NodeType::UncheckedBlock | NodeType::YulBlock => {
143 self.visit_block(node)
144 }
145 NodeType::InlineAssembly => self.visit_block(
147 &node
148 .attribute("AST")
149 .ok_or_else(|| eyre::eyre!("inline assembly block with no AST attribute"))?,
150 ),
151 NodeType::Break
153 | NodeType::Continue
154 | NodeType::EmitStatement
155 | NodeType::RevertStatement
156 | NodeType::YulAssignment
157 | NodeType::YulBreak
158 | NodeType::YulContinue
159 | NodeType::YulLeave
160 | NodeType::YulVariableDeclaration => {
161 self.push_item_kind(CoverageItemKind::Statement, &node.src);
162 Ok(())
163 }
164 NodeType::PlaceholderStatement => Ok(()),
166 NodeType::Return => {
168 self.push_item_kind(CoverageItemKind::Statement, &node.src);
169 if let Some(expr) = node.attribute("expression") {
170 self.visit_expression(&expr)?;
171 }
172 Ok(())
173 }
174 NodeType::VariableDeclarationStatement => {
176 self.push_item_kind(CoverageItemKind::Statement, &node.src);
177 if let Some(expr) = node.attribute("initialValue") {
178 self.visit_expression(&expr)?;
179 }
180 Ok(())
181 }
182 NodeType::DoWhileStatement | NodeType::WhileStatement => {
184 self.visit_expression(
185 &node
186 .attribute("condition")
187 .ok_or_else(|| eyre::eyre!("while statement had no condition"))?,
188 )?;
189
190 let body = node
191 .body
192 .as_deref()
193 .ok_or_else(|| eyre::eyre!("while statement had no body node"))?;
194 self.visit_block_or_statement(body)
195 }
196 NodeType::ForStatement => {
198 if let Some(stmt) = node.attribute("initializationExpression") {
199 self.visit_statement(&stmt)?;
200 }
201 if let Some(expr) = node.attribute("condition") {
202 self.visit_expression(&expr)?;
203 }
204 if let Some(stmt) = node.attribute("loopExpression") {
205 self.visit_statement(&stmt)?;
206 }
207
208 let body = node
209 .body
210 .as_deref()
211 .ok_or_else(|| eyre::eyre!("for statement had no body node"))?;
212 self.visit_block_or_statement(body)
213 }
214 NodeType::ExpressionStatement | NodeType::YulExpressionStatement => self
216 .visit_expression(
217 &node
218 .attribute("expression")
219 .ok_or_else(|| eyre::eyre!("expression statement had no expression"))?,
220 ),
221 NodeType::IfStatement => {
223 self.visit_expression(
224 &node
225 .attribute("condition")
226 .ok_or_else(|| eyre::eyre!("if statement had no condition"))?,
227 )?;
228
229 let true_body: Node = node
230 .attribute("trueBody")
231 .ok_or_else(|| eyre::eyre!("if statement had no true body"))?;
232
233 let branch_id = self.branch_id;
236
237 self.branch_id += 1;
240
241 match node.attribute::<Node>("falseBody") {
242 Some(false_body) => {
244 if has_statements(&true_body) || has_statements(&false_body) {
247 self.push_item_kind(
250 CoverageItemKind::Branch {
251 branch_id,
252 path_id: 0,
253 is_first_opcode: true,
254 },
255 &true_body.src,
256 );
257 self.push_item_kind(
261 CoverageItemKind::Branch {
262 branch_id,
263 path_id: 1,
264 is_first_opcode: false,
265 },
266 &ast::LowFidelitySourceLocation {
267 start: node.src.start,
268 length: false_body.src.length.map(|length| {
269 false_body.src.start - true_body.src.start + length
270 }),
271 index: node.src.index,
272 },
273 );
274
275 self.visit_block_or_statement(&true_body)?;
277 self.visit_block_or_statement(&false_body)?;
279 }
280 }
281 None => {
282 if has_statements(&true_body) {
284 self.push_item_kind(
286 CoverageItemKind::Branch {
287 branch_id,
288 path_id: 0,
289 is_first_opcode: true,
290 },
291 &true_body.src,
292 );
293 self.visit_block_or_statement(&true_body)?;
295 }
296 }
297 }
298
299 Ok(())
300 }
301 NodeType::YulIf => {
302 self.visit_expression(
303 &node
304 .attribute("condition")
305 .ok_or_else(|| eyre::eyre!("yul if statement had no condition"))?,
306 )?;
307 let body = node
308 .body
309 .as_deref()
310 .ok_or_else(|| eyre::eyre!("yul if statement had no body"))?;
311
312 let branch_id = self.branch_id;
315
316 self.branch_id += 1;
319
320 self.push_item_kind(
321 CoverageItemKind::Branch { branch_id, path_id: 0, is_first_opcode: false },
322 &node.src,
323 );
324 self.visit_block(body)?;
325
326 Ok(())
327 }
328 NodeType::TryStatement => {
331 self.visit_expression(
332 &node
333 .attribute("externalCall")
334 .ok_or_else(|| eyre::eyre!("try statement had no call"))?,
335 )?;
336
337 let branch_id = self.branch_id;
338 self.branch_id += 1;
339
340 let mut clauses = node
341 .attribute::<Vec<Node>>("clauses")
342 .ok_or_else(|| eyre::eyre!("try statement had no clauses"))?;
343
344 let try_block = clauses
345 .remove(0)
346 .attribute::<Node>("block")
347 .ok_or_else(|| eyre::eyre!("try statement had no block"))?;
348 self.push_item_kind(
350 CoverageItemKind::Branch { branch_id, path_id: 0, is_first_opcode: true },
351 &ast::LowFidelitySourceLocation {
352 start: node.src.start,
353 length: try_block
354 .src
355 .length
356 .map(|length| try_block.src.start + length - node.src.start),
357 index: node.src.index,
358 },
359 );
360 self.visit_block(&try_block)?;
361
362 let mut path_id = 1;
363 for clause in clauses {
364 if let Some(catch_block) = clause.attribute::<Node>("block") {
365 if has_statements(&catch_block) {
366 self.push_item_kind(
368 CoverageItemKind::Branch {
369 branch_id,
370 path_id,
371 is_first_opcode: true,
372 },
373 &catch_block.src,
374 );
375 self.visit_block(&catch_block)?;
376 path_id += 1;
378 } else if clause.attribute::<Node>("parameters").is_some() {
379 self.push_item_kind(CoverageItemKind::Statement, &clause.src);
383 self.visit_statement(&clause)?;
384 }
385 }
386 }
387
388 Ok(())
389 }
390 NodeType::YulSwitch => {
391 for case in node
393 .attribute::<Vec<Node>>("cases")
394 .ok_or_else(|| eyre::eyre!("yul switch had no case"))?
395 {
396 self.push_item_kind(CoverageItemKind::Statement, &case.src);
397 self.visit_statement(&case)?;
398
399 if let Some(body) = case.body {
400 self.push_item_kind(CoverageItemKind::Statement, &body.src);
401 self.visit_block(&body)?
402 }
403 }
404 Ok(())
405 }
406 NodeType::YulForLoop => {
407 if let Some(condition) = node.attribute("condition") {
408 self.visit_expression(&condition)?;
409 }
410 if let Some(pre) = node.attribute::<Node>("pre") {
411 self.visit_block(&pre)?
412 }
413 if let Some(post) = node.attribute::<Node>("post") {
414 self.visit_block(&post)?
415 }
416
417 if let Some(body) = &node.body {
418 self.push_item_kind(CoverageItemKind::Statement, &body.src);
419 self.visit_block(body)?
420 }
421 Ok(())
422 }
423 NodeType::YulFunctionDefinition => self.visit_modifier_or_yul_fn_definition(node),
424 _ => {
425 warn!("unexpected node type, expected a statement: {:?}", node.node_type);
426 Ok(())
427 }
428 }
429 }
430
431 fn visit_expression(&mut self, node: &Node) -> eyre::Result<()> {
432 match node.node_type {
433 NodeType::Assignment
434 | NodeType::UnaryOperation
435 | NodeType::Conditional
436 | NodeType::YulFunctionCall => {
437 self.push_item_kind(CoverageItemKind::Statement, &node.src);
438 Ok(())
439 }
440 NodeType::FunctionCall => {
441 let kind: Option<String> = node.attribute("kind");
444 if let Some("functionCall") = kind.as_deref() {
445 self.push_item_kind(CoverageItemKind::Statement, &node.src);
446
447 let expr: Option<Node> = node.attribute("expression");
448 if let Some(NodeType::Identifier) = expr.as_ref().map(|expr| &expr.node_type) {
449 let name: Option<String> = expr.and_then(|expr| expr.attribute("name"));
452 if let Some("require") = name.as_deref() {
453 let branch_id = self.branch_id;
454 self.branch_id += 1;
455 self.push_item_kind(
456 CoverageItemKind::Branch {
457 branch_id,
458 path_id: 0,
459 is_first_opcode: false,
460 },
461 &node.src,
462 );
463 self.push_item_kind(
464 CoverageItemKind::Branch {
465 branch_id,
466 path_id: 1,
467 is_first_opcode: false,
468 },
469 &node.src,
470 );
471 }
472 }
473 }
474
475 Ok(())
476 }
477 NodeType::BinaryOperation => {
478 self.push_item_kind(CoverageItemKind::Statement, &node.src);
479
480 if let Some(expr) = node.attribute("leftExpression") {
484 self.visit_expression(&expr)?;
485 }
486
487 if let Some(expr) = node.attribute("rightExpression") {
488 self.visit_expression(&expr)?;
489 }
490
491 Ok(())
492 }
493 NodeType::FunctionCallOptions
495 | NodeType::Identifier
496 | NodeType::IndexAccess
497 | NodeType::IndexRangeAccess
498 | NodeType::Literal
499 | NodeType::YulLiteralValue
500 | NodeType::YulIdentifier => Ok(()),
501 _ => {
502 warn!("unexpected node type, expected an expression: {:?}", node.node_type);
503 Ok(())
504 }
505 }
506 }
507
508 fn visit_block_or_statement(&mut self, node: &Node) -> eyre::Result<()> {
509 match node.node_type {
510 NodeType::Block => self.visit_block(node),
511 NodeType::Break
512 | NodeType::Continue
513 | NodeType::DoWhileStatement
514 | NodeType::EmitStatement
515 | NodeType::ExpressionStatement
516 | NodeType::ForStatement
517 | NodeType::IfStatement
518 | NodeType::InlineAssembly
519 | NodeType::Return
520 | NodeType::RevertStatement
521 | NodeType::TryStatement
522 | NodeType::VariableDeclarationStatement
523 | NodeType::YulVariableDeclaration
524 | NodeType::WhileStatement => self.visit_statement(node),
525 NodeType::PlaceholderStatement => Ok(()),
527 _ => {
528 warn!("unexpected node type, expected block or statement: {:?}", node.node_type);
529 Ok(())
530 }
531 }
532 }
533
534 fn push_item_kind(&mut self, kind: CoverageItemKind, src: &ast::LowFidelitySourceLocation) {
537 let item = CoverageItem { kind, loc: self.source_location_for(src), hits: 0 };
538
539 debug_assert!(!matches!(item.kind, CoverageItemKind::Line));
541 if self.last_line < item.loc.lines.start {
542 self.items.push(CoverageItem {
543 kind: CoverageItemKind::Line,
544 loc: item.loc.clone(),
545 hits: 0,
546 });
547 self.last_line = item.loc.lines.start;
548 }
549
550 self.items.push(item);
551 }
552
553 fn source_location_for(&self, loc: &ast::LowFidelitySourceLocation) -> SourceLocation {
554 let bytes_start = loc.start as u32;
555 let bytes_end = (loc.start + loc.length.unwrap_or(0)) as u32;
556 let bytes = bytes_start..bytes_end;
557
558 let start_line = self.source[..bytes.start as usize].lines().count() as u32;
559 let n_lines = self.source[bytes.start as usize..bytes.end as usize].lines().count() as u32;
560 let lines = start_line..start_line + n_lines;
561 SourceLocation {
562 source_id: self.source_id as usize,
563 contract_name: self.contract_name.clone(),
564 bytes,
565 lines,
566 }
567 }
568}
569
570fn has_statements(node: &Node) -> bool {
572 match node.node_type {
573 NodeType::DoWhileStatement
574 | NodeType::EmitStatement
575 | NodeType::ExpressionStatement
576 | NodeType::ForStatement
577 | NodeType::IfStatement
578 | NodeType::RevertStatement
579 | NodeType::TryStatement
580 | NodeType::VariableDeclarationStatement
581 | NodeType::WhileStatement => true,
582 _ => node.attribute::<Vec<Node>>("statements").is_some_and(|s| !s.is_empty()),
583 }
584}
585
586#[derive(Clone, Debug, Default)]
588pub struct SourceAnalysis {
589 all_items: Vec<CoverageItem>,
591 map: Vec<(u32, u32)>,
593}
594
595impl SourceAnalysis {
596 #[instrument(name = "SourceAnalysis::new", skip_all)]
610 pub fn new(data: &SourceFiles<'_>) -> eyre::Result<Self> {
611 let mut sourced_items = data
612 .sources
613 .par_iter()
614 .flat_map_iter(|(&source_id, SourceFile { source, ast })| {
615 let items = ast.nodes.iter().map(move |node| {
616 if !matches!(node.node_type, NodeType::ContractDefinition) {
617 return Ok(vec![]);
618 }
619
620 let contract_kind: String = node
622 .attribute("contractKind")
623 .ok_or_else(|| eyre::eyre!("Contract has no kind"))?;
624 if contract_kind == "interface" {
625 return Ok(vec![]);
626 }
627
628 let name = node
629 .attribute("name")
630 .ok_or_else(|| eyre::eyre!("Contract has no name"))?;
631 let _guard = debug_span!("visit_contract", %name).entered();
632 let mut visitor = ContractVisitor::new(source_id, &source.content, &name);
633 visitor.visit_contract(node)?;
634 visitor.clear_if_test();
635 visitor.disambiguate_functions();
636 Ok(visitor.items)
637 });
638 items.map(move |items| items.map(|items| (source_id, items)))
639 })
640 .collect::<eyre::Result<Vec<(usize, Vec<CoverageItem>)>>>()?;
641
642 sourced_items.sort_by_key(|(id, items)| (*id, items.first().map(|i| i.loc.bytes.start)));
644 let Some(&(max_idx, _)) = sourced_items.last() else { return Ok(Self::default()) };
645 let len = max_idx + 1;
646 let mut all_items = Vec::new();
647 let mut map = vec![(u32::MAX, 0); len];
648 for (idx, items) in sourced_items {
649 if map[idx].0 == u32::MAX {
651 map[idx].0 = all_items.len() as u32;
652 }
653 map[idx].1 += items.len() as u32;
654 all_items.extend(items);
655 }
656
657 Ok(Self { all_items, map })
658 }
659
660 pub fn all_items(&self) -> &[CoverageItem] {
662 &self.all_items
663 }
664
665 pub fn all_items_mut(&mut self) -> &mut Vec<CoverageItem> {
667 &mut self.all_items
668 }
669
670 pub fn items_for_source_enumerated(
672 &self,
673 source_id: u32,
674 ) -> impl Iterator<Item = (u32, &CoverageItem)> {
675 let (base_id, items) = self.items_for_source(source_id);
676 items.iter().enumerate().map(move |(idx, item)| (base_id + idx as u32, item))
677 }
678
679 pub fn items_for_source(&self, source_id: u32) -> (u32, &[CoverageItem]) {
681 let (mut offset, len) = self.map.get(source_id as usize).copied().unwrap_or_default();
682 if offset == u32::MAX {
683 offset = 0;
684 }
685 (offset, &self.all_items[offset as usize..][..len as usize])
686 }
687
688 #[inline]
690 pub fn get(&self, item_id: u32) -> Option<&CoverageItem> {
691 self.all_items.get(item_id as usize)
692 }
693}
694
695#[derive(Debug, Default)]
697pub struct SourceFiles<'a> {
698 pub sources: HashMap<usize, SourceFile<'a>>,
700}
701
702#[derive(Debug)]
704pub struct SourceFile<'a> {
705 pub source: Source,
707 pub ast: &'a Ast,
709}