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 });

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

@@ -18,18 +18,11 @@ use rustc_middle::ty::TyCtxt;
 use stable_mir::mir::{
     mono::Instance,
     visit::{Location, PlaceContext},
-    BasicBlockIdx, MirVisitor, Operand, Place, Rvalue, Statement, Terminator,
+    BasicBlock, MirVisitor, Operand, Place, Rvalue,
 };
 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>,
     /// Aliasing analysis data.
@@ -44,13 +37,23 @@ impl<'a, 'tcx> TargetFinder for InstrumentationVisitor<'a, 'tcx> {
     fn find_next(
         &mut self,
         body: &MutableBody,
-        bb: BasicBlockIdx,
-        skip_first: bool,
+        source: &SourceInstruction,
     ) -> Option<InitRelevantInstruction> {
-        self.skip_next = skip_first;
-        self.current = SourceInstruction::Statement { idx: 0, bb };
         self.target = None;
-        self.visit_basic_block(&body.blocks()[bb]);
+        match *source {
+            SourceInstruction::Statement { idx, bb } => {
+                let BasicBlock { statements, .. } = &body.blocks()[bb];
+                let stmt = &statements[idx];
+                self.visit_statement(stmt, self.__location_hack_remove_before_merging(stmt.span))
+            }
+            SourceInstruction::Terminator { bb } => {
+                let BasicBlock { terminator, .. } = &body.blocks()[bb];
+                self.visit_terminator(
+                    terminator,
+                    self.__location_hack_remove_before_merging(terminator.span),
+                )
+            }
+        }
         self.target.clone()
     }
 }
@@ -62,19 +65,10 @@ impl<'a, 'tcx> InstrumentationVisitor<'a, 'tcx> {
         current_instance: Instance,
         tcx: TyCtxt<'tcx>,
     ) -> Self {
-        Self {
-            skip_next: false,
-            current: SourceInstruction::Statement { idx: 0, bb: 0 },
-            target: None,
-            points_to,
-            analysis_targets,
-            current_instance,
-            tcx,
-        }
+        Self { target: None, 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![],
         });
@@ -83,25 +77,6 @@ impl<'a, 'tcx> InstrumentationVisitor<'a, 'tcx> {
 }
 
 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);
-        }
-        // Switch to the next statement.
-        let SourceInstruction::Statement { idx, bb } = self.current else { unreachable!() };
-        self.current = SourceInstruction::Statement { idx: idx + 1, bb };
-    }
-
-    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);
-        }
-    }
-
     fn visit_rvalue(&mut self, rvalue: &Rvalue, location: Location) {
         match rvalue {
             Rvalue::AddressOf(..) | Rvalue::Ref(..) => {

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);
                     });
                 }
             }