Add loop scanner to tool-scanner (#3443)

qinheping · web-flow · commit d2141e4c5f59 · 2024-08-20T21:18:55.000Z
This extend #3120 with loop scanner that counts the number of loops in each function and the number of functions that contain loops. By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 and MIT licenses.
diff --git a/Cargo.lock b/Cargo.lock
@@ -9,6 +9,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011"
 dependencies = [
  "cfg-if",
+ "getrandom",
  "once_cell",
  "version_check",
  "zerocopy",
@@ -349,6 +350,12 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9fc0510504f03c51ada170672ac806f1f105a88aa97a5281117e1ddc3368e51a"
 
+[[package]]
+name = "fixedbitset"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
+
 [[package]]
 name = "getopts"
 version = "0.2.21"
@@ -375,6 +382,17 @@ version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
 
+[[package]]
+name = "graph-cycles"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a6ad932c6dd3cfaf16b66754a42f87bbeefd591530c4b6a8334270a7df3e853"
+dependencies = [
+ "ahash",
+ "petgraph",
+ "thiserror",
+]
+
 [[package]]
 name = "hashbrown"
 version = "0.14.5"
@@ -723,6 +741,16 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8835116a5c179084a830efb3adc117ab007512b535bc1a21c991d3b32a6b44dd"
 
+[[package]]
+name = "petgraph"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db"
+dependencies = [
+ "fixedbitset",
+ "indexmap",
+]
+
 [[package]]
 name = "pin-project-lite"
 version = "0.2.14"
@@ -928,6 +956,8 @@ name = "scanner"
 version = "0.0.0"
 dependencies = [
  "csv",
+ "graph-cycles",
+ "petgraph",
  "serde",
  "strum",
  "strum_macros",
diff --git a/tests/script-based-pre/tool-scanner/scanner-test.expected b/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
diff --git a/tests/script-based-pre/tool-scanner/test.rs b/tests/script-based-pre/tool-scanner/test.rs
@@ -33,7 +33,34 @@ pub fn with_iterator(input: &[usize]) -> usize {
         .iter()
         .copied()
         .find(|e| *e == 0)
-        .unwrap_or_else(|| input.iter().fold(0, |acc, i| acc + 1))
+        .unwrap_or_else(|| input.iter().fold(0, |acc, _| acc + 1))
+}
+
+pub fn with_for_loop(input: &[usize]) -> usize {
+    let mut res = 0;
+    for _ 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);
diff --git a/tools/scanner/Cargo.toml b/tools/scanner/Cargo.toml
@@ -15,6 +15,8 @@ csv = "1.3"
 serde = {version = "1", features = ["derive"]}
 strum = "0.26"
 strum_macros = "0.26"
+petgraph = "0.6.5"
+graph-cycles = "0.1.0"
 
 [package.metadata.rust-analyzer]
 # This crate uses rustc crates.
diff --git a/tools/scanner/src/analysis.rs b/tools/scanner/src/analysis.rs
@@ -5,11 +5,13 @@
 
 use crate::info;
 use csv::WriterBuilder;
+use graph_cycles::Cycles;
+use petgraph::graph::Graph;
 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 +161,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 +171,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 +467,26 @@ 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, graph: Vec::new(), current_bbidx: 0 };
         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,12 +497,36 @@ impl FnLoops {
 struct IteratorVisitor<'a> {
     props: FnLoops,
     body: &'a Body,
+    graph: Vec<(u32, u32)>,
+    current_bbidx: u32,
 }
 
 impl<'a> MirVisitor for IteratorVisitor<'a> {
+    fn visit_body(&mut self, body: &Body) {
+        // First visit the body to build the control flow graph
+        self.super_body(body);
+        // Build the petgraph from the adj vec
+        let g = Graph::<(), ()>::from_edges(self.graph.clone());
+        self.props.loops += g.cycles().len();
+    }
+
+    fn visit_basic_block(&mut self, bb: &BasicBlock) {
+        self.current_bbidx = self.body.blocks.iter().position(|b| *b == *bb).unwrap() as u32;
+        self.super_basic_block(bb);
+    }
+
     fn visit_terminator(&mut self, term: &Terminator, location: Location) {
+        // Add edges between basic block into the adj table
+        let successors = term.kind.successors();
+        for target in successors {
+            self.graph.push((self.current_bbidx, target as u32));
+        }
+
         if let TerminatorKind::Call { func, .. } = &term.kind {
             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<_>>();
@@ -505,6 +565,7 @@ impl<'a> MirVisitor for IteratorVisitor<'a> {
                 }
             }
         }
+
         self.super_terminator(term, location)
     }
 }
