Skip to main content

llvm_analysis/
loops.rs

1//! Natural loop detection built on the dominator tree.
2//!
3//! Algorithm:
4//! 1. Walk the CFG with DFS to find back-edges — edges (n → h) where h
5//!    dominates n.
6//! 2. For each back-edge, collect the natural loop body via reverse-CFG BFS
7//!    from the tail up to (and including) the header.
8//! 3. Assign loop nesting: a loop is a child of the smallest loop whose body
9//!    contains the child's header.
10//!
11//! # Reducibility requirement
12//!
13//! This algorithm detects **natural loops** only, which requires the CFG to be
14//! **reducible**.  A CFG is reducible when every strongly-connected component
15//! has a single entry node that dominates all other nodes in the component.
16//! Irreducible CFGs — those containing a
17//! strongly-connected component with two or more entry edges that do not
18//! dominate each other — are not handled correctly: the multi-entry cycle will
19//! not be detected, and `LoopInfo` will report zero loops for it.
20//!
21//! Irreducible CFGs are uncommon in practice (most front-ends and optimisers
22//! avoid them), but the LLVM IR format permits them.  Callers that need
23//! correct results on irreducible graphs should either:
24//! * structurise the CFG before calling `LoopInfo::compute`, or
25//! * use a full SCC-based (Havlak/Sreedhar) algorithm instead.
26
27use crate::cfg::Cfg;
28use crate::dominators::DomTree;
29use llvm_ir::{BlockId, Function};
30use std::collections::{HashMap, HashSet, VecDeque};
31
32/// A single natural loop.
33#[derive(Debug)]
34pub struct Loop {
35    /// The loop header (target of the back-edge).
36    pub header: BlockId,
37    /// All blocks in the loop body, including the header. Sorted by BlockId.
38    pub body: Vec<BlockId>,
39    /// Index of the immediately enclosing loop in `LoopInfo::loops`, if any.
40    pub parent: Option<usize>,
41}
42
43/// Loop nesting information for a single function.
44pub struct LoopInfo {
45    // `loops` field.
46    loops: Vec<Loop>,
47    /// Maps each block to the index of its innermost containing loop.
48    block_loop: HashMap<BlockId, usize>,
49}
50
51impl LoopInfo {
52    /// Detect all natural loops in `func`.
53    ///
54    /// Uses the standard back-edge + reverse-CFG-BFS algorithm, which is
55    /// correct only for **reducible** CFGs.  On an irreducible CFG, cycles
56    /// whose strongly-connected component has no single dominating header will
57    /// not be reported.  See the [module-level documentation](self) for
58    /// details and mitigation options.
59    pub fn compute(func: &Function, cfg: &Cfg, dom: &DomTree) -> Self {
60        if func.num_blocks() == 0 {
61            return LoopInfo {
62                loops: vec![],
63                block_loop: HashMap::new(),
64            };
65        }
66
67        // Find back-edges and build one loop per back-edge.
68        let back_edges = Self::find_back_edges(cfg, dom);
69        let mut loops: Vec<Loop> = back_edges
70            .into_iter()
71            .map(|(tail, header)| Loop {
72                header,
73                body: Self::collect_loop_body(tail, header, cfg),
74                parent: None,
75            })
76            .collect();
77
78        // Sort largest-body first so parent search is straightforward.
79        loops.sort_by(|a, b| b.body.len().cmp(&a.body.len()));
80
81        // Assign parent: the innermost (smallest body) loop that contains
82        // this loop's header, excluding itself.
83        for i in 0..loops.len() {
84            let header = loops[i].header;
85            loops[i].parent = loops
86                .iter()
87                .enumerate()
88                .filter(|(j, l)| *j != i && l.body.contains(&header))
89                .min_by_key(|(_, l)| l.body.len())
90                .map(|(j, _)| j);
91        }
92
93        // Build block → innermost loop map.
94        let mut block_loop: HashMap<BlockId, usize> = HashMap::new();
95        for (i, lp) in loops.iter().enumerate() {
96            for &b in &lp.body {
97                let entry = block_loop.entry(b).or_insert(i);
98                if loops[i].body.len() < loops[*entry].body.len() {
99                    *entry = i;
100                }
101            }
102        }
103
104        LoopInfo { loops, block_loop }
105    }
106
107    /// Find all back-edges (tail, header) using iterative DFS on the CFG.
108    ///
109    /// An edge `(n → h)` is classified as a back-edge when `h` dominates `n`.
110    /// This definition is equivalent to natural-loop back-edges **only for
111    /// reducible CFGs**.  In an irreducible CFG, cross-edges inside a
112    /// multi-entry SCC are not dominance back-edges and will be missed.
113    fn find_back_edges(cfg: &Cfg, dom: &DomTree) -> Vec<(BlockId, BlockId)> {
114        let mut back_edges = Vec::new();
115        let mut visited = HashSet::new();
116        let mut stack = vec![cfg.entry()];
117        while let Some(b) = stack.pop() {
118            if !visited.insert(b) {
119                continue;
120            }
121            for &succ in cfg.successors(b) {
122                if dom.dominates(succ, b) {
123                    back_edges.push((b, succ));
124                } else {
125                    stack.push(succ);
126                }
127            }
128        }
129        back_edges
130    }
131
132    /// Collect all blocks in the natural loop for back-edge (tail → header)
133    /// using reverse-CFG BFS starting at `tail`, stopping at `header`.
134    fn collect_loop_body(tail: BlockId, header: BlockId, cfg: &Cfg) -> Vec<BlockId> {
135        let mut body: HashSet<BlockId> = HashSet::new();
136        body.insert(header);
137        let mut queue = VecDeque::new();
138        if body.insert(tail) {
139            queue.push_back(tail);
140        }
141        while let Some(b) = queue.pop_front() {
142            for &pred in cfg.predecessors(b) {
143                if body.insert(pred) {
144                    queue.push_back(pred);
145                }
146            }
147        }
148        let mut v: Vec<BlockId> = body.into_iter().collect();
149        v.sort_unstable_by_key(|b| b.0);
150        v
151    }
152
153    /// All detected loops.
154    pub fn loops(&self) -> &[Loop] {
155        &self.loops
156    }
157
158    /// Index of the innermost loop containing `bid`, or `None`.
159    pub fn loop_of(&self, bid: BlockId) -> Option<usize> {
160        self.block_loop.get(&bid).copied()
161    }
162
163    /// `true` if `bid` is the header of any detected loop.
164    pub fn is_loop_header(&self, bid: BlockId) -> bool {
165        self.loops.iter().any(|l| l.header == bid)
166    }
167
168    /// Nesting depth of `bid` (0 = not in a loop, 1 = outermost loop, …).
169    pub fn depth(&self, bid: BlockId) -> usize {
170        let mut idx = match self.block_loop.get(&bid) {
171            None => return 0,
172            Some(&i) => i,
173        };
174        let mut d = 1;
175        while let Some(p) = self.loops[idx].parent {
176            idx = p;
177            d += 1;
178        }
179        d
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use llvm_ir::{BasicBlock, Context, Function, InstrKind, Instruction, Linkage, ValueRef};
187
188    fn build_func(num_blocks: usize, edges: &[(usize, Vec<usize>)]) -> (Context, Function) {
189        let mut ctx = Context::new();
190        let fn_ty = ctx.mk_fn_type(ctx.void_ty, vec![], false);
191        let mut func = Function::new("test", fn_ty, vec![], Linkage::External);
192        for i in 0..num_blocks {
193            func.add_block(BasicBlock::new(format!("b{}", i)));
194        }
195        let mut has_term = vec![false; num_blocks];
196        for &(src, ref dsts) in edges {
197            has_term[src] = true;
198            let kind = match dsts.as_slice() {
199                [] => InstrKind::Unreachable,
200                [dst] => InstrKind::Br {
201                    dest: BlockId(*dst as u32),
202                },
203                [t, f] => {
204                    let cond = ValueRef::Constant(ctx.const_int(ctx.i1_ty, 0));
205                    InstrKind::CondBr {
206                        cond,
207                        then_dest: BlockId(*t as u32),
208                        else_dest: BlockId(*f as u32),
209                    }
210                }
211                _ => panic!("max 2 successors"),
212            };
213            let iid = func.alloc_instr(Instruction {
214                name: None,
215                ty: ctx.void_ty,
216                kind,
217            });
218            func.blocks[src].set_terminator(iid);
219        }
220        for (i, &needs_term) in has_term.iter().enumerate() {
221            if !needs_term {
222                let iid = func.alloc_instr(Instruction {
223                    name: None,
224                    ty: ctx.void_ty,
225                    kind: InstrKind::Unreachable,
226                });
227                func.blocks[i].set_terminator(iid);
228            }
229        }
230        (ctx, func)
231    }
232
233    #[test]
234    fn no_loops() {
235        let (_ctx, func) = build_func(3, &[(0, vec![1]), (1, vec![2]), (2, vec![])]);
236        let cfg = Cfg::compute(&func);
237        let dom = DomTree::compute(&func, &cfg);
238        let li = LoopInfo::compute(&func, &cfg, &dom);
239        assert!(li.loops().is_empty());
240        assert_eq!(li.depth(BlockId(0)), 0);
241        assert_eq!(li.depth(BlockId(1)), 0);
242        assert_eq!(li.depth(BlockId(2)), 0);
243    }
244
245    #[test]
246    fn simple_loop() {
247        // 0 -> 1 -> 2 -> {1(back), 3}
248        let (_ctx, func) = build_func(
249            4,
250            &[(0, vec![1]), (1, vec![2]), (2, vec![1, 3]), (3, vec![])],
251        );
252        let cfg = Cfg::compute(&func);
253        let dom = DomTree::compute(&func, &cfg);
254        let li = LoopInfo::compute(&func, &cfg, &dom);
255
256        assert_eq!(li.loops().len(), 1);
257        assert_eq!(li.loops()[0].header, BlockId(1));
258        assert!(li.loops()[0].body.contains(&BlockId(1)));
259        assert!(li.loops()[0].body.contains(&BlockId(2)));
260        assert!(!li.loops()[0].body.contains(&BlockId(0)));
261        assert!(!li.loops()[0].body.contains(&BlockId(3)));
262
263        assert!(li.is_loop_header(BlockId(1)));
264        assert!(!li.is_loop_header(BlockId(0)));
265        assert_eq!(li.depth(BlockId(0)), 0);
266        assert_eq!(li.depth(BlockId(1)), 1);
267        assert_eq!(li.depth(BlockId(2)), 1);
268        assert_eq!(li.depth(BlockId(3)), 0);
269    }
270
271    #[test]
272    fn nested_loops() {
273        // Outer loop header=1: body={1,2,3}
274        // Inner loop header=2: body={2,3}
275        // CFG: 0->1, 1->{2,4}, 2->3, 3->{2(inner back),1(outer back)}, 4 exit
276        let (_ctx, func) = build_func(
277            5,
278            &[
279                (0, vec![1]),
280                (1, vec![2, 4]),
281                (2, vec![3]),
282                (3, vec![2, 1]),
283                (4, vec![]),
284            ],
285        );
286        let cfg = Cfg::compute(&func);
287        let dom = DomTree::compute(&func, &cfg);
288        let li = LoopInfo::compute(&func, &cfg, &dom);
289
290        assert_eq!(li.loops().len(), 2);
291        let headers: Vec<BlockId> = li.loops().iter().map(|l| l.header).collect();
292        assert!(headers.contains(&BlockId(1)));
293        assert!(headers.contains(&BlockId(2)));
294
295        // Outer loop contains blocks 1,2,3; inner contains 2,3.
296        assert_eq!(li.depth(BlockId(1)), 1); // only in outer
297        assert_eq!(li.depth(BlockId(2)), 2); // in both
298        assert_eq!(li.depth(BlockId(3)), 2); // in both
299        assert_eq!(li.depth(BlockId(0)), 0);
300        assert_eq!(li.depth(BlockId(4)), 0);
301    }
302
303    #[test]
304    fn irreducible_cfg_not_detected() {
305        // Irreducible CFG: 0 → 1, 0 → 2, 1 → 2, 2 → 1
306        //
307        // Blocks 1 and 2 form a cycle but neither dominates the other, so no
308        // back-edge is detected by the dominance-based algorithm.  This test
309        // documents the known limitation described in issue #9: natural loop
310        // detection silently under-reports loops in irreducible CFGs.
311        let (_ctx, func) = build_func(3, &[(0, vec![1, 2]), (1, vec![2]), (2, vec![1])]);
312        let cfg = Cfg::compute(&func);
313        let dom = DomTree::compute(&func, &cfg);
314        let li = LoopInfo::compute(&func, &cfg, &dom);
315
316        // The cycle 1 ↔ 2 is NOT reported because neither block dominates the
317        // other (irreducible CFG).  Zero loops is the known, documented
318        // behaviour for this input; a correct SCC-based detector would find 1.
319        assert_eq!(li.loops().len(), 0,
320            "dominance-based algorithm does not detect irreducible cycle (known limitation, see issue #9)");
321    }
322}