Avoid corner-cases by grouping instrumentation into basic blocks and using backward iteration (#3438)

This PR is a follow-up to
#3374.

Its main purpose is to properly handle a corner-case when multiple
instrumentation instructions are added to the same instruction and not
all of them are skipped during subsequent instrumentation.

For example, if instrumentation is added after a terminator, a new basic
block will be created, containing the added instrumentation. However, we
currently only mark the first statement (or the terminator, if there are
none) of the new basic block as skipped for subsequent instrumentation.
That means that if instrumentation in this case contains some
instrumentation targets itself, it will never terminate.

Coincidentally, this is not currently the case, but could lead to subtle
bugs if we decide to change instrumentation. In fact, this bug was only
surfaced when I experimented with checking all memory accesses, which
introduced significantly more checks. Ultimately, this shows that the
current way to avoiding instrumentation is limited and needs to be
changed.

The PR introduces the following changes:

- Group instrumentation into separate basic blocks instead of adding
instructions one-by-one.
- Use backward iteration to avoid having to reason about which
instructions need to be skipped.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.

Author: artemagvanian

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,15 @@ 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>,
+    /// 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 +36,16 @@ 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_target = InitRelevantInstruction {
+                source: SourceInstruction::Statement { idx: 0, bb: bb_idx },
+                before_instruction: vec![],
+                after_instruction: vec![],
+            };
+            self.visit_basic_block(bb);
+        }
+        self.targets
     }
 }
 
@@ -63,43 +57,55 @@ 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_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);
+        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_target.source {
+            self.targets.push(self.current_target.clone());
+            self.current_target = InitRelevantInstruction {
+                source: SourceInstruction::Statement { idx: idx + 1, bb },
+                after_instruction: vec![],
+                before_instruction: vec![],
+            };
+        } 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_target.source {
+            // We don't have to push the previous target, since it already happened in the statement
+            // handling code.
+            self.current_target = InitRelevantInstruction {
+                source: SourceInstruction::Terminator { bb },
+                after_instruction: vec![],
+                before_instruction: vec![],
+            };
+        } else {
+            unreachable!()
         }
+        self.super_terminator(term, location);
+        // Push the current target from the terminator onto the list.
+        self.targets.push(self.current_target.clone());
     }
 
     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);
                     });
                 }
             }