Track local frames incrementally during execution

src/concurrency/thread.rs

@@ -1,5 +1,6 @@
 //! Implements threads.
 
+use std::cell::Cell;
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
 use std::num::TryFromIntError;
@@ -118,6 +119,17 @@ 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`. If `Some`, this field must contain
+    /// the value produced by `reset_top_local`.
+    /// This field is a cache to reduce how often we call that method. The cache is invalidated by
+    /// mutation of the stack, so calls to `active_thread_stack_mut` set this to `None`.
+    ///
+    /// The outer `Option` represents the validity of the cache, and the inner `Option` represents
+    /// whether or not there is a local frame. Being able to differentiate between knowing that we
+    /// have no local frame and not knowing whether there is a local frame reduces how often the
+    /// cache is invalidated.
+    top_local_frame: Cell<Option<Option<usize>>>,
+
     /// The join status.
     join_status: ThreadJoinStatus,
 
@@ -147,6 +159,35 @@ impl<'mir, 'tcx> Thread<'mir, 'tcx> {
     fn thread_name(&self) -> &[u8] {
         if let Some(ref thread_name) = self.thread_name { thread_name } else { b"<unnamed>" }
     }
+
+    fn compute_top_local_frame(&self) -> Option<usize> {
+        self.stack
+            .iter()
+            .enumerate()
+            .rev()
+            .find_map(|(idx, frame)| if frame.extra.is_local { Some(idx) } else { None })
+    }
+
+    pub fn top_local_frame(&self) -> usize {
+        let top_local = if let Some(idx) = self.top_local_frame.get() {
+            debug_assert_eq!(idx, self.compute_top_local_frame());
+            idx
+        } else {
+            let idx = self.compute_top_local_frame();
+            self.top_local_frame.set(Some(idx));
+            idx
+        };
+        top_local.unwrap_or_else(|| self.stack.len().saturating_sub(1))
+    }
+
+    pub fn set_top_local_frame(&mut self, frame_idx: Option<usize>) {
+        if let Some(idx) = frame_idx {
+            self.top_local_frame.set(Some(Some(idx)));
+            debug_assert_eq!(Some(idx), self.compute_top_local_frame());
+        } else {
+            self.top_local_frame.set(None);
+        }
+    }
 }
 
 impl<'mir, 'tcx> std::fmt::Debug for Thread<'mir, 'tcx> {
@@ -167,6 +208,7 @@ impl<'mir, 'tcx> Default for Thread<'mir, 'tcx> {
             state: ThreadState::Enabled,
             thread_name: None,
             stack: Vec::new(),
+            top_local_frame: Cell::new(None),
             join_status: ThreadJoinStatus::Joinable,
             panic_payload: None,
             last_error: None,
@@ -184,8 +226,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 +463,16 @@ 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 of the active thread.
+    /// Mostly called when a new frame is added or removed.
+    pub fn top_local_frame(&self) -> usize {
+        self.active_thread_ref().top_local_frame()
+    }
+
     /// 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

@@ -936,78 +936,35 @@ 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 }
-    }
-}
-
-/// A `CurrentSpan` should be created infrequently (ideally once) per interpreter step. It does
-/// nothing on creation, but when `CurrentSpan::get` is called, searches the current stack for the
-/// topmost frame which corresponds to a local crate, and returns the current span in that frame.
-/// 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>,
-}
-
-impl<'a, 'mir: 'a, 'tcx: 'a + 'mir> CurrentSpan<'a, 'mir, 'tcx> {
-    pub fn machine(&self) -> &'a MiriMachine<'mir, 'tcx> {
-        self.machine
-    }
-
     /// 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)
+    pub fn current_span(&self) -> Span {
+        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
     /// walking down the stack to find the closest frame in a local crate, if the caller of the
     /// current operation is not in a local crate.
     /// This is useful when we are processing something which occurs on function-entry and we want
     /// to point at the call to the function, not the function definition generally.
-    pub fn get_caller(&mut self) -> Span {
+    pub fn caller_span(&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)
     }
 
     fn stack(&self) -> &[Frame<'mir, 'tcx, Provenance, machine::FrameData<'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))
+        self.threads.active_thread_stack()
     }
 
-    // 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.threads.active_thread_ref().top_local_frame()
     }
 }
 

src/lib.rs

@@ -97,7 +97,7 @@ pub use crate::diagnostics::{
 pub use crate::eval::{
     create_ecx, eval_entry, AlignmentCheck, BacktraceStyle, IsolatedOp, MiriConfig, RejectOpWith,
 };
-pub use crate::helpers::{CurrentSpan, EvalContextExt as _};
+pub use crate::helpers::EvalContextExt as _;
 pub use crate::intptrcast::ProvenanceMode;
 pub use crate::machine::{
     AllocExtra, FrameData, MiriInterpCx, MiriInterpCxExt, MiriMachine, MiriMemoryKind,

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);
@@ -895,13 +897,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
 
         let alloc = alloc.into_owned();
         let stacks = ecx.machine.stacked_borrows.as_ref().map(|stacked_borrows| {
-            Stacks::new_allocation(
-                id,
-                alloc.size(),
-                stacked_borrows,
-                kind,
-                ecx.machine.current_span(),
-            )
+            Stacks::new_allocation(id, alloc.size(), stacked_borrows, kind, &ecx.machine)
         });
         let race_alloc = ecx.machine.data_race.as_ref().map(|data_race| {
             data_race::AllocExtra::new_allocation(
@@ -1016,8 +1012,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
                 prov_extra,
                 range,
                 machine.stacked_borrows.as_ref().unwrap(),
-                machine.current_span(),
-                &machine.threads,
+                machine,
             )?;
         }
         if let Some(weak_memory) = &alloc_extra.weak_memory {
@@ -1048,8 +1043,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
                 prov_extra,
                 range,
                 machine.stacked_borrows.as_ref().unwrap(),
-                machine.current_span(),
-                &machine.threads,
+                machine,
             )?;
         }
         if let Some(weak_memory) = &alloc_extra.weak_memory {
@@ -1083,8 +1077,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
                 prove_extra,
                 range,
                 machine.stacked_borrows.as_ref().unwrap(),
-                machine.current_span(),
-                &machine.threads,
+                machine,
             )
         } else {
             Ok(())
@@ -1120,13 +1113,19 @@ 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);
+
         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))
     }
 
@@ -1174,6 +1173,13 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
 
     #[inline(always)]
     fn after_stack_push(ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx> {
+        if ecx.frame().extra.is_local {
+            // We just pushed a local frame, so we know that the topmost local frame is the topmost
+            // frame. If we push a non-local frame, there's no need to do anything.
+            let stack_len = ecx.active_thread_stack().len();
+            ecx.active_thread_mut().set_top_local_frame(Some(stack_len - 1));
+        }
+
         if ecx.machine.stacked_borrows.is_some() { ecx.retag_return_place() } else { Ok(()) }
     }
 
@@ -1183,6 +1189,11 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         mut frame: Frame<'mir, 'tcx, Provenance, FrameData<'tcx>>,
         unwinding: bool,
     ) -> InterpResult<'tcx, StackPopJump> {
+        if frame.extra.is_local {
+            // All that we store is whether or not the frame we just removed is local, so now we
+            // have no idea what if any the next topmost local frame is.
+            ecx.active_thread_mut().set_top_local_frame(None);
+        }
         let timing = frame.extra.timing.take();
         if let Some(stacked_borrows) = &ecx.machine.stacked_borrows {
             stacked_borrows.borrow_mut().end_call(&frame.extra);