1 //! A loop analysis represented as mappings of loops to their header Block
2 //! and parent in the loop tree.
3 
4 use crate::dominator_tree::DominatorTree;
5 use crate::entity::entity_impl;
6 use crate::entity::SecondaryMap;
7 use crate::entity::{Keys, PrimaryMap};
8 use crate::flowgraph::{BlockPredecessor, ControlFlowGraph};
9 use crate::ir::{Block, Function, Layout};
10 use crate::packed_option::PackedOption;
11 use crate::timing;
12 use alloc::vec::Vec;
13 use smallvec::{smallvec, SmallVec};
14 
15 /// A opaque reference to a code loop.
16 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
17 pub struct Loop(u32);
18 entity_impl!(Loop, "loop");
19 
20 /// Loop tree information for a single function.
21 ///
22 /// Loops are referenced by the Loop object, and for each loop you can access its header block,
23 /// its eventual parent in the loop tree and all the block belonging to the loop.
24 pub struct LoopAnalysis {
25     loops: PrimaryMap<Loop, LoopData>,
26     block_loop_map: SecondaryMap<Block, PackedOption<Loop>>,
27     valid: bool,
28 }
29 
30 struct LoopData {
31     header: Block,
32     parent: PackedOption<Loop>,
33     level: LoopLevel,
34 }
35 
36 /// A level in a loop nest.
37 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
38 pub struct LoopLevel(u8);
39 impl LoopLevel {
40     const INVALID: u8 = u8::MAX;
41 
42     /// Get the root level (no loop).
43     pub fn root() -> Self {
44         Self(0)
45     }
46     /// Get the loop level.
47     pub fn level(self) -> usize {
48         self.0 as usize
49     }
50     /// Invalid loop level.
51     pub fn invalid() -> Self {
52         Self(Self::INVALID)
53     }
54     /// One loop level deeper.
55     pub fn inc(self) -> Self {
56         if self.0 == (Self::INVALID - 1) {
57             self
58         } else {
59             Self(self.0 + 1)
60         }
61     }
62     /// A clamped loop level from a larger-width (usize) depth.
63     pub fn clamped(level: usize) -> Self {
64         Self(
65             u8::try_from(std::cmp::min(level, (Self::INVALID as usize) - 1))
66                 .expect("Clamped value must always convert"),
67         )
68     }
69 }
70 
71 impl std::default::Default for LoopLevel {
72     fn default() -> Self {
73         LoopLevel::invalid()
74     }
75 }
76 
77 impl LoopData {
78     /// Creates a `LoopData` object with the loop header and its eventual parent in the loop tree.
79     pub fn new(header: Block, parent: Option<Loop>) -> Self {
80         Self {
81             header,
82             parent: parent.into(),
83             level: LoopLevel::invalid(),
84         }
85     }
86 }
87 
88 /// Methods for querying the loop analysis.
89 impl LoopAnalysis {
90     /// Allocate a new blank loop analysis struct. Use `compute` to compute the loop analysis for
91     /// a function.
92     pub fn new() -> Self {
93         Self {
94             valid: false,
95             loops: PrimaryMap::new(),
96             block_loop_map: SecondaryMap::new(),
97         }
98     }
99 
100     /// Returns all the loops contained in a function.
101     pub fn loops(&self) -> Keys<Loop> {
102         self.loops.keys()
103     }
104 
105     /// Returns the header block of a particular loop.
106     ///
107     /// The characteristic property of a loop header block is that it dominates some of its
108     /// predecessors.
109     pub fn loop_header(&self, lp: Loop) -> Block {
110         self.loops[lp].header
111     }
112 
113     /// Return the eventual parent of a loop in the loop tree.
114     pub fn loop_parent(&self, lp: Loop) -> Option<Loop> {
115         self.loops[lp].parent.expand()
116     }
117 
118     /// Return the innermost loop for a given block.
119     pub fn innermost_loop(&self, block: Block) -> Option<Loop> {
120         self.block_loop_map[block].expand()
121     }
122 
123     /// Determine if a Block is a loop header. If so, return the loop.
124     pub fn is_loop_header(&self, block: Block) -> Option<Loop> {
125         self.innermost_loop(block)
126             .filter(|&lp| self.loop_header(lp) == block)
127     }
128 
129     /// Determine if a Block belongs to a loop by running a finger along the loop tree.
130     ///
131     /// Returns `true` if `block` is in loop `lp`.
132     pub fn is_in_loop(&self, block: Block, lp: Loop) -> bool {
133         let block_loop = self.block_loop_map[block];
134         match block_loop.expand() {
135             None => false,
136             Some(block_loop) => self.is_child_loop(block_loop, lp),
137         }
138     }
139 
140     /// Determines if a loop is contained in another loop.
141     ///
142     /// `is_child_loop(child,parent)` returns `true` if and only if `child` is a child loop of
143     /// `parent` (or `child == parent`).
144     pub fn is_child_loop(&self, child: Loop, parent: Loop) -> bool {
145         let mut finger = Some(child);
146         while let Some(finger_loop) = finger {
147             if finger_loop == parent {
148                 return true;
149             }
150             finger = self.loop_parent(finger_loop);
151         }
152         false
153     }
154 
155     /// Returns the loop-nest level of a given block.
156     pub fn loop_level(&self, block: Block) -> LoopLevel {
157         self.innermost_loop(block)
158             .map_or(LoopLevel(0), |lp| self.loops[lp].level)
159     }
160 }
161 
162 impl LoopAnalysis {
163     /// Detects the loops in a function. Needs the control flow graph and the dominator tree.
164     pub fn compute(&mut self, func: &Function, cfg: &ControlFlowGraph, domtree: &DominatorTree) {
165         let _tt = timing::loop_analysis();
166         self.loops.clear();
167         self.block_loop_map.clear();
168         self.block_loop_map.resize(func.dfg.num_blocks());
169         self.find_loop_headers(cfg, domtree, &func.layout);
170         self.discover_loop_blocks(cfg, domtree, &func.layout);
171         self.assign_loop_levels();
172         self.valid = true;
173     }
174 
175     /// Check if the loop analysis is in a valid state.
176     ///
177     /// Note that this doesn't perform any kind of validity checks. It simply checks if the
178     /// `compute()` method has been called since the last `clear()`. It does not check that the
179     /// loop analysis is consistent with the CFG.
180     pub fn is_valid(&self) -> bool {
181         self.valid
182     }
183 
184     /// Clear all the data structures contained in the loop analysis. This will leave the
185     /// analysis in a similar state to a context returned by `new()` except that allocated
186     /// memory be retained.
187     pub fn clear(&mut self) {
188         self.loops.clear();
189         self.block_loop_map.clear();
190         self.valid = false;
191     }
192 
193     // Traverses the CFG in reverse postorder and create a loop object for every block having a
194     // back edge.
195     fn find_loop_headers(
196         &mut self,
197         cfg: &ControlFlowGraph,
198         domtree: &DominatorTree,
199         layout: &Layout,
200     ) {
201         // We traverse the CFG in reverse postorder
202         for &block in domtree.cfg_postorder().iter().rev() {
203             for BlockPredecessor {
204                 inst: pred_inst, ..
205             } in cfg.pred_iter(block)
206             {
207                 // If the block dominates one of its predecessors it is a back edge
208                 if domtree.dominates(block, pred_inst, layout) {
209                     // This block is a loop header, so we create its associated loop
210                     let lp = self.loops.push(LoopData::new(block, None));
211                     self.block_loop_map[block] = lp.into();
212                     break;
213                     // We break because we only need one back edge to identify a loop header.
214                 }
215             }
216         }
217     }
218 
219     // Intended to be called after `find_loop_headers`. For each detected loop header,
220     // discovers all the block belonging to the loop and its inner loops. After a call to this
221     // function, the loop tree is fully constructed.
222     fn discover_loop_blocks(
223         &mut self,
224         cfg: &ControlFlowGraph,
225         domtree: &DominatorTree,
226         layout: &Layout,
227     ) {
228         let mut stack: Vec<Block> = Vec::new();
229         // We handle each loop header in reverse order, corresponding to a pseudo postorder
230         // traversal of the graph.
231         for lp in self.loops().rev() {
232             for BlockPredecessor {
233                 block: pred,
234                 inst: pred_inst,
235             } in cfg.pred_iter(self.loops[lp].header)
236             {
237                 // We follow the back edges
238                 if domtree.dominates(self.loops[lp].header, pred_inst, layout) {
239                     stack.push(pred);
240                 }
241             }
242             while let Some(node) = stack.pop() {
243                 let continue_dfs: Option<Block>;
244                 match self.block_loop_map[node].expand() {
245                     None => {
246                         // The node hasn't been visited yet, we tag it as part of the loop
247                         self.block_loop_map[node] = PackedOption::from(lp);
248                         continue_dfs = Some(node);
249                     }
250                     Some(node_loop) => {
251                         // We copy the node_loop into a mutable reference passed along the while
252                         let mut node_loop = node_loop;
253                         // The node is part of a loop, which can be lp or an inner loop
254                         let mut node_loop_parent_option = self.loops[node_loop].parent;
255                         while let Some(node_loop_parent) = node_loop_parent_option.expand() {
256                             if node_loop_parent == lp {
257                                 // We have encountered lp so we stop (already visited)
258                                 break;
259                             } else {
260                                 //
261                                 node_loop = node_loop_parent;
262                                 // We lookup the parent loop
263                                 node_loop_parent_option = self.loops[node_loop].parent;
264                             }
265                         }
266                         // Now node_loop_parent is either:
267                         // - None and node_loop is an new inner loop of lp
268                         // - Some(...) and the initial node_loop was a known inner loop of lp
269                         match node_loop_parent_option.expand() {
270                             Some(_) => continue_dfs = None,
271                             None => {
272                                 if node_loop != lp {
273                                     self.loops[node_loop].parent = lp.into();
274                                     continue_dfs = Some(self.loops[node_loop].header)
275                                 } else {
276                                     // If lp is a one-block loop then we make sure we stop
277                                     continue_dfs = None
278                                 }
279                             }
280                         }
281                     }
282                 }
283                 // Now we have handled the popped node and need to continue the DFS by adding the
284                 // predecessors of that node
285                 if let Some(continue_dfs) = continue_dfs {
286                     for BlockPredecessor { block: pred, .. } in cfg.pred_iter(continue_dfs) {
287                         stack.push(pred)
288                     }
289                 }
290             }
291         }
292     }
293 
294     fn assign_loop_levels(&mut self) {
295         let mut stack: SmallVec<[Loop; 8]> = smallvec![];
296         for lp in self.loops.keys() {
297             if self.loops[lp].level == LoopLevel::invalid() {
298                 stack.push(lp);
299                 while let Some(&lp) = stack.last() {
300                     if let Some(parent) = self.loops[lp].parent.into() {
301                         if self.loops[parent].level != LoopLevel::invalid() {
302                             self.loops[lp].level = self.loops[parent].level.inc();
303                             stack.pop();
304                         } else {
305                             stack.push(parent);
306                         }
307                     } else {
308                         self.loops[lp].level = LoopLevel::root().inc();
309                         stack.pop();
310                     }
311                 }
312             }
313         }
314     }
315 }
316 
317 #[cfg(test)]
318 mod tests {
319     use crate::cursor::{Cursor, FuncCursor};
320     use crate::dominator_tree::DominatorTree;
321     use crate::flowgraph::ControlFlowGraph;
322     use crate::ir::{types, Function, InstBuilder};
323     use crate::loop_analysis::{Loop, LoopAnalysis};
324     use alloc::vec::Vec;
325 
326     #[test]
327     fn nested_loops_detection() {
328         let mut func = Function::new();
329         let block0 = func.dfg.make_block();
330         let block1 = func.dfg.make_block();
331         let block2 = func.dfg.make_block();
332         let block3 = func.dfg.make_block();
333         let block4 = func.dfg.make_block();
334         let cond = func.dfg.append_block_param(block0, types::I32);
335 
336         {
337             let mut cur = FuncCursor::new(&mut func);
338 
339             cur.insert_block(block0);
340             cur.ins().jump(block1, &[]);
341 
342             cur.insert_block(block1);
343             cur.ins().jump(block2, &[]);
344 
345             cur.insert_block(block2);
346             cur.ins().brif(cond, block1, &[], block3, &[]);
347 
348             cur.insert_block(block3);
349             cur.ins().brif(cond, block0, &[], block4, &[]);
350 
351             cur.insert_block(block4);
352             cur.ins().return_(&[]);
353         }
354 
355         let mut loop_analysis = LoopAnalysis::new();
356         let mut cfg = ControlFlowGraph::new();
357         let mut domtree = DominatorTree::new();
358         cfg.compute(&func);
359         domtree.compute(&func, &cfg);
360         loop_analysis.compute(&func, &cfg, &domtree);
361 
362         let loops = loop_analysis.loops().collect::<Vec<Loop>>();
363         assert_eq!(loops.len(), 2);
364         assert_eq!(loop_analysis.loop_header(loops[0]), block0);
365         assert_eq!(loop_analysis.loop_header(loops[1]), block1);
366         assert_eq!(loop_analysis.loop_parent(loops[1]), Some(loops[0]));
367         assert_eq!(loop_analysis.loop_parent(loops[0]), None);
368         assert_eq!(loop_analysis.is_in_loop(block0, loops[0]), true);
369         assert_eq!(loop_analysis.is_in_loop(block0, loops[1]), false);
370         assert_eq!(loop_analysis.is_in_loop(block1, loops[1]), true);
371         assert_eq!(loop_analysis.is_in_loop(block1, loops[0]), true);
372         assert_eq!(loop_analysis.is_in_loop(block2, loops[1]), true);
373         assert_eq!(loop_analysis.is_in_loop(block2, loops[0]), true);
374         assert_eq!(loop_analysis.is_in_loop(block3, loops[0]), true);
375         assert_eq!(loop_analysis.is_in_loop(block0, loops[1]), false);
376         assert_eq!(loop_analysis.loop_level(block0).level(), 1);
377         assert_eq!(loop_analysis.loop_level(block1).level(), 2);
378         assert_eq!(loop_analysis.loop_level(block2).level(), 2);
379         assert_eq!(loop_analysis.loop_level(block3).level(), 1);
380     }
381 
382     #[test]
383     fn complex_loop_detection() {
384         let mut func = Function::new();
385         let block0 = func.dfg.make_block();
386         let block1 = func.dfg.make_block();
387         let block2 = func.dfg.make_block();
388         let block3 = func.dfg.make_block();
389         let block4 = func.dfg.make_block();
390         let block5 = func.dfg.make_block();
391         let block6 = func.dfg.make_block();
392         let cond = func.dfg.append_block_param(block0, types::I32);
393 
394         {
395             let mut cur = FuncCursor::new(&mut func);
396 
397             cur.insert_block(block0);
398             cur.ins().brif(cond, block1, &[], block3, &[]);
399 
400             cur.insert_block(block1);
401             cur.ins().jump(block2, &[]);
402 
403             cur.insert_block(block2);
404             cur.ins().brif(cond, block1, &[], block5, &[]);
405 
406             cur.insert_block(block3);
407             cur.ins().jump(block4, &[]);
408 
409             cur.insert_block(block4);
410             cur.ins().brif(cond, block3, &[], block5, &[]);
411 
412             cur.insert_block(block5);
413             cur.ins().brif(cond, block0, &[], block6, &[]);
414 
415             cur.insert_block(block6);
416             cur.ins().return_(&[]);
417         }
418 
419         let mut loop_analysis = LoopAnalysis::new();
420         let mut cfg = ControlFlowGraph::new();
421         let mut domtree = DominatorTree::new();
422         cfg.compute(&func);
423         domtree.compute(&func, &cfg);
424         loop_analysis.compute(&func, &cfg, &domtree);
425 
426         let loops = loop_analysis.loops().collect::<Vec<Loop>>();
427         assert_eq!(loops.len(), 3);
428         assert_eq!(loop_analysis.loop_header(loops[0]), block0);
429         assert_eq!(loop_analysis.loop_header(loops[1]), block1);
430         assert_eq!(loop_analysis.loop_header(loops[2]), block3);
431         assert_eq!(loop_analysis.loop_parent(loops[1]), Some(loops[0]));
432         assert_eq!(loop_analysis.loop_parent(loops[2]), Some(loops[0]));
433         assert_eq!(loop_analysis.loop_parent(loops[0]), None);
434         assert_eq!(loop_analysis.is_in_loop(block0, loops[0]), true);
435         assert_eq!(loop_analysis.is_in_loop(block1, loops[1]), true);
436         assert_eq!(loop_analysis.is_in_loop(block2, loops[1]), true);
437         assert_eq!(loop_analysis.is_in_loop(block3, loops[2]), true);
438         assert_eq!(loop_analysis.is_in_loop(block4, loops[2]), true);
439         assert_eq!(loop_analysis.is_in_loop(block5, loops[0]), true);
440         assert_eq!(loop_analysis.loop_level(block0).level(), 1);
441         assert_eq!(loop_analysis.loop_level(block1).level(), 2);
442         assert_eq!(loop_analysis.loop_level(block2).level(), 2);
443         assert_eq!(loop_analysis.loop_level(block3).level(), 2);
444         assert_eq!(loop_analysis.loop_level(block4).level(), 2);
445         assert_eq!(loop_analysis.loop_level(block5).level(), 1);
446     }
447 }
448