Include #[test] functions as entry points during dead code analysis (…

…FuelLabs#3151) Refactors the collection of flow graph entry points into a function, and includes functions decorated with the `#[test]` attribute. Short-hand methods have been added to `TyFunctionDeclaration` for checking if the function is `main`, a unit test or if it is an entry point. This is another step towards FuelLabs#2985.
srvt1367 · Oct 26, 2022 · 7bee2b2 · 7bee2b2
1 parent 224e729
commit 7bee2b2
Show file tree

Hide file tree

Showing 2 changed files with 102 additions and 82 deletions.
diff --git a/sway-core/src/control_flow_analysis/dead_code_analysis.rs b/sway-core/src/control_flow_analysis/dead_code_analysis.rs
@@ -108,96 +108,98 @@ impl ControlFlowGraph {
 
             leaves = l_leaves;
         }
+        graph.entry_points = entry_points(tree_type, &graph.graph)?;
+        Ok(())
+    }
+}
 
-        // calculate the entry points based on the tree type
-        graph.entry_points = match tree_type {
-            TreeType::Predicate | TreeType::Script => {
-                // a predicate or script have a main function as the only entry point
-                let mut ret = vec![];
-                for i in graph.graph.node_indices() {
-                    let count_it = match &graph.graph[i] {
-                        ControlFlowGraphNode::OrganizationalDominator(_) => false,
-                        ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
-                            span,
-                            content:
-                                ty::TyAstNodeContent::Declaration(
-                                    ty::TyDeclaration::FunctionDeclaration(decl_id),
-                                ),
-                            ..
-                        }) => {
-                            let decl = de_get_function(decl_id.clone(), span)?;
-                            decl.name.as_str() == "main"
+/// Collect all entry points into the graph based on the tree type.
+fn entry_points(
+    tree_type: &TreeType,
+    graph: &flow_graph::Graph,
+) -> Result<Vec<flow_graph::EntryPoint>, CompileError> {
+    let mut entry_points = vec![];
+    match tree_type {
+        TreeType::Predicate | TreeType::Script => {
+            // Predicates and scripts have main and test functions as entry points.
+            for i in graph.node_indices() {
+                match &graph[i] {
+                    ControlFlowGraphNode::OrganizationalDominator(_) => continue,
+                    ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
+                        span,
+                        content:
+                            ty::TyAstNodeContent::Declaration(ty::TyDeclaration::FunctionDeclaration(
+                                decl_id,
+                            )),
+                        ..
+                    }) => {
+                        let decl = de_get_function(decl_id.clone(), span)?;
+                        if !decl.is_entry() {
+                            continue;
                         }
-                        _ => false,
-                    };
-                    if count_it {
-                        ret.push(i);
                     }
-                }
-                ret
+                    _ => continue,
+                };
+                entry_points.push(i);
             }
-            TreeType::Contract | TreeType::Library { .. } => {
-                let mut ret = vec![];
-                for i in graph.graph.node_indices() {
-                    let count_it = match &graph.graph[i] {
-                        ControlFlowGraphNode::OrganizationalDominator(_) => false,
-                        ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
-                            content:
-                                ty::TyAstNodeContent::Declaration(
-                                    ty::TyDeclaration::FunctionDeclaration(decl_id),
-                                ),
-                            ..
-                        }) => {
-                            let function_decl = de_get_function(decl_id.clone(), &decl_id.span())?;
-                            function_decl.visibility == Visibility::Public
-                        }
-                        ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
-                            content:
-                                ty::TyAstNodeContent::Declaration(ty::TyDeclaration::TraitDeclaration(
-                                    decl_id,
-                                )),
-                            ..
-                        }) => de_get_trait(decl_id.clone(), &decl_id.span())?
-                            .visibility
-                            .is_public(),
-                        ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
-                            content:
-                                ty::TyAstNodeContent::Declaration(ty::TyDeclaration::StructDeclaration(
-                                    decl_id,
-                                )),
-                            ..
-                        }) => {
-                            let struct_decl = de_get_struct(decl_id.clone(), &decl_id.span())?;
-                            struct_decl.visibility == Visibility::Public
-                        }
-                        ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
-                            content:
-                                ty::TyAstNodeContent::Declaration(ty::TyDeclaration::ImplTrait {
-                                    ..
-                                }),
-                            ..
-                        }) => true,
-                        ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
-                            content:
-                                ty::TyAstNodeContent::Declaration(
-                                    ty::TyDeclaration::ConstantDeclaration(decl_id),
-                                ),
-                            ..
-                        }) => {
-                            let decl = de_get_constant(decl_id.clone(), &decl_id.span())?;
-                            decl.visibility.is_public()
-                        }
-                        _ => false,
-                    };
-                    if count_it {
-                        ret.push(i);
+        }
+        TreeType::Contract | TreeType::Library { .. } => {
+            for i in graph.node_indices() {
+                let is_entry = match &graph[i] {
+                    ControlFlowGraphNode::OrganizationalDominator(_) => continue,
+                    ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
+                        content:
+                            ty::TyAstNodeContent::Declaration(ty::TyDeclaration::FunctionDeclaration(
+                                decl_id,
+                            )),
+                        ..
+                    }) => {
+                        let decl = de_get_function(decl_id.clone(), &decl_id.span())?;
+                        decl.visibility == Visibility::Public || decl.is_test()
+                    }
+                    ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
+                        content:
+                            ty::TyAstNodeContent::Declaration(ty::TyDeclaration::TraitDeclaration(
+                                decl_id,
+                            )),
+                        ..
+                    }) => de_get_trait(decl_id.clone(), &decl_id.span())?
+                        .visibility
+                        .is_public(),
+                    ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
+                        content:
+                            ty::TyAstNodeContent::Declaration(ty::TyDeclaration::StructDeclaration(
+                                decl_id,
+                            )),
+                        ..
+                    }) => {
+                        let struct_decl = de_get_struct(decl_id.clone(), &decl_id.span())?;
+                        struct_decl.visibility == Visibility::Public
+                    }
+                    ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
+                        content:
+                            ty::TyAstNodeContent::Declaration(ty::TyDeclaration::ImplTrait { .. }),
+                        ..
+                    }) => true,
+                    ControlFlowGraphNode::ProgramNode(ty::TyAstNode {
+                        content:
+                            ty::TyAstNodeContent::Declaration(ty::TyDeclaration::ConstantDeclaration(
+                                decl_id,
+                            )),
+                        ..
+                    }) => {
+                        let decl = de_get_constant(decl_id.clone(), &decl_id.span())?;
+                        decl.visibility.is_public()
                     }
+                    _ => continue,
+                };
+                if is_entry {
+                    entry_points.push(i);
                 }
-                ret
             }
-        };
-        Ok(())
+        }
     }
+    Ok(entry_points)
 }
 
 fn connect_node(

diff --git a/sway-core/src/language/ty/declaration/function.rs b/sway-core/src/language/ty/declaration/function.rs
@@ -253,6 +253,24 @@ impl TyFunctionDeclaration {
             },
         }
     }
+
+    /// Whether or not this function is the default entry point.
+    pub fn is_main_entry(&self) -> bool {
+        // NOTE: We may want to make this check more sophisticated or customisable in the future,
+        // but for now this assumption is baked in throughout the compiler.
+        self.name.as_str() == sway_types::constants::DEFAULT_ENTRY_POINT_FN_NAME
+    }
+
+    /// Whether or not this function is a unit test, i.e. decorated with `#[test]`.
+    pub fn is_test(&self) -> bool {
+        self.attributes
+            .contains_key(&transform::AttributeKind::Test)
+    }
+
+    /// Whether or not this function describes a program entry point.
+    pub fn is_entry(&self) -> bool {
+        self.is_main_entry() || self.is_test()
+    }
 }
 
 #[derive(Debug, Clone, Eq)]