duskphantom_middle/analysis/
call_graph.rs

1// Copyright 2024 Duskphantom Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15// SPDX-License-Identifier: Apache-2.0
16
17use 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            // Caller should not be library function
52            if func.is_lib() {
53                continue;
54            }
55
56            // Iterate all instructions
57            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                        // Callee should not be library function
63                        if call.func.is_lib() {
64                            continue;
65                        }
66
67                        // Construct and add call edge
68                        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}