1use std::fmt;
10
11use rustc_data_structures::graph::{
12 ControlFlowGraph, DirectedGraph, Predecessors, StartNode, Successors, dominators,
13 dominators::Dominators, iterate, vec_graph::VecGraph,
14};
15use rustc_index::{
16 Idx,
17 bit_set::{DenseBitSet, MixedBitSet, SparseBitMatrix},
18};
19use smallvec::SmallVec;
20
21struct ReversedGraph<'a, G: ControlFlowGraph> {
22 graph: &'a G,
23 exit: G::Node,
24 unreachable: MixedBitSet<G::Node>,
25}
26
27impl<G: ControlFlowGraph> DirectedGraph for ReversedGraph<'_, G> {
28 type Node = G::Node;
29
30 fn num_nodes(&self) -> usize {
31 self.graph.num_nodes()
32 }
33}
34
35impl<G: ControlFlowGraph> StartNode for ReversedGraph<'_, G> {
36 fn start_node(&self) -> Self::Node {
37 self.exit
38 }
39}
40
41impl<G: ControlFlowGraph> Successors for ReversedGraph<'_, G> {
42 fn successors(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> {
43 self
44 .graph
45 .predecessors(node)
46 .filter(|bb| !self.unreachable.contains(*bb))
47 .collect::<SmallVec<[G::Node; 4]>>()
51 .into_iter()
52 }
53}
54
55impl<G: ControlFlowGraph> Predecessors for ReversedGraph<'_, G> {
56 fn predecessors(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> {
57 self
58 .graph
59 .successors(node)
60 .filter(|bb| !self.unreachable.contains(*bb))
61 .collect::<SmallVec<[G::Node; 4]>>()
62 .into_iter()
63 }
64}
65
66pub struct PostDominators<Node: Idx> {
68 dominators: Dominators<Node>,
69 num_nodes: usize,
70}
71
72impl<Node: Idx> PostDominators<Node> {
73 pub fn build<G: ControlFlowGraph<Node = Node>>(graph: &G, exit: Node) -> Self {
75 let num_nodes = graph.num_nodes();
76 let mut reversed = ReversedGraph {
77 graph,
78 exit,
79 unreachable: MixedBitSet::new_empty(num_nodes),
80 };
81
82 let reachable = iterate::post_order_from(&reversed, exit);
83 reversed.unreachable.insert_all();
84 for n in &reachable {
85 reversed.unreachable.remove(*n);
86 }
87
88 let dominators = dominators::dominators(&reversed);
89 PostDominators {
90 dominators,
91 num_nodes,
92 }
93 }
94
95 pub fn immediate_post_dominator(&self, node: Node) -> Option<Node> {
97 self.dominators.immediate_dominator(node)
98 }
99
100 pub fn post_dominators(&self, node: Node) -> Option<impl Iterator<Item = Node> + '_> {
102 let reachable = self.dominators.is_reachable(node);
103 reachable.then(move || {
104 (0 .. self.num_nodes)
105 .map(Node::new)
106 .filter(move |other| self.dominators.dominates(*other, node))
107 })
108 }
109}
110
111pub struct ControlDependencies<Node: Idx>(SparseBitMatrix<Node, Node>);
113
114impl<Node: Idx> fmt::Debug for ControlDependencies<Node> {
115 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116 for (i, (bb, bbs)) in self
117 .0
118 .rows()
119 .enumerate()
120 .filter_map(|(i, bb)| self.0.row(bb).map(move |bbs| (i, (bb, bbs))))
121 {
122 if i > 0 {
123 write!(f, ", ")?;
124 }
125 write!(f, "{bb:?}: {{")?;
126 for (j, bb2) in bbs.iter().enumerate() {
127 if j > 0 {
128 write!(f, ", ")?;
129 }
130 write!(f, "{bb2:?}")?;
131 }
132 write!(f, "}}")?;
133 }
134 Ok(())
135 }
136}
137
138impl<Node: Idx + Ord> ControlDependencies<Node> {
139 fn build<G: ControlFlowGraph<Node = Node>>(graph: &G, exit: Node) -> Self {
143 let post_dominators = PostDominators::build(graph, exit);
144 let idom = |node| post_dominators.immediate_post_dominator(node);
145
146 let n = graph.num_nodes();
147 let edges = (0 .. n)
148 .filter_map(|i| {
149 let node = Node::new(i);
150 Some((idom(node)?, node))
151 })
152 .collect::<Vec<_>>();
153
154 let dominator_tree: VecGraph<Node, true> = VecGraph::new(n, edges);
155
156 let traversal = iterate::post_order_from(&dominator_tree, exit);
157
158 let mut df = SparseBitMatrix::new(n);
160 for x in traversal {
161 let local = graph.predecessors(x);
162 let up = dominator_tree
163 .successors(x)
164 .iter()
165 .flat_map(|z| df.row(*z).into_iter().flat_map(|set| set.iter()));
166 let frontier = local
167 .chain(up)
168 .filter(|y| idom(*y).map(|yd| yd != x).unwrap_or(false))
169 .collect::<Vec<_>>();
170
171 for y in frontier {
172 df.insert(x, y);
173 }
174 }
175
176 ControlDependencies(df)
177 }
178
179 pub fn build_many<G: ControlFlowGraph<Node = Node>>(
181 graph: &G,
182 exits: impl IntoIterator<Item = Node>,
183 ) -> Self {
184 let mut all_deps = SparseBitMatrix::new(graph.num_nodes());
185 for exit in exits {
186 let deps = ControlDependencies::build(graph, exit);
187 for node in deps.0.rows() {
188 if let Some(set) = deps.0.row(node) {
189 all_deps.union_row(node, set);
190 }
191 }
192 }
193 ControlDependencies(all_deps)
194 }
195
196 pub fn dependent_on(&self, node: Node) -> Option<&DenseBitSet<Node>> {
198 self.0.row(node)
199 }
200}
201
202#[cfg(test)]
203mod test {
204 use log::debug;
205 use rustc_data_structures::fx::{FxHashMap as HashMap, FxHashSet as HashSet};
206 use rustc_middle::mir::Location;
207 use test_log::test;
208
209 use crate::{BodyExt, test_utils};
210
211 #[test]
212 fn test_control_dependencies() {
213 let input = r"
214fn main() {
215 let mut x = 1;
216 x = 2;
217 if true { x = 3; }
218 for _ in 0 .. 1 { x = 4; }
219 x = 5;
220}";
221 test_utils::compile_body(input, move |tcx, _, body_with_facts| {
222 let body = &body_with_facts.body;
223 let control_deps = body.control_dependencies();
224
225 let mut snippet_to_loc: HashMap<_, Vec<_>> = HashMap::default();
226 for loc in body.all_locations() {
227 let snippet = tcx
228 .sess
229 .source_map()
230 .span_to_snippet(body.source_info(loc).span)
231 .unwrap();
232 snippet_to_loc.entry(snippet).or_default().push(loc);
233 }
234 debug!("snippet_to_loc: {snippet_to_loc:#?}");
235 let pair = |s| (s, &snippet_to_loc[s]);
236
237 let x_eq_1 = pair("mut x");
238 let x_eq_2 = pair("x = 2");
239 let if_true = pair("true");
240 let x_eq_3 = pair("x = 3");
241 let for_in = pair("0 .. 1");
242 let x_eq_4 = pair("x = 4");
243 let x_eq_5 = pair("x = 5");
244
245 let is_dep_loc = |l1: Location, l2: Location| {
246 let is_terminator = body.stmt_at(l2).is_right();
247 is_terminator
248 && control_deps
249 .dependent_on(l1.block)
250 .map(|deps| deps.contains(l2.block))
251 .unwrap_or(false)
252 };
253
254 let is_dep = |l1: &[Location], l2: &[Location]| {
255 l1.iter().any(|l1| l2.iter().any(|l2| is_dep_loc(*l1, *l2)))
256 };
257
258 let all_locs = [x_eq_1, x_eq_2, if_true, x_eq_3, for_in, x_eq_4, x_eq_5]
259 .into_iter()
260 .collect::<HashSet<_>>();
261 let all_deps: &[(_, &[_])] = &[
262 (x_eq_1, &[]),
263 (x_eq_2, &[]),
264 (if_true, &[x_eq_3]),
265 (x_eq_3, &[]),
266 (for_in, &[x_eq_4]),
267 (x_eq_5, &[]),
268 ];
269
270 for ((parent_name, parent_locs), desired) in all_deps {
271 let desired = desired.iter().copied().collect::<HashSet<_>>();
272 for (child_name, child_locs) in &desired {
273 assert!(
274 is_dep(child_locs, parent_locs),
275 "{child_name} was not dependent on {parent_name}, but should be"
276 );
277 }
278
279 let remaining = &all_locs - &desired;
280 for (child_name, child_locs) in remaining {
281 if *parent_name == child_name {
282 continue;
283 }
284
285 assert!(
286 !is_dep(child_locs, parent_locs),
287 "{child_name} was dependent on {parent_name}, but should not be"
288 );
289 }
290 }
291 });
292 }
293}