Avoid corner-cases by grouping instrumentation into basic blocks and using backward iteration

kani-compiler/src/kani_middle/transform/body.rs

@@ -310,16 +310,16 @@ impl MutableBody {
                 // Update the source to point at the terminator.
                 *source = SourceInstruction::Terminator { bb: orig_bb };
             }
-            // Make the terminator at `source` point at the new block,
-            // the terminator of which is a simple Goto instruction.
+            // Make the terminator at `source` point at the new block, the terminator of which is
+            // provided by the caller.
             SourceInstruction::Terminator { bb } => {
                 let current_term = &mut self.blocks.get_mut(*bb).unwrap().terminator;
                 let target_bb = get_mut_target_ref(current_term);
                 let new_target_bb = get_mut_target_ref(&mut new_term);
-                // Set the new terminator to point where previous terminator pointed.
-                *new_target_bb = *target_bb;
-                // Point the current terminator to the new terminator's basic block.
-                *target_bb = new_bb_idx;
+                // Swap the targets of the newly inserted terminator and the original one. This is
+                // an easy way to make the original terminator point to the new basic block with the
+                // new terminator.
+                std::mem::swap(new_target_bb, target_bb);
                 // Update the source to point at the terminator.
                 *bb = new_bb_idx;
                 self.blocks.push(BasicBlock { statements: vec![], terminator: new_term });
@@ -484,7 +484,7 @@ impl CheckType {
 }
 
 /// We store the index of an instruction to avoid borrow checker issues and unnecessary copies.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum SourceInstruction {
     Statement { idx: usize, bb: BasicBlockIdx },
     Terminator { bb: BasicBlockIdx },

kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/instrumentation_visitor.rs

@@ -18,20 +18,17 @@ use rustc_middle::ty::TyCtxt;
 use stable_mir::mir::{
     mono::Instance,
     visit::{Location, PlaceContext},
-    BasicBlockIdx, MirVisitor, Operand, Place, Rvalue, Statement, Terminator,
+    MirVisitor, Operand, Place, Rvalue, Statement, Terminator,
 };
 use std::collections::HashSet;
 
 pub struct InstrumentationVisitor<'a, 'tcx> {
-    /// Whether we should skip the next instruction, since it might've been instrumented already.
-    /// When we instrument an instruction, we partition the basic block, and the instruction that
-    /// may trigger UB becomes the first instruction of the basic block, which we need to skip
-    /// later.
-    skip_next: bool,
-    /// The instruction being visited at a given point.
-    current: SourceInstruction,
-    /// The target instruction that should be verified.
-    pub target: Option<InitRelevantInstruction>,
+    /// All target instructions in the body.
+    targets: Vec<InitRelevantInstruction>,
+    /// Currently analyzed instruction.
+    current_instruction: SourceInstruction,
+    /// Current analysis target, eventually needs to be added to a list of all targets.
+    current_target: InitRelevantInstruction,
     /// Aliasing analysis data.
     points_to: &'a PointsToGraph<'tcx>,
     /// The list of places we should be looking for, ignoring others
@@ -41,17 +38,14 @@ pub struct InstrumentationVisitor<'a, 'tcx> {
 }
 
 impl<'a, 'tcx> TargetFinder for InstrumentationVisitor<'a, 'tcx> {
-    fn find_next(
-        &mut self,
-        body: &MutableBody,
-        bb: BasicBlockIdx,
-        skip_first: bool,
-    ) -> Option<InitRelevantInstruction> {
-        self.skip_next = skip_first;
-        self.current = SourceInstruction::Statement { idx: 0, bb };
-        self.target = None;
-        self.visit_basic_block(&body.blocks()[bb]);
-        self.target.clone()
+    fn find_all(&mut self, body: &MutableBody) -> Vec<InitRelevantInstruction> {
+        for (bb_idx, bb) in body.blocks().iter().enumerate() {
+            self.current_instruction = SourceInstruction::Statement { idx: 0, bb: bb_idx };
+            self.visit_basic_block(bb);
+        }
+        // Push the last current target into the list.
+        self.targets.push(self.current_target.clone());
+        self.targets.clone()
     }
 }
 
@@ -63,43 +57,52 @@ impl<'a, 'tcx> InstrumentationVisitor<'a, 'tcx> {
         tcx: TyCtxt<'tcx>,
     ) -> Self {
         Self {
-            skip_next: false,
-            current: SourceInstruction::Statement { idx: 0, bb: 0 },
-            target: None,
+            targets: vec![],
+            current_instruction: SourceInstruction::Statement { idx: 0, bb: 0 },
+            current_target: InitRelevantInstruction {
+                source: SourceInstruction::Statement { idx: 0, bb: 0 },
+                before_instruction: vec![],
+                after_instruction: vec![],
+            },
             points_to,
             analysis_targets,
             current_instance,
             tcx,
         }
     }
+
     fn push_target(&mut self, source_op: MemoryInitOp) {
-        let target = self.target.get_or_insert_with(|| InitRelevantInstruction {
-            source: self.current,
-            after_instruction: vec![],
-            before_instruction: vec![],
-        });
-        target.push_operation(source_op);
+        // If we switched to the next instruction, push the old one onto the list of targets.
+        if self.current_target.source != self.current_instruction {
+            self.targets.push(self.current_target.clone());
+            self.current_target = InitRelevantInstruction {
+                source: self.current_instruction,
+                after_instruction: vec![],
+                before_instruction: vec![],
+            };
+        }
+        self.current_target.push_operation(source_op);
     }
 }
 
 impl<'a, 'tcx> MirVisitor for InstrumentationVisitor<'a, 'tcx> {
     fn visit_statement(&mut self, stmt: &Statement, location: Location) {
-        if self.skip_next {
-            self.skip_next = false;
-        } else if self.target.is_none() {
-            // Check all inner places.
-            self.super_statement(stmt, location);
-        }
+        self.super_statement(stmt, location);
         // Switch to the next statement.
-        let SourceInstruction::Statement { idx, bb } = self.current else { unreachable!() };
-        self.current = SourceInstruction::Statement { idx: idx + 1, bb };
+        if let SourceInstruction::Statement { idx, bb } = self.current_instruction {
+            self.current_instruction = SourceInstruction::Statement { idx: idx + 1, bb }
+        } else {
+            unreachable!()
+        }
     }
 
     fn visit_terminator(&mut self, term: &Terminator, location: Location) {
-        if !(self.skip_next || self.target.is_some()) {
-            self.current = SourceInstruction::Terminator { bb: self.current.bb() };
-            self.super_terminator(term, location);
+        if let SourceInstruction::Statement { bb, .. } = self.current_instruction {
+            self.current_instruction = SourceInstruction::Terminator { bb };
+        } else {
+            unreachable!()
         }
+        self.super_terminator(term, location);
     }
 
     fn visit_rvalue(&mut self, rvalue: &Rvalue, location: Location) {

kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/mod.rs

@@ -11,9 +11,8 @@ use crate::kani_middle::{
     points_to::{run_points_to_analysis, MemLoc, PointsToGraph},
     reachability::CallGraph,
     transform::{
-        body::{CheckType, MutableBody},
-        check_uninit::UninitInstrumenter,
-        BodyTransformation, GlobalPass, TransformationResult,
+        body::CheckType, check_uninit::UninitInstrumenter, BodyTransformation, GlobalPass,
+        TransformationResult,
     },
 };
 use crate::kani_queries::QueryDb;
@@ -112,25 +111,27 @@ impl GlobalPass for DelayedUbPass {
 
             // Instrument each instance based on the final targets we found.
             for instance in instances {
-                let mut instrumenter = UninitInstrumenter {
-                    check_type: self.check_type.clone(),
-                    mem_init_fn_cache: &mut self.mem_init_fn_cache,
-                };
                 // Retrieve the body with all local instrumentation passes applied.
-                let body = MutableBody::from(transformer.body(tcx, instance));
+                let body = transformer.body(tcx, instance);
                 // Instrument for delayed UB.
                 let target_finder = InstrumentationVisitor::new(
                     &global_points_to_graph,
                     &analysis_targets,
                     instance,
                     tcx,
                 );
-                let (instrumentation_added, body) =
-                    instrumenter.instrument(tcx, body, instance, target_finder);
+                let (instrumentation_added, body) = UninitInstrumenter::run(
+                    body,
+                    tcx,
+                    instance,
+                    self.check_type.clone(),
+                    &mut self.mem_init_fn_cache,
+                    target_finder,
+                );
                 // If some instrumentation has been performed, update the cached body in the local transformer.
                 if instrumentation_added {
                     transformer.cache.entry(instance).and_modify(|transformation_result| {
-                        *transformation_result = TransformationResult::Modified(body.into());
+                        *transformation_result = TransformationResult::Modified(body);
                     });
                 }
             }