Global pre-processing for leveraging callgraphs

aderyn_core/Cargo.toml

@@ -23,6 +23,7 @@ serde_repr = "0.1.12"
 strum = { version = "0.26", features = ["derive"] }
 toml = "0.8.2"
 cyfrin-foundry-compilers = { version = "0.3.20-aderyn", features = ["svm-solc"]  }
+derive_more = "0.99.18"
 
 [dev-dependencies]
 serial_test = "3.0.0"

aderyn_core/src/audit/public_functions_no_sender.rs

@@ -60,7 +60,7 @@ impl AuditorDetector for PublicFunctionsNoSenderChecksDetector {
                         });
                 // Check if the function has a `msg.sender` BinaryOperation check
                 let has_msg_sender_binary_operation =
-                    has_msg_sender_binary_operation(function_definition);
+                    has_msg_sender_binary_operation(&((*function_definition).into()));
                 // TODO Check if the function has a hasRole identifier with msg.sender as an arg
                 does_not_have_an_owner_modifier && !has_msg_sender_binary_operation
             });

aderyn_core/src/context/browser/extractor.rs

@@ -88,3 +88,40 @@ impl ASTConstVisitor for ExtractImmediateChildrenIDs {
         Ok(())
     }
 }
+
+// Extract Reference Declaration IDs
+#[derive(Default)]
+pub struct ExtractReferencedDeclarations {
+    pub extracted: Vec<NodeID>,
+}
+
+impl ExtractReferencedDeclarations {
+    pub fn from<T: Node + ?Sized>(node: &T) -> Self {
+        let mut extractor: ExtractReferencedDeclarations = Self::default();
+        node.accept(&mut extractor).unwrap_or_default();
+        extractor
+    }
+}
+
+impl ASTConstVisitor for ExtractReferencedDeclarations {
+    fn visit_member_access(&mut self, node: &MemberAccess) -> Result<bool> {
+        if let Some(referenced_id) = node.referenced_declaration {
+            self.extracted.push(referenced_id);
+        }
+        Ok(true)
+    }
+    fn visit_identifier(&mut self, node: &Identifier) -> Result<bool> {
+        if let Some(referenced_id) = node.referenced_declaration {
+            self.extracted.push(referenced_id);
+        }
+        Ok(true)
+    }
+    fn visit_identifier_path(&mut self, node: &IdentifierPath) -> Result<bool> {
+        self.extracted.push(node.referenced_declaration as i64);
+        Ok(true)
+    }
+    fn visit_user_defined_type_name(&mut self, node: &UserDefinedTypeName) -> Result<bool> {
+        self.extracted.push(node.referenced_declaration);
+        Ok(true)
+    }
+}

aderyn_core/src/context/graph/mod.rs

@@ -0,0 +1,32 @@
+pub mod traits;
+mod workspace_callgraph;
+
+pub use workspace_callgraph::*;
+
+use derive_more::From;
+
+pub type Result<T> = core::result::Result<T, Error>;
+
+#[derive(Debug, From)]
+pub enum Error {
+    #[from]
+    Custom(String),
+
+    // region: -- standard::* errors
+    WorkspaceCallGraphDFSError,
+    // endregion
+}
+
+impl core::fmt::Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{self:?}")
+    }
+}
+
+impl From<&str> for Error {
+    fn from(value: &str) -> Self {
+        Error::Custom(value.to_string())
+    }
+}
+
+impl std::error::Error for Error {}

aderyn_core/src/context/graph/traits.rs

@@ -0,0 +1,4 @@
+/// Trait to support reversing of callgraph. (Because, direct impl is not allowed on Foreign Types)
+pub trait Transpose {
+    fn reverse(&self) -> Self;
+}

aderyn_core/src/context/graph/workspace_callgraph.rs

