diff --git a/rap/rust-toolchain.toml.bak b/rap/rust-toolchain.toml.bak index 389ade2..e3337a7 100644 --- a/rap/rust-toolchain.toml.bak +++ b/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"] diff --git a/rap/src/analysis/core.rs b/rap/src/analysis/core.rs index e995d72..36ec480 100644 --- a/rap/src/analysis/core.rs +++ b/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; diff --git a/rap/src/analysis/core/call_graph.rs b/rap/src/analysis/core/call_graph.rs new file mode 100644 index 0000000..b433771 --- /dev/null +++ b/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(); + } +} diff --git a/rap/src/analysis/core/call_graph/call_graph_helper.rs b/rap/src/analysis/core/call_graph/call_graph_helper.rs new file mode 100644 index 0000000..5404221 --- /dev/null +++ b/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, // id -> node + // pub function_calls: Vec<(usize, usize)>, // (id, id) + pub function_calls: HashMap>, + pub node_registry: HashMap, // 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 { + 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()); + } + } +} diff --git a/rap/src/analysis/core/call_graph/call_graph_visitor.rs b/rap/src/analysis/core/call_graph/call_graph_visitor.rs new file mode 100644 index 0000000..103a5b2 --- /dev/null +++ b/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); + } + } + } + } + } +} diff --git a/rap/src/analysis/core/control_flow.rs b/rap/src/analysis/core/control_flow.rs index b3ed24b..8b13789 100644 --- a/rap/src/analysis/core/control_flow.rs +++ b/rap/src/analysis/core/control_flow.rs @@ -1 +1 @@ -pub mod callgraph; + diff --git a/rap/src/analysis/core/control_flow/callgraph.rs b/rap/src/analysis/core/control_flow/callgraph.rs deleted file mode 100644 index d489f7b..0000000 --- a/rap/src/analysis/core/control_flow/callgraph.rs +++ /dev/null @@ -1,98 +0,0 @@ -use crate::{rap_debug, rap_info}; -use rustc_data_structures::fx::FxHashMap; -use rustc_hir::{def_id::DefId, intravisit::Visitor, BodyId, HirId, ItemKind}; -use rustc_middle::mir::{Operand, TerminatorKind}; -use rustc_middle::ty::{self, TyCtxt}; -use rustc_span::Span; -use std::collections::HashSet; - -/* - The graph simply records all pairs of callers and callees; - TODO: it can be extended, e.g., - 1) to manage the graph as a linked list of function nodes - 2) to record all attributes of each function -*/ -pub struct CallGraph<'tcx> { - pub tcx: TyCtxt<'tcx>, - pub edges: HashSet<(DefId, DefId)>, -} - -impl<'tcx> CallGraph<'tcx> { - pub fn new(tcx: TyCtxt<'tcx>) -> Self { - Self { - tcx, - edges: HashSet::new(), - } - } - - pub fn start(&mut self) { - rap_info!("Start callgraph analysis"); - let fn_items = FnCollector::collect(self.tcx); - rap_debug!("{:?}", fn_items); - for (_, &ref vec) in &fn_items { - for (body_id, _) in vec { - let body_did = self.tcx.hir().body_owner_def_id(*body_id).to_def_id(); - self.find_callees(body_did); - } - } - rap_info!("Show all edges of the call graph:"); - for (caller, callee) in &self.edges { - rap_info!( - " {} -> {}", - self.tcx.def_path_str(*caller), - self.tcx.def_path_str(*callee) - ); - } - } - - pub fn find_callees(&mut self, def_id: DefId) { - let tcx = self.tcx; - if tcx.is_mir_available(def_id) { - let body = tcx.optimized_mir(def_id); - for bb in body.basic_blocks.iter() { - match &bb.terminator().kind { - TerminatorKind::Call { func, .. } => { - if let Operand::Constant(func_constant) = func { - if let ty::FnDef(ref callee_def_id, _) = - func_constant.const_.ty().kind() - { - self.edges.insert((def_id, *callee_def_id)); - } - } - } - _ => {} - } - } - } - } -} - -/// Maps `HirId` of a type to `BodyId` of related impls. -pub type FnMap = FxHashMap, Vec<(BodyId, Span)>>; - -pub struct FnCollector { - fn_map: FnMap, -} - -impl FnCollector { - pub fn collect<'tcx>(tcx: TyCtxt<'tcx>) -> FnMap { - let mut collector = FnCollector { - fn_map: FnMap::default(), - }; - tcx.hir().visit_all_item_likes_in_crate(&mut collector); - collector.fn_map - } -} - -impl<'tcx> Visitor<'tcx> for FnCollector { - fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) { - match &item.kind { - ItemKind::Fn(_fn_sig, _generics, body_id) => { - let key = Some(body_id.hir_id); - let entry = self.fn_map.entry(key).or_insert(Vec::new()); - entry.push((*body_id, item.span)); - } - _ => (), - } - } -} diff --git a/rap/src/lib.rs b/rap/src/lib.rs index 1e84520..b02113b 100644 --- a/rap/src/lib.rs +++ b/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(); + } }