Track local frames incrementally during execution

src/concurrency/thread.rs

@@ -118,6 +118,9 @@ pub struct Thread<'mir, 'tcx> {
     /// The virtual call stack.
     stack: Vec<Frame<'mir, 'tcx, Provenance, FrameData<'tcx>>>,
 
+    /// The index of the topmost local frame in `stack`
+    pub(crate) top_local_frame: usize,
+
     /// The join status.
     join_status: ThreadJoinStatus,
 
@@ -167,6 +170,7 @@ impl<'mir, 'tcx> Default for Thread<'mir, 'tcx> {
             state: ThreadState::Enabled,
             thread_name: None,
             stack: Vec::new(),
+            top_local_frame: 0,
             join_status: ThreadJoinStatus::Joinable,
             panic_payload: None,
             last_error: None,
@@ -184,8 +188,15 @@ impl<'mir, 'tcx> Thread<'mir, 'tcx> {
 
 impl VisitTags for Thread<'_, '_> {
     fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
-        let Thread { panic_payload, last_error, stack, state: _, thread_name: _, join_status: _ } =
-            self;
+        let Thread {
+            panic_payload,
+            last_error,
+            stack,
+            top_local_frame: _,
+            state: _,
+            thread_name: _,
+            join_status: _,
+        } = self;
 
         panic_payload.visit_tags(visit);
         last_error.visit_tags(visit);
@@ -414,10 +425,22 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
     }
 
     /// Get a shared borrow of the currently active thread.
-    fn active_thread_ref(&self) -> &Thread<'mir, 'tcx> {
+    pub fn active_thread_ref(&self) -> &Thread<'mir, 'tcx> {
         &self.threads[self.active_thread]
     }
 
+    /// Reset the current-span caching logic. Mostly called when a new frame is added or removed.
+    pub fn reset_top_local(&mut self) {
+        let idx = self
+            .active_thread_stack()
+            .iter()
+            .enumerate()
+            .rev()
+            .find_map(|(idx, frame)| if frame.extra.is_local { Some(idx) } else { None })
+            .unwrap_or(self.active_thread_stack().len().saturating_sub(1));
+        self.active_thread_mut().top_local_frame = idx;
+    }
+
     /// Mark the thread as detached, which means that no other thread will try
     /// to join it and the thread is responsible for cleaning up.
     ///

src/helpers.rs

@@ -945,7 +945,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
 impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
     pub fn current_span(&self) -> CurrentSpan<'_, 'mir, 'tcx> {
-        CurrentSpan { current_frame_idx: None, machine: self }
+        CurrentSpan { machine: self }
     }
 }
 
@@ -955,7 +955,6 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
 /// The result of that search is cached so that later calls are approximately free.
 #[derive(Clone)]
 pub struct CurrentSpan<'a, 'mir, 'tcx> {
-    current_frame_idx: Option<usize>,
     machine: &'a MiriMachine<'mir, 'tcx>,
 }
 
@@ -967,8 +966,10 @@ impl<'a, 'mir: 'a, 'tcx: 'a + 'mir> CurrentSpan<'a, 'mir, 'tcx> {
     /// Get the current span, skipping non-local frames.
     /// This function is backed by a cache, and can be assumed to be very fast.
     pub fn get(&mut self) -> Span {
-        let idx = self.current_frame_idx();
-        self.stack().get(idx).map(Frame::current_span).unwrap_or(rustc_span::DUMMY_SP)
+        self.stack()
+            .get(self.top_local_frame())
+            .map(Frame::current_span)
+            .unwrap_or(rustc_span::DUMMY_SP)
     }
 
     /// Returns the span of the *caller* of the current operation, again
@@ -979,7 +980,7 @@ impl<'a, 'mir: 'a, 'tcx: 'a + 'mir> CurrentSpan<'a, 'mir, 'tcx> {
     pub fn get_caller(&mut self) -> Span {
         // We need to go down at least to the caller (len - 2), or however
         // far we have to go to find a frame in a local crate.
-        let local_frame_idx = self.current_frame_idx();
+        let local_frame_idx = self.top_local_frame();
         let stack = self.stack();
         let idx = cmp::min(local_frame_idx, stack.len().saturating_sub(2));
         stack.get(idx).map(Frame::current_span).unwrap_or(rustc_span::DUMMY_SP)
@@ -989,33 +990,8 @@ impl<'a, 'mir: 'a, 'tcx: 'a + 'mir> CurrentSpan<'a, 'mir, 'tcx> {
         self.machine.threads.active_thread_stack()
     }
 
-    fn current_frame_idx(&mut self) -> usize {
-        *self
-            .current_frame_idx
-            .get_or_insert_with(|| Self::compute_current_frame_index(self.machine))
-    }
-
-    // Find the position of the inner-most frame which is part of the crate being
-    // compiled/executed, part of the Cargo workspace, and is also not #[track_caller].
-    #[inline(never)]
-    fn compute_current_frame_index(machine: &MiriMachine<'_, '_>) -> usize {
-        machine
-            .threads
-            .active_thread_stack()
-            .iter()
-            .enumerate()
-            .rev()
-            .find_map(|(idx, frame)| {
-                let def_id = frame.instance.def_id();
-                if (def_id.is_local() || machine.local_crates.contains(&def_id.krate))
-                    && !frame.instance.def.requires_caller_location(machine.tcx)
-                {
-                    Some(idx)
-                } else {
-                    None
-                }
-            })
-            .unwrap_or(0)
+    fn top_local_frame(&self) -> usize {
+        self.machine.threads.active_thread_ref().top_local_frame
     }
 }
 

src/machine.rs

@@ -50,12 +50,14 @@ pub struct FrameData<'tcx> {
     /// for the start of this frame. When we finish executing this frame,
     /// we use this to register a completed event with `measureme`.
     pub timing: Option<measureme::DetachedTiming>,
+
+    pub is_local: bool,
 }
 
 impl<'tcx> std::fmt::Debug for FrameData<'tcx> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         // Omitting `timing`, it does not support `Debug`.
