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 });
@@ -503,6 +503,10 @@ impl SourceInstruction {
             SourceInstruction::Statement { bb, .. } | SourceInstruction::Terminator { bb } => bb,
         }
     }
+
+    pub fn is_terminator(&self) -> bool {
+        matches!(self, SourceInstruction::Terminator { .. })
+    }
 }
 
 fn find_instance(tcx: TyCtxt, diagnostic: &str) -> Option<Instance> {

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