Add loop scanner to tool-scanner

tests/script-based-pre/tool-scanner/scanner-test.expected

@@ -1,6 +1,6 @@
-2 test_scan_fn_loops.csv
-16 test_scan_functions.csv
+5 test_scan_fn_loops.csv
+19 test_scan_functions.csv
 5 test_scan_input_tys.csv
-14 test_scan_overall.csv
+16 test_scan_overall.csv
 3 test_scan_recursion.csv
 5 test_scan_unsafe_ops.csv

tests/script-based-pre/tool-scanner/test.rs

@@ -36,6 +36,33 @@ pub fn with_iterator(input: &[usize]) -> usize {
         .unwrap_or_else(|| input.iter().fold(0, |acc, i| acc + 1))
 }
 
+pub fn with_for_loop(input: &[usize]) -> usize {
+    let mut res = 0;
+    for n in input {
+        res += 1;
+    }
+    res
+}
+
+pub fn with_while_loop(input: &[usize]) -> usize {
+    let mut res = 0;
+    while res < input.len() {
+        res += 1;
+    }
+    return res;
+}
+
+pub fn with_loop_loop(input: &[usize]) -> usize {
+    let mut res = 0;
+    loop {
+        if res == input.len() {
+            break;
+        }
+        res += 1;
+    }
+    res
+}
+
 static mut COUNTER: Option<usize> = Some(0);
 static OK: bool = true;
 

tools/scanner/src/analysis.rs

@@ -9,7 +9,7 @@ use serde::{ser::SerializeStruct, Serialize, Serializer};
 use stable_mir::mir::mono::Instance;
 use stable_mir::mir::visit::{Location, PlaceContext, PlaceRef};
 use stable_mir::mir::{
-    Body, MirVisitor, Mutability, ProjectionElem, Safety, Terminator, TerminatorKind,
+    BasicBlock, Body, MirVisitor, Mutability, ProjectionElem, Safety, Terminator, TerminatorKind,
 };
 use stable_mir::ty::{AdtDef, AdtKind, FnDef, GenericArgs, MirConst, RigidTy, Ty, TyKind};
 use stable_mir::visitor::{Visitable, Visitor};
