duskphantom_middle/analysis/
call_graph.rs1use crate::ir::instruction::downcast_ref;
18use crate::{
19 ir::{
20 instruction::{misc_inst::Call, InstType},
21 FunPtr, InstPtr,
22 },
23 Program,
24};
25use std::collections::{HashMap, HashSet};
26
27#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
28pub struct CallEdge {
29 pub inst: InstPtr,
30 pub caller: FunPtr,
31 pub callee: FunPtr,
32}
33
34#[allow(unused)]
35pub struct CallGraph {
36 main: Option<FunPtr>,
37 calls: HashMap<FunPtr, HashSet<CallEdge>>,
38 called_by: HashMap<FunPtr, HashSet<CallEdge>>,
39}
40
41impl CallGraph {
42 pub fn new(program: &Program) -> Self {
43 let mut calls = HashMap::new();
44 let mut called_by = HashMap::new();
45 let mut main = None;
46 for func in program.module.functions.clone() {
47 if func.name == "main" {
48 main = Some(func);
49 }
50
51 if func.is_lib() {
53 continue;
54 }
55
56 for bb in func.dfs_iter() {
58 for inst in bb.iter() {
59 if inst.get_type() == InstType::Call {
60 let call = downcast_ref::<Call>(inst.as_ref().as_ref());
61
62 if call.func.is_lib() {
64 continue;
65 }
66
67 let call_edge = CallEdge {
69 inst,
70 caller: func,
71 callee: call.func,
72 };
73 calls
74 .entry(func)
75 .or_insert(HashSet::new())
76 .insert(call_edge);
77 called_by
78 .entry(call.func)
79 .or_insert(HashSet::new())
80 .insert(call_edge);
81 }
82 }
83 }
84 }
85 CallGraph {
86 main,
87 calls,
88 called_by,
89 }
90 }
91
92 pub fn get_calls(&self, func: FunPtr) -> HashSet<CallEdge> {
93 self.calls.get(&func).cloned().unwrap_or_default()
94 }
95
96 pub fn get_called_by(&self, func: FunPtr) -> HashSet<CallEdge> {
97 self.called_by.get(&func).cloned().unwrap_or_default()
98 }
99
100 pub fn remove(&mut self, func: FunPtr) {
101 if let Some(calls) = self.calls.remove(&func) {
102 for call in calls {
103 self.called_by.get_mut(&call.caller).unwrap().remove(&call);
104 }
105 }
106 if let Some(called_by) = self.called_by.remove(&func) {
107 for call in called_by {
108 self.calls.get_mut(&call.caller).unwrap().remove(&call);
109 }
110 }
111 }
112}