diff --git a/kani-compiler/src/kani_middle/transform/body.rs b/kani-compiler/src/kani_middle/transform/body.rs index d3f2afc15d31..e3de35588475 100644 --- a/kani-compiler/src/kani_middle/transform/body.rs +++ b/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 { diff --git a/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/instrumentation_visitor.rs b/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/instrumentation_visitor.rs index 2fb43172e05e..21fb3f42bc77 100644 --- a/kani-compiler/src/kani_middle/transform/check_uninit/delayed_ub/instrumentation_visitor.rs +++ b/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