llvm_analysis/
call_graph.rs

1//! Module-level call graph with edge kinds and SCC traversal.
2
3use llvm_ir::{Context, FunctionId, InstrKind, Module, ValueRef};
4use std::collections::{BTreeSet, HashSet};
5
6/// Public API for `CallEdgeKind`.
7#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
8pub enum CallEdgeKind {
9    /// `Direct` variant.
10    Direct,
11    /// `Indirect` variant.
12    Indirect,
13    /// `External` variant.
14    External,
15}
16
17/// Public API for `CallEdge`.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
19pub struct CallEdge {
20    /// Public API for `to`.
21    pub to: Option<FunctionId>,
22    /// Public API for `kind`.
23    pub kind: CallEdgeKind,
24}
25
26/// Call graph over all functions in a module.
27pub struct CallGraph {
28    // `direct_callees` field.
29    direct_callees: Vec<Vec<FunctionId>>,
30    // `direct_callers` field.
31    direct_callers: Vec<Vec<FunctionId>>,
32    // `edges` field.
33    edges: Vec<Vec<CallEdge>>,
34}
35
36impl CallGraph {
37    /// Public API for `build`.
38    pub fn build(_ctx: &Context, module: &Module) -> Self {
39        let n = module.functions.len();
40        let mut direct_callees_sets: Vec<BTreeSet<FunctionId>> = vec![BTreeSet::new(); n];
41        let mut direct_callers_sets: Vec<BTreeSet<FunctionId>> = vec![BTreeSet::new(); n];
42        let mut edges: Vec<Vec<CallEdge>> = vec![Vec::new(); n];
43
44        for (src_idx, func) in module.functions.iter().enumerate() {
45            for instr in &func.instructions {
46                let InstrKind::Call { callee, .. } = &instr.kind else {
47                    continue;
48                };
49                let edge = match callee {
50                    ValueRef::Global(gid) => {
51                        let fid = FunctionId(gid.0);
52                        if (fid.0 as usize) < n {
53                            let kind = if module.functions[fid.0 as usize].is_declaration {
54                                CallEdgeKind::External
55                            } else {
56                                CallEdgeKind::Direct
57                            };
58                            direct_callees_sets[src_idx].insert(fid);
59                            direct_callers_sets[fid.0 as usize].insert(FunctionId(src_idx as u32));
60                            CallEdge {
61                                to: Some(fid),
62                                kind,
63                            }
64                        } else {
65                            CallEdge {
66                                to: None,
67                                kind: CallEdgeKind::External,
68                            }
69                        }
70                    }
71                    _ => CallEdge {
72                        to: None,
73                        kind: CallEdgeKind::Indirect,
74                    },
75                };
76                edges[src_idx].push(edge);
77            }
78        }
79
80        let direct_callees = direct_callees_sets
81            .into_iter()
82            .map(|s| s.into_iter().collect())
83            .collect();
84        let direct_callers = direct_callers_sets
85            .into_iter()
86            .map(|s| s.into_iter().collect())
87            .collect();
88        Self {
89            direct_callees,
90            direct_callers,
91            edges,
92        }
93    }
94
95    /// Public API for `callees`.
96    pub fn callees(&self, f: FunctionId) -> &[FunctionId] {
97        &self.direct_callees[f.0 as usize]
98    }
99
100    /// Public API for `callers`.
101    pub fn callers(&self, f: FunctionId) -> &[FunctionId] {
102        &self.direct_callers[f.0 as usize]
103    }
104
105    /// Public API for `edges`.
106    pub fn edges(&self, f: FunctionId) -> &[CallEdge] {
107        &self.edges[f.0 as usize]
108    }
109
110    /// Return SCCs in bottom-up order over the direct-call graph.
111    pub fn sccs(&self) -> Vec<Vec<FunctionId>> {
112        let n = self.direct_callees.len();
113        let mut order = Vec::with_capacity(n);
114        let mut seen = vec![false; n];
115        for i in 0..n {
116            if !seen[i] {
117                self.dfs_post(i, &mut seen, &mut order);
118            }
119        }
120
121        let mut rev_adj: Vec<Vec<usize>> = vec![Vec::new(); n];
122        for src in 0..n {
123            for &dst in &self.direct_callees[src] {
124                rev_adj[dst.0 as usize].push(src);
125            }
126        }
127
128        let mut comp = vec![usize::MAX; n];
129        let mut comp_nodes: Vec<Vec<usize>> = Vec::new();
130        for &u in order.iter().rev() {
131            if comp[u] != usize::MAX {
132                continue;
133            }
134            let id = comp_nodes.len();
135            let mut nodes = Vec::new();
136            let mut stack = vec![u];
137            comp[u] = id;
138            while let Some(x) = stack.pop() {
139                nodes.push(x);
140                for &p in &rev_adj[x] {
141                    if comp[p] == usize::MAX {
142                        comp[p] = id;
143                        stack.push(p);
144                    }
145                }
146            }
147            comp_nodes.push(nodes);
148        }
149
150        let c = comp_nodes.len();
151        let mut comp_succ: Vec<HashSet<usize>> = vec![HashSet::new(); c];
152        let mut indeg = vec![0usize; c];
153        for src in 0..n {
154            for &dst in &self.direct_callees[src] {
155                let a = comp[src];
156                let b = comp[dst.0 as usize];
157                if a != b && comp_succ[a].insert(b) {
158                    indeg[b] += 1;
159                }
160            }
161        }
162
163        // Topological order from roots to leaves; reverse for bottom-up.
164        let mut q: Vec<usize> = (0..c).filter(|&i| indeg[i] == 0).collect();
165        q.sort_unstable();
166        let mut topo = Vec::with_capacity(c);
167        while let Some(x) = q.pop() {
168            topo.push(x);
169            let mut succs: Vec<usize> = comp_succ[x].iter().copied().collect();
170            succs.sort_unstable();
171            for y in succs {
172                indeg[y] -= 1;
173                if indeg[y] == 0 {
174                    q.push(y);
175                }
176            }
177        }
178
179        topo.reverse();
180        topo.into_iter()
181            .map(|cid| {
182                let mut ids: Vec<FunctionId> = comp_nodes[cid]
183                    .iter()
184                    .map(|&x| FunctionId(x as u32))
185                    .collect();
186                ids.sort_unstable_by_key(|f| f.0);
187                ids
188            })
189            .collect()
190    }
191
192    fn dfs_post(&self, u: usize, seen: &mut [bool], order: &mut Vec<usize>) {
193        if seen[u] {
194            return;
195        }
196        seen[u] = true;
197        for &v in &self.direct_callees[u] {
198            self.dfs_post(v.0 as usize, seen, order);
199        }
200        order.push(u);
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207    use llvm_ir::{Builder, GlobalId, Linkage};
208
209    fn build_graph_module() -> (Context, Module) {
210        let mut ctx = Context::new();
211        let mut module = Module::new("m");
212        let mut b = Builder::new(&mut ctx, &mut module);
213        let i64_ty = b.ctx.i64_ty;
214        let fn_ty = b.ctx.mk_fn_type(i64_ty, vec![i64_ty], false);
215
216        b.add_function("a", i64_ty, vec![i64_ty], vec!["x".into()], false, Linkage::External);
217        let a_entry = b.add_block("a.entry");
218        b.position_at_end(a_entry);
219        let x = b.get_arg(0);
220        let c1 = b.build_call("c1", i64_ty, fn_ty, ValueRef::Global(GlobalId(1)), vec![x]);
221        let c2 = b.build_call("c2", i64_ty, fn_ty, ValueRef::Global(GlobalId(2)), vec![x]);
222        let sum = b.build_add("sum", c1, c2);
223        b.build_ret(sum);
224
225        b.add_function("b", i64_ty, vec![i64_ty], vec!["x".into()], false, Linkage::External);
226        let b_entry = b.add_block("b.entry");
227        b.position_at_end(b_entry);
228        let bx = b.get_arg(0);
229        let bcall = b.build_call("r", i64_ty, fn_ty, ValueRef::Global(GlobalId(2)), vec![bx]);
230        b.build_ret(bcall);
231
232        b.add_function("c", i64_ty, vec![i64_ty], vec!["x".into()], false, Linkage::External);
233        let c_entry = b.add_block("c.entry");
234        b.position_at_end(c_entry);
235        let cx = b.get_arg(0);
236        b.build_ret(cx);
237
238        (ctx, module)
239    }
240
241    #[test]
242    fn call_graph_builds_direct_edges_for_dag() {
243        let (ctx, module) = build_graph_module();
244        let cg = CallGraph::build(&ctx, &module);
245
246        assert_eq!(cg.callees(FunctionId(0)), &[FunctionId(1), FunctionId(2)]);
247        assert_eq!(cg.callees(FunctionId(1)), &[FunctionId(2)]);
248        assert!(cg.callees(FunctionId(2)).is_empty());
249        assert_eq!(cg.callers(FunctionId(2)), &[FunctionId(0), FunctionId(1)]);
250    }
251
252    #[test]
253    fn call_graph_sccs_group_cycle_together() {
254        let mut ctx = Context::new();
255        let mut module = Module::new("m");
256        let mut b = Builder::new(&mut ctx, &mut module);
257        let i64_ty = b.ctx.i64_ty;
258        let fn_ty = b.ctx.mk_fn_type(i64_ty, vec![i64_ty], false);
259        for name in ["a", "b", "c"] {
260            b.add_function(name, i64_ty, vec![i64_ty], vec!["x".into()], false, Linkage::External);
261            let entry = b.add_block(format!("{name}.entry"));
262            b.position_at_end(entry);
263            let x = b.get_arg(0);
264            let callee = match name {
265                "a" => ValueRef::Global(GlobalId(1)),
266                "b" => ValueRef::Global(GlobalId(0)),
267                _ => ValueRef::Global(GlobalId(2)),
268            };
269            let r = b.build_call("r", i64_ty, fn_ty, callee, vec![x]);
270            b.build_ret(r);
271        }
272
273        let cg = CallGraph::build(&ctx, &module);
274        let sccs = cg.sccs();
275        assert!(
276            sccs.iter()
277                .any(|s| s.len() == 2 && s.contains(&FunctionId(0)) && s.contains(&FunctionId(1)))
278        );
279    }
280
281    #[test]
282    fn call_graph_classifies_indirect_and_external_calls() {
283        let mut ctx = Context::new();
284        let mut module = Module::new("m");
285        let mut b = Builder::new(&mut ctx, &mut module);
286        let i64_ty = b.ctx.i64_ty;
287        let fn_ty = b.ctx.mk_fn_type(i64_ty, vec![i64_ty], false);
288        b.add_declaration("ext", i64_ty, vec![i64_ty], false);
289
290        b.add_function("f", i64_ty, vec![i64_ty], vec!["x".into()], false, Linkage::External);
291        let entry = b.add_block("entry");
292        b.position_at_end(entry);
293        let x = b.get_arg(0);
294        let _ext = b.build_call("ext", i64_ty, fn_ty, ValueRef::Global(GlobalId(0)), vec![x]);
295        let _ind = b.build_call("ind", i64_ty, fn_ty, x, vec![x]);
296        b.build_ret(x);
297
298        let cg = CallGraph::build(&ctx, &module);
299        let edges = cg.edges(FunctionId(1));
300        assert!(edges.iter().any(|e| e.kind == CallEdgeKind::External));
301        assert!(edges.iter().any(|e| e.kind == CallEdgeKind::Indirect));
302    }
303}
llvm_analysis/call_graph.rs

llvm_analysis/
call_graph.rs
