From 0566426c392acb2df8ed75b98e11af7e1d3fdb67 Mon Sep 17 00:00:00 2001
From: Arsenii Kulikov <klkvrr@gmail.com>
Date: Mon, 15 Apr 2024 18:40:37 +0400
Subject: [PATCH] feat: coverage for modifiers

---
 crates/evm/coverage/src/analysis.rs | 34 +++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/crates/evm/coverage/src/analysis.rs b/crates/evm/coverage/src/analysis.rs
index a136f2d79cd2..5b86a6f0362a 100644
--- a/crates/evm/coverage/src/analysis.rs
+++ b/crates/evm/coverage/src/analysis.rs
@@ -33,8 +33,14 @@ impl<'a> ContractVisitor<'a> {
     pub fn visit(mut self, contract_ast: Node) -> eyre::Result<Self> {
         // Find all functions and walk their AST
         for node in contract_ast.nodes {
-            if node.node_type == NodeType::FunctionDefinition {
-                self.visit_function_definition(node.clone())?;
+            match &node.node_type {
+                NodeType::FunctionDefinition => {
+                    self.visit_function_definition(node)?;
+                }
+                NodeType::ModifierDefinition => {
+                    self.visit_modifier_definition(node)?;
+                }
+                _ => {}
             }
         }
 
@@ -68,6 +74,23 @@ impl<'a> ContractVisitor<'a> {
         }
     }
 
+    fn visit_modifier_definition(&mut self, mut node: Node) -> eyre::Result<()> {
+        let name: String =
+            node.attribute("name").ok_or_else(|| eyre::eyre!("Modifier has no name"))?;
+
+        match node.body.take() {
+            Some(body) => {
+                self.push_item(CoverageItem {
+                    kind: CoverageItemKind::Function { name },
+                    loc: self.source_location_for(&node.src),
+                    hits: 0,
+                });
+                self.visit_block(*body)
+            }
+            _ => Ok(()),
+        }
+    }
+
     fn visit_block(&mut self, node: Node) -> eyre::Result<()> {
         let statements: Vec<Node> = node.attribute("statements").unwrap_or_default();
 
@@ -93,7 +116,6 @@ impl<'a> ContractVisitor<'a> {
             NodeType::Break |
             NodeType::Continue |
             NodeType::EmitStatement |
-            NodeType::PlaceholderStatement |
             NodeType::RevertStatement |
             NodeType::YulAssignment |
             NodeType::YulBreak |
@@ -107,6 +129,9 @@ impl<'a> ContractVisitor<'a> {
                 Ok(())
             }
 
+            // Skip placeholder statements as they are never referenced in source maps.
+            NodeType::PlaceholderStatement => Ok(()),
+
             // Return with eventual subcall
             NodeType::Return => {
                 self.push_item(CoverageItem {
@@ -341,12 +366,13 @@ impl<'a> ContractVisitor<'a> {
             NodeType::ForStatement |
             NodeType::IfStatement |
             NodeType::InlineAssembly |
-            NodeType::PlaceholderStatement |
             NodeType::Return |
             NodeType::RevertStatement |
             NodeType::TryStatement |
             NodeType::VariableDeclarationStatement |
             NodeType::WhileStatement => self.visit_statement(node),
+            // Skip placeholder statements as they are never referenced in source maps.
+            NodeType::PlaceholderStatement => Ok(()),
             _ => {
                 warn!("unexpected node type, expected block or statement: {:?}", node.node_type);
                 Ok(())