-        let FrameData { stacked_borrows, catch_unwind, timing: _ } = self;
+        let FrameData { stacked_borrows, catch_unwind, timing: _, is_local: _ } = self;
         f.debug_struct("FrameData")
             .field("stacked_borrows", stacked_borrows)
             .field("catch_unwind", catch_unwind)
@@ -65,7 +67,7 @@ impl<'tcx> std::fmt::Debug for FrameData<'tcx> {
 
 impl VisitTags for FrameData<'_> {
     fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
-        let FrameData { catch_unwind, stacked_borrows, timing: _ } = self;
+        let FrameData { catch_unwind, stacked_borrows, timing: _, is_local: _ } = self;
 
         catch_unwind.visit_tags(visit);
         stacked_borrows.visit_tags(visit);
@@ -1004,6 +1006,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         (alloc_id, prov_extra): (AllocId, Self::ProvenanceExtra),
         range: AllocRange,
     ) -> InterpResult<'tcx> {
+        let mut current_span = machine.current_span();
         if let Some(data_race) = &alloc_extra.data_race {
             data_race.read(
                 alloc_id,
@@ -1018,7 +1021,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
                 prov_extra,
                 range,
                 machine.stacked_borrows.as_ref().unwrap(),
-                machine.current_span(),
+                &mut current_span,
                 &machine.threads,
             )?;
         }
@@ -1036,21 +1039,23 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         (alloc_id, prov_extra): (AllocId, Self::ProvenanceExtra),
         range: AllocRange,
     ) -> InterpResult<'tcx> {
