From 3f7c5915244e3e5ba34592ab5a662a5e96a245aa Mon Sep 17 00:00:00 2001 From: sean Date: Thu, 24 Oct 2024 11:18:17 +0800 Subject: [PATCH] Substitute vec with HashSet --- rap/rust-toolchain.toml | 2 +- rap/src/analysis/core.rs | 2 +- rap/src/analysis/core/call_graph.rs | 6 +-- .../core/call_graph/call_graph_helper.rs | 48 +++++++++++-------- .../core/call_graph/call_graph_visitor.rs | 47 +++++++++++------- rap/src/analysis/core/control_flow.rs | 1 + rap/src/lib.rs | 2 +- 7 files changed, 62 insertions(+), 46 deletions(-) diff --git a/rap/rust-toolchain.toml b/rap/rust-toolchain.toml index 46d8578..e3337a7 100644 --- a/rap/rust-toolchain.toml +++ b/rap/rust-toolchain.toml @@ -1,5 +1,5 @@ [toolchain] # The default version of the rustc compiler #channel = "nightly-2023-10-05-x86_64-unknown-linux-gnu" -channel = "nightly-2024-06-30" +channel = "nightly-2024-06-30-x86_64-unknown-linux-gnu" components = ["rustc-dev", "rust-src", "llvm-tools-preview"] diff --git a/rap/src/analysis/core.rs b/rap/src/analysis/core.rs index 041cc07..36ec480 100644 --- a/rap/src/analysis/core.rs +++ b/rap/src/analysis/core.rs @@ -1,5 +1,5 @@ pub mod alias; +pub mod call_graph; pub mod control_flow; pub mod dataflow; pub mod heap_item; -pub mod call_graph; diff --git a/rap/src/analysis/core/call_graph.rs b/rap/src/analysis/core/call_graph.rs index ccd19fd..b433771 100644 --- a/rap/src/analysis/core/call_graph.rs +++ b/rap/src/analysis/core/call_graph.rs @@ -5,8 +5,6 @@ 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, @@ -23,10 +21,10 @@ impl<'tcx> CallGraph<'tcx> { 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); + 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 index 5520856..5404221 100644 --- a/rap/src/analysis/core/call_graph/call_graph_helper.rs +++ b/rap/src/analysis/core/call_graph/call_graph_helper.rs @@ -1,6 +1,6 @@ -use std::{collections::HashMap, hash::Hash}; 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 { @@ -26,16 +26,17 @@ impl Node { } pub struct CallGraphInfo { - pub functions: HashMap, // id -> node - pub function_calls: Vec<(usize, usize)>, // (id, id) - pub node_registry: HashMap, // path -> id + 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: Vec::new(), + function_calls: HashMap::new(), node_registry: HashMap::new(), } } @@ -44,7 +45,7 @@ impl CallGraphInfo { self.functions.len() } - pub fn add_node(& mut self, def_id: DefId, def_path: &String) { + 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); @@ -53,29 +54,37 @@ impl CallGraphInfo { } } - pub fn add_funciton_call_edge(& mut self, caller_id: usize, callee_id: usize) { - self.function_calls.push((caller_id, callee_id)); + 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 { + } else { None } } pub fn print_call_graph(&self) { println!("CallGraph Analysis:"); - println!("There are {} functions calls!", self.function_calls.len()); - for call in self.function_calls.clone() { - let caller_id = call.0; - let callee_id = call.1; + // 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) { - 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!("{}:{} -> {}:{}", call.0, caller_def_path, call.1, callee_def_path); + 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 + ); + } } } } @@ -85,6 +94,3 @@ impl CallGraphInfo { } } } - - - diff --git a/rap/src/analysis/core/call_graph/call_graph_visitor.rs b/rap/src/analysis/core/call_graph/call_graph_visitor.rs index 9c2956b..103a5b2 100644 --- a/rap/src/analysis/core/call_graph/call_graph_visitor.rs +++ b/rap/src/analysis/core/call_graph/call_graph_visitor.rs @@ -1,8 +1,7 @@ use super::call_graph_helper::CallGraphInfo; -use rustc_middle::mir; -use rustc_middle::ty::{TyCtxt, Instance, FnDef, InstanceKind}; 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>, @@ -12,7 +11,12 @@ pub struct CallGraphVisitor<'b, 'tcx> { } 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 { + 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, @@ -21,14 +25,22 @@ impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> { } } - pub fn add_in_call_graph(&mut self, caller_def_path: &String, callee_def_id: DefId, callee_def_path: &String) { + 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); + 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); + 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); + self.call_graph_info + .add_funciton_call_edge(caller_id, callee_id); } } } @@ -40,8 +52,8 @@ impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> { 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); @@ -56,18 +68,18 @@ impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> { } fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>) { - if let mir::TerminatorKind::Call { - func, - .. - } = &terminator.kind { + 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) { + 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, _) => { + InstanceKind::Intrinsic(def_id) + | InstanceKind::CloneShim(def_id, _) => { if !self.tcx.is_closure_like(def_id) { // Not a closure Some(def_id) @@ -87,6 +99,5 @@ impl<'b, 'tcx> CallGraphVisitor<'b, 'tcx> { } } } - } - + } } diff --git a/rap/src/analysis/core/control_flow.rs b/rap/src/analysis/core/control_flow.rs index e69de29..8b13789 100644 --- a/rap/src/analysis/core/control_flow.rs +++ b/rap/src/analysis/core/control_flow.rs @@ -0,0 +1 @@ + diff --git a/rap/src/lib.rs b/rap/src/lib.rs index 91680e6..b02113b 100644 --- a/rap/src/lib.rs +++ b/rap/src/lib.rs @@ -18,8 +18,8 @@ extern crate rustc_span; extern crate rustc_target; use analysis::core::alias::mop::MopAlias; -use analysis::core::dataflow::DataFlow; use analysis::core::call_graph::CallGraph; +use analysis::core::dataflow::DataFlow; use analysis::rcanary::rCanary; use analysis::safedrop::SafeDrop; use analysis::senryx::SenryxCheck;