1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
//! Tools for interprocedural optimizations (aka "IPO"s).
// FIXME(eddyb) perhaps make all IPOs sub-modules of this module?
use indexmap::IndexSet;
use rspirv::dr::Module;
use rspirv::spirv::Op;
use rustc_data_structures::fx::FxHashMap;
// FIXME(eddyb) use newtyped indices and `IndexVec`.
type FuncIdx = usize;
pub struct CallGraph {
pub entry_points: IndexSet<FuncIdx>,
/// `callees[i].contains(j)` implies `functions[i]` calls `functions[j]`.
pub callees: Vec<IndexSet<FuncIdx>>,
}
impl CallGraph {
pub fn collect(module: &Module) -> Self {
let func_id_to_idx: FxHashMap<_, _> = module
.functions
.iter()
.enumerate()
.map(|(i, func)| (func.def_id().unwrap(), i))
.collect();
let entry_points = module
.entry_points
.iter()
.map(|entry| {
assert_eq!(entry.class.opcode, Op::EntryPoint);
func_id_to_idx[&entry.operands[1].unwrap_id_ref()]
})
.collect();
let callees = module
.functions
.iter()
.map(|func| {
func.all_inst_iter()
.filter(|inst| inst.class.opcode == Op::FunctionCall)
.filter_map(|inst| {
// FIXME(eddyb) `func_id_to_idx` should always have an
// entry for a callee ID, but when ran early enough
// (before zombie removal), the callee ID might not
// point to an `OpFunction` (unsure what, `OpUndef`?).
func_id_to_idx
.get(&inst.operands[0].unwrap_id_ref())
.copied()
})
.collect()
})
.collect();
Self {
entry_points,
callees,
}
}
/// Order functions using a post-order traversal, i.e. callees before callers.
// FIXME(eddyb) replace this with `rustc_data_structures::graph::iterate`
// (or similar).
pub fn post_order(&self) -> Vec<FuncIdx> {
let num_funcs = self.callees.len();
// FIXME(eddyb) use a proper bitset.
let mut visited = vec![false; num_funcs];
let mut post_order = Vec::with_capacity(num_funcs);
// Visit the call graph with entry points as roots.
for &entry in &self.entry_points {
self.post_order_step(entry, &mut visited, &mut post_order);
}
// Also visit any functions that were not reached from entry points
// (they might be dead but they should be processed nonetheless).
for func in 0..num_funcs {
if !visited[func] {
self.post_order_step(func, &mut visited, &mut post_order);
}
}
post_order
}
fn post_order_step(&self, func: FuncIdx, visited: &mut [bool], post_order: &mut Vec<FuncIdx>) {
if visited[func] {
return;
}
visited[func] = true;
for &callee in &self.callees[func] {
self.post_order_step(callee, visited, post_order);
}
post_order.push(func);
}
}