rustc_codegen_spirv/linker/
ipo.rs1use indexmap::IndexSet;
6use rspirv::dr::Module;
7use rspirv::spirv::{Op, Word};
8use rustc_data_structures::fx::FxHashMap;
9
10type FuncIdx = usize;
12
13pub struct CallGraph {
14 pub entry_points: IndexSet<FuncIdx>,
15
16 pub callees: Vec<IndexSet<FuncIdx>>,
18}
19
20impl CallGraph {
21 pub fn collect(module: &Module) -> Self {
22 Self::collect_with_func_id_to_idx(module).0
23 }
24 pub fn collect_with_func_id_to_idx(module: &Module) -> (Self, FxHashMap<Word, FuncIdx>) {
25 let func_id_to_idx: FxHashMap<_, _> = module
26 .functions
27 .iter()
28 .enumerate()
29 .map(|(i, func)| (func.def_id().unwrap(), i))
30 .collect();
31 let entry_points = module
32 .entry_points
33 .iter()
34 .map(|entry| {
35 assert_eq!(entry.class.opcode, Op::EntryPoint);
36 func_id_to_idx[&entry.operands[1].unwrap_id_ref()]
37 })
38 .collect();
39 let callees = module
40 .functions
41 .iter()
42 .map(|func| {
43 func.all_inst_iter()
44 .filter(|inst| inst.class.opcode == Op::FunctionCall)
45 .filter_map(|inst| {
46 func_id_to_idx
51 .get(&inst.operands[0].unwrap_id_ref())
52 .copied()
53 })
54 .collect()
55 })
56 .collect();
57 (
58 Self {
59 entry_points,
60 callees,
61 },
62 func_id_to_idx,
63 )
64 }
65
66 pub fn post_order(&self) -> Vec<FuncIdx> {
70 let num_funcs = self.callees.len();
71
72 let mut visited = vec![false; num_funcs];
74 let mut post_order = Vec::with_capacity(num_funcs);
75
76 for &entry in &self.entry_points {
78 self.post_order_step(entry, &mut visited, &mut post_order);
79 }
80
81 for func in 0..num_funcs {
84 if !visited[func] {
85 self.post_order_step(func, &mut visited, &mut post_order);
86 }
87 }
88
89 post_order
90 }
91
92 fn post_order_step(&self, func: FuncIdx, visited: &mut [bool], post_order: &mut Vec<FuncIdx>) {
93 if visited[func] {
94 return;
95 }
96 visited[func] = true;
97
98 for &callee in &self.callees[func] {
99 self.post_order_step(callee, visited, post_order);
100 }
101
102 post_order.push(func);
103 }
104}