@@ -0,0 +1,134 @@
+use std::collections::{hash_map, HashMap, HashSet};
+
+use crate::{
+    ast::{Expression, IdentifierOrIdentifierPath, NodeID, NodeType},
+    context::{
+        browser::{ExtractFunctionCalls, ExtractModifierInvocations},
+        workspace_context::WorkspaceContext,
+    },
+};
+
+use super::traits::Transpose;
+
+#[derive(Debug)]
+pub struct WorkspaceCallGraph {
+    pub graph: CallGraph,
+}
+
+/**
+* Every NodeID in CallGraph should corresponds to [`crate::ast::FunctionDefinition`] or [`crate::ast::ModifierDefinition`]
+*/
+pub type CallGraph = HashMap<NodeID, Vec<NodeID>>;
+
+impl WorkspaceCallGraph {
+    /// Formula to create [`WorkspaceCallGraph`] for global preprocessing .
+    pub fn from_context(context: &WorkspaceContext) -> super::Result<WorkspaceCallGraph> {
+        let mut graph: CallGraph = HashMap::new();
+        let mut visited: HashSet<NodeID> = HashSet::new();
+
+        let funcs = context
+            .function_definitions()
+            .into_iter()
+            .filter(|func| func.implemented)
+            .collect::<Vec<_>>();
+
+        let modifier_definitions = context.modifier_definitions();
+
+        for func in funcs {
+            dfs_to_create_graph(func.id, &mut graph, &mut visited, context)
+                .map_err(|_| super::Error::WorkspaceCallGraphDFSError)?;
+        }
+
+        for modifier in modifier_definitions {
+            dfs_to_create_graph(modifier.id, &mut graph, &mut visited, context)
+                .map_err(|_| super::Error::WorkspaceCallGraphDFSError)?;
+        }
+
+        Ok(WorkspaceCallGraph { graph })
+    }
+}
+
+/// Make connections from each of the nodes of [`crate::ast::FunctionDefinition`] and [`crate::ast::ModifierDefinition`]
+/// with their connected counterparts.
+fn dfs_to_create_graph(
+    id: NodeID,
+    graph: &mut CallGraph,
+    visited: &mut HashSet<NodeID>,
+    context: &WorkspaceContext,
+) -> super::Result<()> {
+    if visited.contains(&id) {
+        return Ok(());
+    }
+
+    visited.insert(id);
+
+    // Only deal with `id`s that are in scope right now
+    if let Some(from_node) = context.nodes.get(&id) {
+        // referenced_declarations from previous calls in the recursion stack need to be vetted
+        if from_node.node_type() != NodeType::FunctionDefinition
+            && from_node.node_type() != NodeType::ModifierDefinition
+        {
+            return Ok(());
+        }
+
+        // connections to FunctionDefinition
+        let function_calls = ExtractFunctionCalls::from(from_node).extracted;
+        for function_call in function_calls {
+            if let Expression::Identifier(identifier) = function_call.expression.as_ref() {
+                if let Some(referenced_function_id) = identifier.referenced_declaration {
+                    create_connection_if_not_exsits(id, referenced_function_id, graph);
+                    dfs_to_create_graph(referenced_function_id, graph, visited, context)?;
+                }
+            }
+        }
+
+        // connections to ModifierDefinition
+        let modifier_invocations = ExtractModifierInvocations::from(from_node).extracted;
+        for modifier_invocation in &modifier_invocations {
+            match &modifier_invocation.modifier_name {
+                IdentifierOrIdentifierPath::Identifier(identifier) => {
+                    if let Some(reference_modifier_id) = identifier.referenced_declaration {
+                        create_connection_if_not_exsits(id, reference_modifier_id, graph);
+                        dfs_to_create_graph(reference_modifier_id, graph, visited, context)?;
+                    }
+                }
+                IdentifierOrIdentifierPath::IdentifierPath(identifier_path) => {
+                    let referenced_modifier_id = identifier_path.referenced_declaration;
+                    create_connection_if_not_exsits(id, referenced_modifier_id as i64, graph);
+                    dfs_to_create_graph(referenced_modifier_id as i64, graph, visited, context)?;
+                }
+            }
+        }
+    }
+
+    // Change the default return to error later in "strict mode" maybe, because if we
+    // can't find the node that means, the file was not in scope and hence it is not
+    // available in the context although references to it exist.
+    Ok(())
+}
+
+fn create_connection_if_not_exsits(from_id: NodeID, to_id: NodeID, graph: &mut CallGraph) {
+    match graph.entry(from_id) {
+        hash_map::Entry::Occupied(mut o) => {
+            // Performance Tip: Maybe later use binary search (it requires keeping ascending order while inserting tho)
+            if !o.get().contains(&to_id) {
+                o.get_mut().push(to_id);
+            }
+        }
+        hash_map::Entry::Vacant(v) => {
+            v.insert(vec![to_id]);
+        }
+    }
+}
+
+impl Transpose for CallGraph {
+    fn reverse(&self) -> Self {
+        let mut reversed_callgraph = CallGraph::default();
+        for (from_id, tos) in self {
+            for to_id in tos {
+                create_connection_if_not_exsits(*to_id, *from_id, &mut reversed_callgraph);
+            }
+        }
+        reversed_callgraph
+    }
+}