Merge pull request Artisan-Lab#64 from SeanXunX/callgraph_feature

Add Callgraph feature

rap/rust-toolchain.toml.bak

@@ -2,4 +2,4 @@
 # The default version of the rustc compiler
 #channel = "nightly-2023-10-05-x86_64-unknown-linux-gnu"
 channel = "nightly-2024-06-30-x86_64-unknown-linux-gnu"
-components = ["rustc-dev", "rustc-src", "llvm-tools-preview"]
+components = ["rustc-dev", "rust-src", "llvm-tools-preview"]

rap/src/analysis/core.rs

@@ -1,4 +1,5 @@
 pub mod alias;
+pub mod call_graph;
 pub mod control_flow;
 pub mod dataflow;
 pub mod heap_item;

rap/src/analysis/core/call_graph.rs

@@ -0,0 +1,30 @@
+pub mod call_graph_helper;
+pub mod call_graph_visitor;
+
+use call_graph_helper::CallGraphInfo;
+use call_graph_visitor::CallGraphVisitor;
+use rustc_middle::ty::TyCtxt;
+
+pub struct CallGraph<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+    pub graph: CallGraphInfo,
+}
+
+impl<'tcx> CallGraph<'tcx> {
+    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
+        Self {
+            tcx: tcx,
+            graph: CallGraphInfo::new(),
+        }
+    }
+
+    pub fn start(&mut self) {
+        for &def_id in self.tcx.mir_keys(()).iter() {
+            let body = &self.tcx.optimized_mir(def_id);
+            let mut call_graph_visitor =
+                CallGraphVisitor::new(self.tcx, def_id.into(), body, &mut self.graph);
+            call_graph_visitor.visit();
+        }
+        self.graph.print_call_graph();
+    }
+}

rap/src/analysis/core/call_graph/call_graph_helper.rs

@@ -0,0 +1,96 @@
+use rustc_hir::def_id::DefId;
+use std::collections::HashSet;
+use std::{collections::HashMap, hash::Hash};
+
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+pub struct Node {
+    def_id: DefId,
+    def_path: String,
+}
+
+impl Node {
+    pub fn new(def_id: DefId, def_path: &String) -> Self {
+        Self {
+            def_id: def_id,
+            def_path: def_path.clone(),
+        }
+    }
+
+    pub fn get_def_id(&self) -> DefId {
+        self.def_id
+    }
+
+    pub fn get_def_path(&self) -> String {
+        self.def_path.clone()
+    }
+}
+
+pub struct CallGraphInfo {
+    pub functions: HashMap<usize, Node>, // id -> node
+    // pub function_calls: Vec<(usize, usize)>,   // (id, id)
+    pub function_calls: HashMap<usize, HashSet<usize>>,
+    pub node_registry: HashMap<String, usize>, // path -> id
+}
+
+impl CallGraphInfo {
+    pub fn new() -> Self {
+        Self {
+            functions: HashMap::new(),
+            function_calls: HashMap::new(),
+            node_registry: HashMap::new(),
+        }
+    }
+
+    pub fn get_node_num(&self) -> usize {
+        self.functions.len()
+    }
+
+    pub fn add_node(&mut self, def_id: DefId, def_path: &String) {
+        if let None = self.get_noed_by_path(def_path) {
+            let id = self.node_registry.len();
+            let node = Node::new(def_id, def_path);
+            self.node_registry.insert(def_path.clone(), id);
+            self.functions.insert(id, node);
+        }
+    }
+
+    pub fn add_funciton_call_edge(&mut self, caller_id: usize, callee_id: usize) {
+        if !self.function_calls.contains_key(&caller_id) {
+            self.function_calls.insert(caller_id, HashSet::new());
+        }
+        if let Some(callees) = self.function_calls.get_mut(&caller_id) {
+            callees.insert(callee_id);
+        }
+    }
+
+    pub fn get_noed_by_path(&self, def_path: &String) -> Option<usize> {
+        if let Some(&id) = self.node_registry.get(def_path) {
+            Some(id)
+        } else {
+            None
+        }
+    }
+
+    pub fn print_call_graph(&self) {
+        println!("CallGraph Analysis:");
+        // println!("There are {} functions calls!", self.function_calls.len());
+        for (caller_id, callees) in self.function_calls.clone() {
+            if let Some(caller_node) = self.functions.get(&caller_id) {
+                for callee_id in callees {
+                    if let Some(callee_node) = self.functions.get(&callee_id) {
+                        let caller_def_path = caller_node.get_def_path();
+                        let callee_def_path = callee_node.get_def_path();
+                        println!(
+                            "{}:{} -> {}:{}",
+                            caller_id, caller_def_path, callee_id, callee_def_path
+                        );
+                    }
+                }
+            }
+        }
+        println!("There are {} functions", self.functions.len());
+        for (id, node) in self.functions.clone() {
+            println!("{}:{}", id, node.get_def_path());
+        }
+    }
+}

rap/src/analysis/core/call_graph/call_graph_visitor.rs

