1use crate::cfg::analysis::find_entry;
4use crate::cfg::Cfg;
5use petgraph::algo::dominators::simple_fast;
6use petgraph::graph::NodeIndex;
7use petgraph::visit::EdgeRef;
8use std::collections::{HashSet, VecDeque};
9
10#[derive(Debug, Clone)]
15pub struct NaturalLoop {
16 pub header: NodeIndex,
18 pub back_edge: (NodeIndex, NodeIndex),
20 pub body: HashSet<NodeIndex>,
22}
23
24impl NaturalLoop {
25 pub fn contains(&self, node: NodeIndex) -> bool {
27 self.body.contains(&node)
28 }
29
30 pub fn size(&self) -> usize {
32 self.body.len()
33 }
34
35 pub fn nesting_level(&self, all_loops: &[NaturalLoop]) -> usize {
39 let mut level = 0;
40 for other in all_loops {
41 if other.header != self.header && other.body.contains(&self.header) {
42 level = level.max(other.nesting_level(all_loops) + 1);
43 }
44 }
45 level
46 }
47}
48
49pub fn detect_natural_loops(cfg: &Cfg) -> Vec<NaturalLoop> {
70 let entry = match find_entry(cfg) {
71 Some(e) => e,
72 None => return vec![],
73 };
74
75 let dominators = simple_fast(cfg, entry);
77
78 let mut loops = Vec::new();
79
80 for edge in cfg.edge_references() {
82 let tail = edge.source();
83 let header = edge.target();
84
85 if let Some(mut tail_dominators) = dominators.dominators(tail) {
88 if tail_dominators.any(|d| d == header) {
89 let body = compute_loop_body(cfg, header, tail);
90 loops.push(NaturalLoop {
91 header,
92 back_edge: (tail, header),
93 body,
94 });
95 }
96 }
97 }
98
99 loops
100}
101
102pub fn apply_loop_nesting_depths(cfg: &mut Cfg) {
119 let loops = detect_natural_loops(cfg);
120
121 for node in cfg.node_indices() {
123 if let Some(block) = cfg.node_weight_mut(node) {
124 block.coord_y = 0;
125 }
126 }
127
128 for loop_ in &loops {
130 for node in &loop_.body {
131 if let Some(block) = cfg.node_weight_mut(*node) {
132 block.coord_y += 1;
133 }
134 }
135 }
136}
137
138pub fn detect_natural_loops_with_depths(cfg: &mut Cfg) -> Vec<NaturalLoop> {
160 let loops = detect_natural_loops(cfg);
161 apply_loop_nesting_depths_from_loops(cfg, &loops);
162 loops
163}
164
165fn apply_loop_nesting_depths_from_loops(cfg: &mut Cfg, loops: &[NaturalLoop]) {
167 for node in cfg.node_indices() {
169 if let Some(block) = cfg.node_weight_mut(node) {
170 block.coord_y = 0;
171 }
172 }
173
174 for loop_ in loops {
176 for node in &loop_.body {
177 if let Some(block) = cfg.node_weight_mut(*node) {
178 block.coord_y += 1;
179 }
180 }
181 }
182}
183
184fn compute_loop_body(cfg: &Cfg, header: NodeIndex, tail: NodeIndex) -> HashSet<NodeIndex> {
193 let mut body = HashSet::new();
194 let mut worklist = VecDeque::new();
195
196 worklist.push_back(tail);
197
198 while let Some(node) = worklist.pop_front() {
199 if node == header {
200 continue;
201 }
202
203 if body.contains(&node) {
204 continue;
205 }
206
207 body.insert(node);
208
209 for pred in cfg.neighbors_directed(node, petgraph::Direction::Incoming) {
211 if pred != header && !body.contains(&pred) {
212 worklist.push_back(pred);
213 }
214 }
215 }
216
217 body.insert(header); body
219}
220
221pub fn find_loop_headers(cfg: &Cfg) -> HashSet<NodeIndex> {
235 detect_natural_loops(cfg)
236 .into_iter()
237 .map(|loop_| loop_.header)
238 .collect()
239}
240
241pub fn is_loop_header(cfg: &Cfg, node: NodeIndex) -> bool {
245 find_loop_headers(cfg).contains(&node)
246}
247
248pub fn loops_containing(cfg: &Cfg, node: NodeIndex) -> Vec<NaturalLoop> {
252 detect_natural_loops(cfg)
253 .into_iter()
254 .filter(|loop_| loop_.body.contains(&node))
255 .collect()
256}
257
258pub fn find_nested_loops(cfg: &Cfg) -> Vec<(NaturalLoop, NaturalLoop)> {
262 let loops = detect_natural_loops(cfg);
263 let mut nested = Vec::new();
264
265 for (i, outer) in loops.iter().enumerate() {
266 for inner in loops.iter().skip(i + 1) {
267 if outer.body.contains(&inner.header) {
269 nested.push((outer.clone(), inner.clone()));
270 }
271 }
272 }
273
274 nested
275}
276
277#[cfg(test)]
278mod tests {
279 use super::*;
280 use crate::cfg::{BasicBlock, BlockKind, EdgeType, Terminator};
281 use petgraph::graph::DiGraph;
282
283 fn create_simple_loop_cfg() -> Cfg {
286 let mut g = DiGraph::new();
287
288 let b0 = g.add_node(BasicBlock {
290 id: 0,
291 db_id: None,
292 kind: BlockKind::Entry,
293 statements: vec![],
294 terminator: Terminator::Goto { target: 1 },
295 source_location: None,
296 coord_x: 0,
297 coord_y: 0,
298 coord_z: 0,
299 });
300
301 let b1 = g.add_node(BasicBlock {
303 id: 1,
304 db_id: None,
305 kind: BlockKind::Normal,
306 statements: vec![],
307 terminator: Terminator::SwitchInt {
308 targets: vec![2],
309 otherwise: 3,
310 },
311 source_location: None,
312 coord_x: 0,
313 coord_y: 0,
314 coord_z: 0,
315 });
316
317 let b2 = g.add_node(BasicBlock {
319 id: 2,
320 db_id: None,
321 kind: BlockKind::Normal,
322 statements: vec!["loop body".to_string()],
323 terminator: Terminator::Goto { target: 1 },
324 source_location: None,
325 coord_x: 0,
326 coord_y: 0,
327 coord_z: 0,
328 });
329
330 let b3 = g.add_node(BasicBlock {
332 id: 3,
333 db_id: None,
334 kind: BlockKind::Exit,
335 statements: vec![],
336 terminator: Terminator::Return,
337 source_location: None,
338 coord_x: 0,
339 coord_y: 0,
340 coord_z: 0,
341 });
342
343 g.add_edge(b0, b1, EdgeType::Fallthrough);
344 g.add_edge(b1, b2, EdgeType::TrueBranch);
345 g.add_edge(b1, b3, EdgeType::FalseBranch);
346 g.add_edge(b2, b1, EdgeType::LoopBack);
347
348 g
349 }
350
351 #[test]
352 fn test_detect_simple_loop() {
353 let cfg = create_simple_loop_cfg();
354 let loops = detect_natural_loops(&cfg);
355
356 assert_eq!(loops.len(), 1);
357
358 let loop_ = &loops[0];
359 assert_eq!(loop_.header.index(), 1); assert_eq!(loop_.back_edge.0.index(), 2); assert_eq!(loop_.back_edge.1.index(), 1); assert!(loop_.contains(NodeIndex::new(1)));
363 assert!(loop_.contains(NodeIndex::new(2)));
364 assert!(!loop_.contains(NodeIndex::new(0))); assert!(!loop_.contains(NodeIndex::new(3))); }
367
368 #[test]
369 fn test_find_loop_headers() {
370 let cfg = create_simple_loop_cfg();
371 let headers = find_loop_headers(&cfg);
372
373 assert_eq!(headers.len(), 1);
374 assert!(headers.contains(&NodeIndex::new(1)));
375 }
376
377 #[test]
378 fn test_is_loop_header() {
379 let cfg = create_simple_loop_cfg();
380
381 assert!(is_loop_header(&cfg, NodeIndex::new(1)));
382 assert!(!is_loop_header(&cfg, NodeIndex::new(0)));
383 assert!(!is_loop_header(&cfg, NodeIndex::new(2)));
384 }
385
386 #[test]
387 fn test_no_loops() {
388 let mut g = DiGraph::new();
389
390 let b0 = g.add_node(BasicBlock {
392 id: 0,
393 db_id: None,
394 kind: BlockKind::Entry,
395 statements: vec![],
396 terminator: Terminator::Goto { target: 1 },
397 source_location: None,
398 coord_x: 0,
399 coord_y: 0,
400 coord_z: 0,
401 });
402
403 let b1 = g.add_node(BasicBlock {
404 id: 1,
405 db_id: None,
406 kind: BlockKind::Normal,
407 statements: vec![],
408 terminator: Terminator::Goto { target: 2 },
409 source_location: None,
410 coord_x: 0,
411 coord_y: 0,
412 coord_z: 0,
413 });
414
415 let b2 = g.add_node(BasicBlock {
416 id: 2,
417 db_id: None,
418 kind: BlockKind::Exit,
419 statements: vec![],
420 terminator: Terminator::Return,
421 source_location: None,
422 coord_x: 0,
423 coord_y: 0,
424 coord_z: 0,
425 });
426
427 g.add_edge(b0, b1, EdgeType::Fallthrough);
428 g.add_edge(b1, b2, EdgeType::Fallthrough);
429
430 let loops = detect_natural_loops(&g);
431 assert!(loops.is_empty());
432 }
433
434 #[test]
435 fn test_nested_loops() {
436 let mut g = DiGraph::new();
437
438 let b0 = g.add_node(BasicBlock {
446 id: 0,
447 db_id: None,
448 kind: BlockKind::Entry,
449 statements: vec![],
450 terminator: Terminator::Goto { target: 1 },
451 source_location: None,
452 coord_x: 0,
453 coord_y: 0,
454 coord_z: 0,
455 });
456
457 let b1 = g.add_node(BasicBlock {
458 id: 1,
459 db_id: None,
460 kind: BlockKind::Normal,
461 statements: vec![],
462 terminator: Terminator::SwitchInt {
463 targets: vec![2],
464 otherwise: 4,
465 },
466 source_location: None,
467 coord_x: 0,
468 coord_y: 0,
469 coord_z: 0,
470 });
471
472 let b2 = g.add_node(BasicBlock {
473 id: 2,
474 db_id: None,
475 kind: BlockKind::Normal,
476 statements: vec![],
477 terminator: Terminator::SwitchInt {
478 targets: vec![3],
479 otherwise: 1,
480 },
481 source_location: None,
482 coord_x: 0,
483 coord_y: 0,
484 coord_z: 0,
485 });
486
487 let b3 = g.add_node(BasicBlock {
488 id: 3,
489 db_id: None,
490 kind: BlockKind::Normal,
491 statements: vec![],
492 terminator: Terminator::Goto { target: 2 },
493 source_location: None,
494 coord_x: 0,
495 coord_y: 0,
496 coord_z: 0,
497 });
498
499 let b4 = g.add_node(BasicBlock {
500 id: 4,
501 db_id: None,
502 kind: BlockKind::Exit,
503 statements: vec![],
504 terminator: Terminator::Return,
505 source_location: None,
506 coord_x: 0,
507 coord_y: 0,
508 coord_z: 0,
509 });
510
511 g.add_edge(b0, b1, EdgeType::Fallthrough);
512 g.add_edge(b1, b2, EdgeType::TrueBranch);
513 g.add_edge(b1, b4, EdgeType::FalseBranch);
514 g.add_edge(b2, b3, EdgeType::TrueBranch);
515 g.add_edge(b2, b1, EdgeType::LoopBack); g.add_edge(b3, b2, EdgeType::LoopBack); let loops = detect_natural_loops(&g);
519 assert_eq!(loops.len(), 2); let nested = find_nested_loops(&g);
522 assert_eq!(nested.len(), 1); }
524
525 #[test]
526 fn test_empty_cfg() {
527 let cfg: Cfg = DiGraph::new();
528 assert!(detect_natural_loops(&cfg).is_empty());
529 assert!(find_loop_headers(&cfg).is_empty());
530 }
531
532 #[test]
533 fn test_loops_containing() {
534 let cfg = create_simple_loop_cfg();
535
536 let loops_2 = loops_containing(&cfg, NodeIndex::new(2));
538 assert_eq!(loops_2.len(), 1);
539
540 let loops_0 = loops_containing(&cfg, NodeIndex::new(0));
542 assert_eq!(loops_0.len(), 0);
543 }
544
545 #[test]
546 fn test_loop_size() {
547 let cfg = create_simple_loop_cfg();
548 let loops = detect_natural_loops(&cfg);
549
550 assert_eq!(loops.len(), 1);
551 assert_eq!(loops[0].size(), 2); }
553
554 #[test]
555 fn test_nesting_level() {
556 let cfg = create_simple_loop_cfg();
557 let loops = detect_natural_loops(&cfg);
558
559 assert_eq!(loops.len(), 1);
560 assert_eq!(loops[0].nesting_level(&loops), 0);
562 }
563
564 #[test]
565 fn test_nesting_level_nested() {
566 let mut g = DiGraph::new();
567
568 let b0 = g.add_node(BasicBlock {
570 id: 0,
571 db_id: None,
572 kind: BlockKind::Entry,
573 statements: vec![],
574 terminator: Terminator::Goto { target: 1 },
575 source_location: None,
576 coord_x: 0,
577 coord_y: 0,
578 coord_z: 0,
579 });
580
581 let b1 = g.add_node(BasicBlock {
582 id: 1,
583 db_id: None,
584 kind: BlockKind::Normal,
585 statements: vec![],
586 terminator: Terminator::SwitchInt {
587 targets: vec![2],
588 otherwise: 4,
589 },
590 source_location: None,
591 coord_x: 0,
592 coord_y: 0,
593 coord_z: 0,
594 });
595
596 let b2 = g.add_node(BasicBlock {
597 id: 2,
598 db_id: None,
599 kind: BlockKind::Normal,
600 statements: vec![],
601 terminator: Terminator::SwitchInt {
602 targets: vec![3],
603 otherwise: 1,
604 },
605 source_location: None,
606 coord_x: 0,
607 coord_y: 0,
608 coord_z: 0,
609 });
610
611 let b3 = g.add_node(BasicBlock {
612 id: 3,
613 db_id: None,
614 kind: BlockKind::Normal,
615 statements: vec![],
616 terminator: Terminator::Goto { target: 2 },
617 source_location: None,
618 coord_x: 0,
619 coord_y: 0,
620 coord_z: 0,
621 });
622
623 let b4 = g.add_node(BasicBlock {
624 id: 4,
625 db_id: None,
626 kind: BlockKind::Exit,
627 statements: vec![],
628 terminator: Terminator::Return,
629 source_location: None,
630 coord_x: 0,
631 coord_y: 0,
632 coord_z: 0,
633 });
634
635 g.add_edge(b0, b1, EdgeType::Fallthrough);
636 g.add_edge(b1, b2, EdgeType::TrueBranch);
637 g.add_edge(b1, b4, EdgeType::FalseBranch);
638 g.add_edge(b2, b3, EdgeType::TrueBranch);
639 g.add_edge(b2, b1, EdgeType::LoopBack);
640 g.add_edge(b3, b2, EdgeType::LoopBack);
641
642 let loops = detect_natural_loops(&g);
643 assert_eq!(loops.len(), 2);
644
645 let outer_loop = loops.iter().find(|l| l.header.index() == 1).unwrap();
647 let inner_loop = loops.iter().find(|l| l.header.index() == 2).unwrap();
648
649 assert_eq!(outer_loop.nesting_level(&loops), 0);
651 assert_eq!(inner_loop.nesting_level(&loops), 1);
653 }
654
655 #[test]
656 fn test_apply_loop_nesting_depths_simple_loop() {
657 let mut cfg = create_simple_loop_cfg();
659
660 apply_loop_nesting_depths(&mut cfg);
662
663 assert_eq!(
666 cfg[NodeIndex::new(0)].coord_y,
667 0,
668 "Entry should not be in loop"
669 );
670 assert_eq!(
671 cfg[NodeIndex::new(1)].coord_y,
672 1,
673 "Loop header should have depth 1"
674 );
675 assert_eq!(
676 cfg[NodeIndex::new(2)].coord_y,
677 1,
678 "Loop body should have depth 1"
679 );
680 assert_eq!(
681 cfg[NodeIndex::new(3)].coord_y,
682 0,
683 "Exit should not be in loop"
684 );
685 }
686
687 #[test]
688 fn test_apply_loop_nesting_depths_nested_loops() {
689 let mut g = DiGraph::new();
691
692 let b0 = g.add_node(BasicBlock {
693 id: 0,
694 db_id: None,
695 kind: BlockKind::Entry,
696 statements: vec![],
697 terminator: Terminator::Goto { target: 1 },
698 source_location: None,
699 coord_x: 0,
700 coord_y: 0,
701 coord_z: 0,
702 });
703
704 let b1 = g.add_node(BasicBlock {
705 id: 1,
706 db_id: None,
707 kind: BlockKind::Normal,
708 statements: vec![],
709 terminator: Terminator::SwitchInt {
710 targets: vec![2],
711 otherwise: 4,
712 },
713 source_location: None,
714 coord_x: 0,
715 coord_y: 0,
716 coord_z: 0,
717 });
718
719 let b2 = g.add_node(BasicBlock {
720 id: 2,
721 db_id: None,
722 kind: BlockKind::Normal,
723 statements: vec![],
724 terminator: Terminator::SwitchInt {
725 targets: vec![3],
726 otherwise: 1,
727 },
728 source_location: None,
729 coord_x: 0,
730 coord_y: 0,
731 coord_z: 0,
732 });
733
734 let b3 = g.add_node(BasicBlock {
735 id: 3,
736 db_id: None,
737 kind: BlockKind::Normal,
738 statements: vec![],
739 terminator: Terminator::Goto { target: 2 },
740 source_location: None,
741 coord_x: 0,
742 coord_y: 0,
743 coord_z: 0,
744 });
745
746 let b4 = g.add_node(BasicBlock {
747 id: 4,
748 db_id: None,
749 kind: BlockKind::Exit,
750 statements: vec![],
751 terminator: Terminator::Return,
752 source_location: None,
753 coord_x: 0,
754 coord_y: 0,
755 coord_z: 0,
756 });
757
758 g.add_edge(b0, b1, EdgeType::Fallthrough);
759 g.add_edge(b1, b2, EdgeType::TrueBranch);
760 g.add_edge(b1, b4, EdgeType::FalseBranch);
761 g.add_edge(b2, b3, EdgeType::TrueBranch);
762 g.add_edge(b2, b1, EdgeType::Fallthrough); g.add_edge(b3, b2, EdgeType::Fallthrough); apply_loop_nesting_depths(&mut g);
767
768 assert_eq!(g[b0].coord_y, 0, "Entry should not be in a loop");
770 assert_eq!(g[b1].coord_y, 1, "Outer loop header should have depth 1");
771 assert_eq!(g[b2].coord_y, 2, "Inner loop header should have depth 2");
772 assert_eq!(g[b3].coord_y, 2, "Inner loop body should have depth 2");
773 assert_eq!(g[b4].coord_y, 0, "Exit should not be in a loop");
774 }
775
776 #[test]
777 fn test_detect_natural_loops_with_depths_simple_loop() {
778 let mut cfg = create_simple_loop_cfg();
780
781 let loops = detect_natural_loops_with_depths(&mut cfg);
783
784 assert_eq!(loops.len(), 1, "Should detect exactly one loop");
786 assert_eq!(
787 cfg[NodeIndex::new(0)].coord_y,
788 0,
789 "Entry should not be in loop"
790 );
791 assert_eq!(
792 cfg[NodeIndex::new(1)].coord_y,
793 1,
794 "Loop header should have depth 1"
795 );
796 assert_eq!(
797 cfg[NodeIndex::new(2)].coord_y,
798 1,
799 "Loop body should have depth 1"
800 );
801 }
802
803 #[test]
804 fn test_apply_loop_nesting_depths_no_loops() {
805 let mut g = DiGraph::new();
807
808 let b0 = g.add_node(BasicBlock {
809 id: 0,
810 db_id: None,
811 kind: BlockKind::Entry,
812 statements: vec![],
813 terminator: Terminator::Goto { target: 1 },
814 source_location: None,
815 coord_x: 0,
816 coord_y: 0,
817 coord_z: 0,
818 });
819
820 let b1 = g.add_node(BasicBlock {
821 id: 1,
822 db_id: None,
823 kind: BlockKind::Exit,
824 statements: vec![],
825 terminator: Terminator::Return,
826 source_location: None,
827 coord_x: 0,
828 coord_y: 0,
829 coord_z: 0,
830 });
831
832 g.add_edge(b0, b1, EdgeType::Fallthrough);
833
834 apply_loop_nesting_depths(&mut g);
836
837 assert_eq!(g[b0].coord_y, 0, "Entry should not be in a loop");
839 assert_eq!(g[b1].coord_y, 0, "Exit should not be in a loop");
840 }
841
842 #[test]
843 fn test_loop_nesting_matches_nesting_level_method() {
844 let mut g = DiGraph::new();
846
847 let b0 = g.add_node(BasicBlock {
848 id: 0,
849 db_id: None,
850 kind: BlockKind::Entry,
851 statements: vec![],
852 terminator: Terminator::Goto { target: 1 },
853 source_location: None,
854 coord_x: 0,
855 coord_y: 0,
856 coord_z: 0,
857 });
858
859 let b1 = g.add_node(BasicBlock {
860 id: 1,
861 db_id: None,
862 kind: BlockKind::Normal,
863 statements: vec![],
864 terminator: Terminator::SwitchInt {
865 targets: vec![2],
866 otherwise: 4,
867 },
868 source_location: None,
869 coord_x: 0,
870 coord_y: 0,
871 coord_z: 0,
872 });
873
874 let b2 = g.add_node(BasicBlock {
875 id: 2,
876 db_id: None,
877 kind: BlockKind::Normal,
878 statements: vec![],
879 terminator: Terminator::SwitchInt {
880 targets: vec![3],
881 otherwise: 1,
882 },
883 source_location: None,
884 coord_x: 0,
885 coord_y: 0,
886 coord_z: 0,
887 });
888
889 let b3 = g.add_node(BasicBlock {
890 id: 3,
891 db_id: None,
892 kind: BlockKind::Normal,
893 statements: vec![],
894 terminator: Terminator::Goto { target: 2 },
895 source_location: None,
896 coord_x: 0,
897 coord_y: 0,
898 coord_z: 0,
899 });
900
901 let b4 = g.add_node(BasicBlock {
902 id: 4,
903 db_id: None,
904 kind: BlockKind::Exit,
905 statements: vec![],
906 terminator: Terminator::Return,
907 source_location: None,
908 coord_x: 0,
909 coord_y: 0,
910 coord_z: 0,
911 });
912
913 g.add_edge(b0, b1, EdgeType::Fallthrough);
914 g.add_edge(b1, b2, EdgeType::TrueBranch);
915 g.add_edge(b1, b4, EdgeType::FalseBranch);
916 g.add_edge(b2, b3, EdgeType::TrueBranch);
917 g.add_edge(b2, b1, EdgeType::Fallthrough);
918 g.add_edge(b3, b2, EdgeType::Fallthrough);
919
920 let loops = detect_natural_loops_with_depths(&mut g);
922
923 for loop_ in &loops {
925 let expected_coord_y = (loop_.nesting_level(&loops) + 1) as i64;
926 let actual_coord_y = g[loop_.header].coord_y;
927 assert_eq!(
928 actual_coord_y, expected_coord_y,
929 "Loop header {:?} coord_y should match nesting_level + 1",
930 loop_.header
931 );
932 }
933 }
934}