@@ -159,6 +159,7 @@ impl OverallStats {
     pub fn loops(&mut self, filename: PathBuf) {
         let all_items = stable_mir::all_local_items();
         let (has_loops, no_loops) = all_items
+            .clone()
             .into_iter()
             .filter_map(|item| {
                 let kind = item.ty().kind();
@@ -168,9 +169,37 @@ impl OverallStats {
                 Some(FnLoops::new(item.name()).collect(&item.body()))
             })
             .partition::<Vec<_>, _>(|props| props.has_loops());
-        self.counters
-            .extend_from_slice(&[("has_loops", has_loops.len()), ("no_loops", no_loops.len())]);
-        dump_csv(filename, &has_loops);
+
+        let (has_iterators, no_iterators) = all_items
+            .clone()
+            .into_iter()
+            .filter_map(|item| {
+                let kind = item.ty().kind();
+                if !kind.is_fn() {
+                    return None;
+                };
+                Some(FnLoops::new(item.name()).collect(&item.body()))
+            })
+            .partition::<Vec<_>, _>(|props| props.has_iterators());
+
+        let (has_either, _) = all_items
+            .into_iter()
+            .filter_map(|item| {
+                let kind = item.ty().kind();
+                if !kind.is_fn() {
+                    return None;
+                };
+                Some(FnLoops::new(item.name()).collect(&item.body()))
+            })
+            .partition::<Vec<_>, _>(|props| props.has_iterators() || props.has_loops());
+
+        self.counters.extend_from_slice(&[
+            ("has_loops", has_loops.len()),
+            ("no_loops", no_loops.len()),
+            ("has_iterators", has_iterators.len()),
+            ("no_iterators", no_iterators.len()),
+        ]);
+        dump_csv(filename, &has_either);
     }
 
     /// Create a callgraph for this crate and try to find recursive calls.
@@ -436,21 +465,25 @@ impl<'a> MirVisitor for BodyVisitor<'a> {
 fn_props! {
     struct FnLoops {
         iterators,
-        nested_loops,
-        /// TODO: Collect loops.
         loops,
+        // TODO: Collect nested loops.
+        nested_loops,
     }
 }
 
 impl FnLoops {
     pub fn collect(self, body: &Body) -> FnLoops {
-        let mut visitor = IteratorVisitor { props: self, body };
+        let mut visitor = IteratorVisitor { props: self, body, visited_blocks: HashSet::new() };
         visitor.visit_body(body);
         visitor.props
     }
 
     pub fn has_loops(&self) -> bool {
-        (self.iterators + self.loops + self.nested_loops) > 0
+        (self.loops + self.nested_loops) > 0
+    }
+
+    pub fn has_iterators(&self) -> bool {
+        (self.iterators) > 0
     }
 }
 
@@ -461,50 +494,94 @@ impl FnLoops {
 struct IteratorVisitor<'a> {
     props: FnLoops,
     body: &'a Body,
+    visited_blocks: HashSet<usize>,
 }
 
 impl<'a> MirVisitor for IteratorVisitor<'a> {
+    fn visit_basic_block(&mut self, bb: &BasicBlock) {
+        self.visited_blocks.insert(self.body.blocks.iter().position(|b| *b == *bb).unwrap());
+        self.super_basic_block(bb);
+    }
+
     fn visit_terminator(&mut self, term: &Terminator, location: Location) {
-        if let TerminatorKind::Call { func, .. } = &term.kind {
-            let kind = func.ty(self.body.locals()).unwrap().kind();
-            if let TyKind::RigidTy(RigidTy::FnDef(def, _)) = kind {
-                let fullname = def.name();
-                let names = fullname.split("::").collect::<Vec<_>>();
-                if let [.., s_last, last] = names.as_slice() {
-                    if *s_last == "Iterator"
-                        && [
-                            "for_each",
-                            "collect",
-                            "advance_by",
-                            "all",
-                            "any",
-                            "partition",
-                            "partition_in_place",
-                            "fold",
-                            "try_fold",
-                            "spec_fold",
-                            "spec_try_fold",
-                            "try_for_each",
-                            "for_each",
-                            "try_reduce",
-                            "reduce",
-                            "find",
-                            "find_map",
-                            "try_find",
-                            "position",
-                            "rposition",
-                            "nth",
-                            "count",
-                            "last",
-                            "find",
-                        ]
-                        .contains(last)
-                    {
-                        self.props.iterators += 1;
+        match &term.kind {
+            TerminatorKind::Goto { target } => {
+                // A goto is identified as a loop latch if its target has been visited.
+                if self.visited_blocks.contains(target) {
+                    self.props.loops += 1;
+                }
+            }
+            TerminatorKind::SwitchInt { targets, .. } => {
+                let successors = targets.all_targets();
+                // Check if any of the targets of SwitchInt has been visited.
+                if self.visited_blocks.contains(&successors[0])
+                    || self.visited_blocks.contains(&successors[1])
+                {
+                    self.props.loops += 1;
+                }
+            }
+            TerminatorKind::Drop { target, .. } => {
+                if self.visited_blocks.contains(target) {
+                    self.props.loops += 1;
+                }
+            }
+            TerminatorKind::Call { func, .. } => {
+                let kind = func.ty(self.body.locals()).unwrap().kind();
+                // Check if the target is a visited block.
+
+                // Check if the call is an iterator function that contains loops.
+                if let TyKind::RigidTy(RigidTy::FnDef(def, _)) = kind {
+                    let fullname = def.name();
+                    let names = fullname.split("::").collect::<Vec<_>>();
+                    if let [.., s_last, last] = names.as_slice() {
+                        if *s_last == "Iterator"
+                            && [
+                                "for_each",
+                                "collect",
+                                "advance_by",
+                                "all",
+                                "any",
+                                "partition",
+                                "partition_in_place",
+                                "fold",
+                                "try_fold",
+                                "spec_fold",
+                                "spec_try_fold",
+                                "try_for_each",
+                                "for_each",
+                                "try_reduce",
+                                "reduce",
+                                "find",
+                                "find_map",
+                                "try_find",
+                                "position",
+                                "rposition",
+                                "nth",
+                                "count",
+                                "last",
+                                "find",
+                            ]
+                            .contains(last)
+                        {
+                            self.props.iterators += 1;
+                        }
                     }
                 }
             }
+            TerminatorKind::Assert { target, .. } => {
+                if self.visited_blocks.contains(target) {
+                    self.props.loops += 1;
+                }
+            }
+            TerminatorKind::InlineAsm { destination: Some(target), .. } => {
+                if self.visited_blocks.contains(target) {
+                    self.props.loops += 1;
+                }
+            }
+            // No targets in other terminators.
+            _ => {}
         }
+
         self.super_terminator(term, location)
     }
 }