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(core::cmp::min(level, (Self::INVALID as usize) - 1)) 66 .expect("Clamped value must always convert"), 67 ) 68 } 69 } 70 71 impl core::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