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::SecondaryMap;
6 use crate::entity::entity_impl;
7 use crate::entity::{Keys, PrimaryMap};
8 use crate::flowgraph::ControlFlowGraph;
9 use crate::ir::{Block, Function};
10 use crate::packed_option::PackedOption;
11 use crate::timing;
12 use alloc::vec::Vec;
13 use 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);
170         self.discover_loop_blocks(cfg, domtree);
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     // Determines if a block dominates any predecessor
194     // and thus is a loop header.
195     fn is_block_loop_header(block: Block, cfg: &ControlFlowGraph, domtree: &DominatorTree) -> bool {
196         // A block is a loop header if it dominates any of its predecessors.
197         cfg.pred_iter(block)
198             .any(|pred| domtree.block_dominates(block, pred.block))
199     }
200 
201     // Traverses the CFG in reverse postorder and create a loop object for every block having a
202     // back edge.
203     fn find_loop_headers(&mut self, cfg: &ControlFlowGraph, domtree: &DominatorTree) {
204         for &block in domtree
205             .cfg_rpo()
206             .filter(|&&block| Self::is_block_loop_header(block, cfg, domtree))
207         {
208             // This block is a loop header, so we create its associated loop
209             let lp = self.loops.push(LoopData::new(block, None));
210             self.block_loop_map[block] = lp.into();
211         }
212     }
213 
214     // Intended to be called after `find_loop_headers`. For each detected loop header,
215     // discovers all the block belonging to the loop and its inner loops. After a call to this
216     // function, the loop tree is fully constructed.
217     fn discover_loop_blocks(&mut self, cfg: &ControlFlowGraph, domtree: &DominatorTree) {
218         let mut stack: Vec<Block> = Vec::new();
219         // We handle each loop header in reverse order, corresponding to a pseudo postorder
220         // traversal of the graph.
221         for lp in self.loops().rev() {
222             // Push all predecessors of this header that it dominates onto the stack.
223             stack.extend(
224                 cfg.pred_iter(self.loops[lp].header)
225                     .filter(|pred| {
226                         // We follow the back edges
227                         domtree.block_dominates(self.loops[lp].header, pred.block)
228                     })
229                     .map(|pred| pred.block),
230             );
231             while let Some(node) = stack.pop() {
232                 let continue_dfs: Option<Block>;
233                 match self.block_loop_map[node].expand() {
234                     None => {
235                         // The node hasn't been visited yet, we tag it as part of the loop
236                         self.block_loop_map[node] = PackedOption::from(lp);
237                         continue_dfs = Some(node);
238                     }
239                     Some(node_loop) => {
240                         // We copy the node_loop into a mutable reference passed along the while
241                         let mut node_loop = node_loop;
242                         // The node is part of a loop, which can be lp or an inner loop
243                         let mut node_loop_parent_option = self.loops[node_loop].parent;
244                         while let Some(node_loop_parent) = node_loop_parent_option.expand() {
245                             if node_loop_parent == lp {
246                                 // We have encountered lp so we stop (already visited)
247                                 break;
248                             } else {
249                                 //
250                                 node_loop = node_loop_parent;
251                                 // We lookup the parent loop
252                                 node_loop_parent_option = self.loops[node_loop].parent;
253                             }
254                         }
255                         // Now node_loop_parent is either:
256                         // - None and node_loop is an new inner loop of lp
257                         // - Some(...) and the initial node_loop was a known inner loop of lp
258                         match node_loop_parent_option.expand() {
259                             Some(_) => continue_dfs = None,
260                             None => {
261                                 if node_loop != lp {
262                                     self.loops[node_loop].parent = lp.into();
263                                     continue_dfs = Some(self.loops[node_loop].header)
264                                 } else {
265                                     // If lp is a one-block loop then we make sure we stop
266                                     continue_dfs = None
267                                 }
268                             }
269                         }
270                     }
271                 }
272                 // Now we have handled the popped node and need to continue the DFS by adding the
273                 // predecessors of that node
274                 if let Some(continue_dfs) = continue_dfs {
275                     stack.extend(cfg.pred_iter(continue_dfs).map(|pred| pred.block));
276                 }
277             }
278         }
279     }
280 
281     fn assign_loop_levels(&mut self) {
282         let mut stack: SmallVec<[Loop; 8]> = SmallVec::new();
283         for lp in self.loops.keys() {
284             if self.loops[lp].level == LoopLevel::invalid() {
285                 stack.push(lp);
286                 while let Some(&lp) = stack.last() {
287                     if let Some(parent) = self.loops[lp].parent.into() {
288                         if self.loops[parent].level != LoopLevel::invalid() {
289                             self.loops[lp].level = self.loops[parent].level.inc();
290                             stack.pop();
291                         } else {
292                             stack.push(parent);
293                         }
294                     } else {
295                         self.loops[lp].level = LoopLevel::root().inc();
296                         stack.pop();
297                     }
298                 }
299             }
300         }
301     }
302 }
303 
304 #[cfg(test)]
305 mod tests {
306     use crate::cursor::{Cursor, FuncCursor};
307     use crate::dominator_tree::DominatorTree;
308     use crate::flowgraph::ControlFlowGraph;
309     use crate::ir::{Function, InstBuilder, types};
310     use crate::loop_analysis::{Loop, LoopAnalysis};
311     use alloc::vec::Vec;
312 
313     #[test]
314     fn nested_loops_detection() {
315         let mut func = Function::new();
316         let block0 = func.dfg.make_block();
317         let block1 = func.dfg.make_block();
318         let block2 = func.dfg.make_block();
319         let block3 = func.dfg.make_block();
320         let block4 = func.dfg.make_block();
321         let cond = func.dfg.append_block_param(block0, types::I32);
322 
323         {
324             let mut cur = FuncCursor::new(&mut func);
325 
326             cur.insert_block(block0);
327             cur.ins().jump(block1, &[]);
328 
329             cur.insert_block(block1);
330             cur.ins().jump(block2, &[]);
331 
332             cur.insert_block(block2);
333             cur.ins().brif(cond, block1, &[], block3, &[]);
334 
335             cur.insert_block(block3);
336             cur.ins().brif(cond, block0, &[], block4, &[]);
337 
338             cur.insert_block(block4);
339             cur.ins().return_(&[]);
340         }
341 
342         let mut loop_analysis = LoopAnalysis::new();
343         let mut cfg = ControlFlowGraph::new();
344         let mut domtree = DominatorTree::new();
345         cfg.compute(&func);
346         domtree.compute(&func, &cfg);
347         loop_analysis.compute(&func, &cfg, &domtree);
348 
349         let loops = loop_analysis.loops().collect::<Vec<Loop>>();
350         assert_eq!(loops.len(), 2);
351         assert_eq!(loop_analysis.loop_header(loops[0]), block0);
352         assert_eq!(loop_analysis.loop_header(loops[1]), block1);
353         assert_eq!(loop_analysis.loop_parent(loops[1]), Some(loops[0]));
354         assert_eq!(loop_analysis.loop_parent(loops[0]), None);
355         assert_eq!(loop_analysis.is_in_loop(block0, loops[0]), true);
356         assert_eq!(loop_analysis.is_in_loop(block0, loops[1]), false);
357         assert_eq!(loop_analysis.is_in_loop(block1, loops[1]), true);
358         assert_eq!(loop_analysis.is_in_loop(block1, loops[0]), true);
359         assert_eq!(loop_analysis.is_in_loop(block2, loops[1]), true);
360         assert_eq!(loop_analysis.is_in_loop(block2, loops[0]), true);
361         assert_eq!(loop_analysis.is_in_loop(block3, loops[0]), true);
362         assert_eq!(loop_analysis.is_in_loop(block0, loops[1]), false);
363         assert_eq!(loop_analysis.loop_level(block0).level(), 1);
364         assert_eq!(loop_analysis.loop_level(block1).level(), 2);
365         assert_eq!(loop_analysis.loop_level(block2).level(), 2);
366         assert_eq!(loop_analysis.loop_level(block3).level(), 1);
367     }
368 
369     #[test]
370     fn complex_loop_detection() {
371         let mut func = Function::new();
372         let block0 = func.dfg.make_block();
373         let block1 = func.dfg.make_block();
374         let block2 = func.dfg.make_block();
375         let block3 = func.dfg.make_block();
376         let block4 = func.dfg.make_block();
377         let block5 = func.dfg.make_block();
378         let block6 = func.dfg.make_block();
379         let cond = func.dfg.append_block_param(block0, types::I32);
380 
381         {
382             let mut cur = FuncCursor::new(&mut func);
383 
384             cur.insert_block(block0);
385             cur.ins().brif(cond, block1, &[], block3, &[]);
386 
387             cur.insert_block(block1);
388             cur.ins().jump(block2, &[]);
389 
390             cur.insert_block(block2);
391             cur.ins().brif(cond, block1, &[], block5, &[]);
392 
393             cur.insert_block(block3);
394             cur.ins().jump(block4, &[]);
395 
396             cur.insert_block(block4);
397             cur.ins().brif(cond, block3, &[], block5, &[]);
398 
399             cur.insert_block(block5);
400             cur.ins().brif(cond, block0, &[], block6, &[]);
401 
402             cur.insert_block(block6);
403             cur.ins().return_(&[]);
404         }
405 
406         let mut loop_analysis = LoopAnalysis::new();
407         let cfg = ControlFlowGraph::with_function(&func);
408         let domtree = DominatorTree::with_function(&func, &cfg);
409         loop_analysis.compute(&func, &cfg, &domtree);
410 
411         let loops = loop_analysis.loops().collect::<Vec<Loop>>();
412         assert_eq!(loops.len(), 3);
413         assert_eq!(loop_analysis.loop_header(loops[0]), block0);
414         assert_eq!(loop_analysis.loop_header(loops[1]), block3);
415         assert_eq!(loop_analysis.loop_header(loops[2]), block1);
416         assert_eq!(loop_analysis.loop_parent(loops[1]), Some(loops[0]));
417         assert_eq!(loop_analysis.loop_parent(loops[2]), Some(loops[0]));
418         assert_eq!(loop_analysis.loop_parent(loops[0]), None);
419         assert_eq!(loop_analysis.is_in_loop(block0, loops[0]), true);
420         assert_eq!(loop_analysis.is_in_loop(block1, loops[2]), true);
421         assert_eq!(loop_analysis.is_in_loop(block2, loops[2]), true);
422         assert_eq!(loop_analysis.is_in_loop(block3, loops[1]), true);
423         assert_eq!(loop_analysis.is_in_loop(block4, loops[1]), true);
424         assert_eq!(loop_analysis.is_in_loop(block5, loops[0]), true);
425         assert_eq!(loop_analysis.loop_level(block0).level(), 1);
426         assert_eq!(loop_analysis.loop_level(block1).level(), 2);
427         assert_eq!(loop_analysis.loop_level(block2).level(), 2);
428         assert_eq!(loop_analysis.loop_level(block3).level(), 2);
429         assert_eq!(loop_analysis.loop_level(block4).level(), 2);
430         assert_eq!(loop_analysis.loop_level(block5).level(), 1);
431     }
432 }
433