1use crate::cfg::Cfg;
28use crate::dominators::DomTree;
29use llvm_ir::{BlockId, Function};
30use std::collections::{HashMap, HashSet, VecDeque};
31
32#[derive(Debug)]
34pub struct Loop {
35 pub header: BlockId,
37 pub body: Vec<BlockId>,
39 pub parent: Option<usize>,
41}
42
43pub struct LoopInfo {
45 loops: Vec<Loop>,
47 block_loop: HashMap<BlockId, usize>,
49}
50
51impl LoopInfo {
52 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 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 loops.sort_by(|a, b| b.body.len().cmp(&a.body.len()));
80
81 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 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 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 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 pub fn loops(&self) -> &[Loop] {
155 &self.loops
156 }
157
158 pub fn loop_of(&self, bid: BlockId) -> Option<usize> {
160 self.block_loop.get(&bid).copied()
161 }
162
163 pub fn is_loop_header(&self, bid: BlockId) -> bool {
165 self.loops.iter().any(|l| l.header == bid)
166 }
167
168 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 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 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 assert_eq!(li.depth(BlockId(1)), 1); assert_eq!(li.depth(BlockId(2)), 2); assert_eq!(li.depth(BlockId(3)), 2); 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 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 assert_eq!(li.loops().len(), 0,
320 "dominance-based algorithm does not detect irreducible cycle (known limitation, see issue #9)");
321 }
322}