-        if let Some(data_race) = &mut alloc_extra.data_race {
-            data_race.write(
+        let mut current_span = machine.current_span();
+        if let Some(stacked_borrows) = &mut alloc_extra.stacked_borrows {
+            stacked_borrows.get_mut().before_memory_write(
                 alloc_id,
+                prov_extra,
                 range,
-                machine.data_race.as_mut().unwrap(),
+                machine.stacked_borrows.as_ref().unwrap(),
+                &mut current_span,
                 &machine.threads,
             )?;
         }
-        if let Some(stacked_borrows) = &mut alloc_extra.stacked_borrows {
-            stacked_borrows.get_mut().before_memory_write(
+        if let Some(data_race) = &mut alloc_extra.data_race {
+            let _span = current_span.get();
+            data_race.write(
                 alloc_id,
-                prov_extra,
                 range,
-                machine.stacked_borrows.as_ref().unwrap(),
-                machine.current_span(),
+                machine.data_race.as_mut().unwrap(),
                 &machine.threads,
             )?;
         }
@@ -1071,26 +1076,26 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         if machine.tracked_alloc_ids.contains(&alloc_id) {
             machine.emit_diagnostic(NonHaltingDiagnostic::FreedAlloc(alloc_id));
         }
-        if let Some(data_race) = &mut alloc_extra.data_race {
-            data_race.deallocate(
+        let mut current_span = machine.current_span();
+        if let Some(stacked_borrows) = &mut alloc_extra.stacked_borrows {
+            stacked_borrows.get_mut().before_memory_deallocation(
                 alloc_id,
+                prove_extra,
                 range,
-                machine.data_race.as_mut().unwrap(),
+                machine.stacked_borrows.as_ref().unwrap(),
+                &mut current_span,
                 &machine.threads,
             )?;
         }
-        if let Some(stacked_borrows) = &mut alloc_extra.stacked_borrows {
-            stacked_borrows.get_mut().before_memory_deallocation(
+        if let Some(data_race) = &mut alloc_extra.data_race {
+            data_race.deallocate(
                 alloc_id,
-                prove_extra,
                 range,
-                machine.stacked_borrows.as_ref().unwrap(),
-                machine.current_span(),
+                machine.data_race.as_mut().unwrap(),
                 &machine.threads,
-            )
-        } else {
-            Ok(())
+            )?;
         }
+        Ok(())
     }
 
     #[inline(always)]
@@ -1122,13 +1127,24 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
             None
         };
 
+        let def_id = frame.instance.def_id();
+        let is_local = (def_id.is_local() || ecx.machine.local_crates.contains(&def_id.krate))
+            && !frame.instance.def.requires_caller_location(ecx.machine.tcx);
+
+        if is_local {
+            let stack_len = ecx.active_thread_stack().len();
+            ecx.active_thread_mut().top_local_frame = stack_len;
+        }
+
         let stacked_borrows = ecx.machine.stacked_borrows.as_ref();
 
         let extra = FrameData {
             stacked_borrows: stacked_borrows.map(|sb| sb.borrow_mut().new_frame(&ecx.machine)),
             catch_unwind: None,
             timing,
+            is_local,
         };
+
         Ok(frame.with_extra(extra))
     }
 
@@ -1189,6 +1205,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         if let Some(stacked_borrows) = &ecx.machine.stacked_borrows {
             stacked_borrows.borrow_mut().end_call(&frame.extra);
         }
+        ecx.machine.threads.reset_top_local();
         let res = ecx.handle_stack_pop_unwind(frame.extra, unwinding);
         if let Some(profiler) = ecx.machine.profiler.as_ref() {
             profiler.finish_recording_interval_event(timing.unwrap());

src/stacked_borrows/mod.rs

@@ -621,7 +621,7 @@ impl Stacks {
         tag: ProvenanceExtra,
         range: AllocRange,
         state: &GlobalState,
-        mut current_span: CurrentSpan<'ecx, 'mir, 'tcx>,
+        current_span: &mut CurrentSpan<'ecx, 'mir, 'tcx>,
         threads: &'ecx ThreadManager<'mir, 'tcx>,
     ) -> InterpResult<'tcx>
     where
@@ -633,7 +633,7 @@ impl Stacks {
             Pointer::new(alloc_id, range.start),
             range.size.bytes()
         );
-        let dcx = DiagnosticCxBuilder::read(&mut current_span, threads, tag, range);
+        let dcx = DiagnosticCxBuilder::read(current_span, threads, tag, range);
         let mut state = state.borrow_mut();
         self.for_each(range, dcx, |stack, dcx, exposed_tags| {
             stack.access(AccessKind::Read, tag, &mut state, dcx, exposed_tags)
@@ -647,7 +647,7 @@ impl Stacks {
         tag: ProvenanceExtra,
         range: AllocRange,
         state: &GlobalState,
-        mut current_span: CurrentSpan<'ecx, 'mir, 'tcx>,
+        current_span: &mut CurrentSpan<'ecx, 'mir, 'tcx>,
         threads: &'ecx ThreadManager<'mir, 'tcx>,
     ) -> InterpResult<'tcx> {
         trace!(
@@ -656,7 +656,7 @@ impl Stacks {
             Pointer::new(alloc_id, range.start),
             range.size.bytes()
         );
-        let dcx = DiagnosticCxBuilder::write(&mut current_span, threads, tag, range);
+        let dcx = DiagnosticCxBuilder::write(current_span, threads, tag, range);
         let mut state = state.borrow_mut();
         self.for_each(range, dcx, |stack, dcx, exposed_tags| {
             stack.access(AccessKind::Write, tag, &mut state, dcx, exposed_tags)
@@ -670,11 +670,11 @@ impl Stacks {
         tag: ProvenanceExtra,
         range: AllocRange,
         state: &GlobalState,
-        mut current_span: CurrentSpan<'ecx, 'mir, 'tcx>,
+        current_span: &mut CurrentSpan<'ecx, 'mir, 'tcx>,
         threads: &'ecx ThreadManager<'mir, 'tcx>,
     ) -> InterpResult<'tcx> {
         trace!("deallocation with tag {:?}: {:?}, size {}", tag, alloc_id, range.size.bytes());
-        let dcx = DiagnosticCxBuilder::dealloc(&mut current_span, threads, tag);
+        let dcx = DiagnosticCxBuilder::dealloc(current_span, threads, tag);
         let state = state.borrow();
         self.for_each(range, dcx, |stack, dcx, exposed_tags| {
             stack.dealloc(tag, &state, dcx, exposed_tags)