@@ -0,0 +1,103 @@
+use super::call_graph_helper::CallGraphInfo;
+use rustc_hir::def_id::DefId;
+use rustc_middle::mir;
+use rustc_middle::ty::{FnDef, Instance, InstanceKind, TyCtxt};
+
+pub struct CallGraphVisitor<'b, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+    body: &'tcx mir::Body<'tcx>,
+    call_graph_info: &'b mut CallGraphInfo,
+}
+
+impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> {
+    pub fn new(
+        tcx: TyCtxt<'tcx>,
+        def_id: DefId,
+        body: &'tcx mir::Body<'tcx>,
+        call_graph_info: &'b mut CallGraphInfo,
+    ) -> Self {
+        Self {
+            tcx: tcx,
+            def_id: def_id,
+            body: body,
+            call_graph_info: call_graph_info,
+        }
+    }
+
+    pub fn add_in_call_graph(
+        &mut self,
+        caller_def_path: &String,
+        callee_def_id: DefId,
+        callee_def_path: &String,
+    ) {
+        if let Some(caller_id) = self.call_graph_info.get_noed_by_path(caller_def_path) {
+            if let Some(callee_id) = self.call_graph_info.get_noed_by_path(callee_def_path) {
+                self.call_graph_info
+                    .add_funciton_call_edge(caller_id, callee_id);
+            } else {
+                self.call_graph_info
+                    .add_node(callee_def_id, callee_def_path);
+                if let Some(callee_id) = self.call_graph_info.get_noed_by_path(callee_def_path) {
+                    self.call_graph_info
+                        .add_funciton_call_edge(caller_id, callee_id);
+                }
+            }
+        }
+    }
+
+    pub fn visit(&mut self) {
+        let caller_path_str = self.tcx.def_path_str(self.def_id);
+        self.call_graph_info.add_node(self.def_id, &caller_path_str);
+        for (_, data) in self.body.basic_blocks.iter().enumerate() {
+            let terminator = data.terminator();
+            self.visit_terminator(&terminator);
+        }
+    }
+
+    fn add_to_call_graph(&mut self, callee_def_id: DefId) {
+        let caller_def_path = self.tcx.def_path_str(self.def_id);
+        let callee_def_path = self.tcx.def_path_str(callee_def_id);
+        // let callee_location = self.tcx.def_span(callee_def_id);
+        if callee_def_id == self.def_id {
+            // Recursion
+            println!("Warning! Find a recursion function which may cause stackoverflow!")
+        }
+        self.add_in_call_graph(&caller_def_path, callee_def_id, &callee_def_path);
+        println!("")
+    }
+
+    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>) {
+        if let mir::TerminatorKind::Call { func, .. } = &terminator.kind {
+            if let mir::Operand::Constant(constant) = func {
+                if let FnDef(callee_def_id, callee_substs) = constant.const_.ty().kind() {
+                    let param_env = self.tcx.param_env(self.def_id);
+                    if let Ok(Some(instance)) =
+                        Instance::resolve(self.tcx, param_env, *callee_def_id, callee_substs)
+                    {
+                        // Try to analysis the specific type of callee.
+                        let instance_def_id = match instance.def {
+                            InstanceKind::Item(def_id) => Some(def_id),
+                            InstanceKind::Intrinsic(def_id)
+                            | InstanceKind::CloneShim(def_id, _) => {
+                                if !self.tcx.is_closure_like(def_id) {
+                                    // Not a closure
+                                    Some(def_id)
+                                } else {
+                                    None
+                                }
+                            }
+                            _ => None,
+                        };
+                        if let Some(instance_def_id) = instance_def_id {
+                            self.add_to_call_graph(instance_def_id);
+                        }
+                    } else {
+                        // Although failing to get specific type, callee is still useful.
+                        self.add_to_call_graph(*callee_def_id);
+                    }
+                }
+            }
+        }
+    }
+}

rap/src/analysis/core/control_flow.rs

@@ -1 +1 @@
-pub mod callgraph;
+

rap/src/lib.rs

@@ -18,7 +18,7 @@ extern crate rustc_span;
 extern crate rustc_target;
 
 use analysis::core::alias::mop::MopAlias;
-use analysis::core::control_flow::callgraph::CallGraph;
+use analysis::core::call_graph::CallGraph;
 use analysis::core::dataflow::DataFlow;
 use analysis::rcanary::rCanary;
 use analysis::safedrop::SafeDrop;
@@ -227,10 +227,6 @@ pub fn start_analyzer(tcx: TyCtxt, callback: RapCallback) {
         SenryxCheck::new(tcx, 2).start();
     }
 
-    if callback.is_callgraph_enabled() {
-        CallGraph::new(tcx).start();
-    }
-
     if callback.is_show_mir_enabled() {
         ShowMir::new(tcx).start();
     }
@@ -240,4 +236,8 @@ pub fn start_analyzer(tcx: TyCtxt, callback: RapCallback) {
         2 => DataFlow::new(tcx, true).start(),
         _ => {}
     }
+
+    if callback.is_callgraph_enabled() {
+        CallGraph::new(tcx).start();
+    }
 }