Fix an issue with terminator handling in InstrumentationVisitor

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

@@ -469,6 +469,12 @@ impl SourceInstruction {
             SourceInstruction::Terminator { bb } => blocks[bb].terminator.span,
         }
     }
+
+    pub fn bb(&self) -> usize {
+        match self {
+            SourceInstruction::Statement { bb, .. } | SourceInstruction::Terminator { bb } => *bb,
+        }
+    }
 }
 
 fn find_instance(tcx: TyCtxt, diagnostic: &str) -> Option<Instance> {

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

@@ -16,7 +16,7 @@ use rustc_hir::def_id::DefId as InternalDefId;
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
 use stable_mir::mir::visit::{Location, PlaceContext};
-use stable_mir::mir::{BasicBlockIdx, MirVisitor, Operand, Place, Statement};
+use stable_mir::mir::{BasicBlockIdx, MirVisitor, Operand, Place, Statement, Terminator};
 use std::collections::HashSet;
 
 pub struct InstrumentationVisitor<'a, 'tcx> {
@@ -92,6 +92,13 @@ impl<'a, 'tcx> MirVisitor for InstrumentationVisitor<'a, 'tcx> {
         }
     }
 
+    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_place(&mut self, place: &Place, ptx: PlaceContext, location: Location) {
         // Match the place by whatever it is pointing to and find an intersection with the targets.
         if self