Auto merge of #2139 - saethlin:lazy-current-span, r=RalfJung

Factor current-span logic into a caching handle

After #2030 and while working on #1935 it became quite clear that we need to do some caching here, because some retag operations generate many calls to `log_invalidation`, and would thus search the current thread's stack _many_ times for a local crate. This caching fixes that. This handle type also has the nice benefit of tucking away all the `ThreadManager` + `CrateNum` logic.

src/helpers.rs

@@ -2,7 +2,6 @@ pub mod convert;
 
 use std::mem;
 use std::num::NonZeroUsize;
-use std::rc::Rc;
 use std::time::Duration;
 
 use log::trace;
@@ -14,7 +13,7 @@ use rustc_middle::ty::{
     layout::{LayoutOf, TyAndLayout},
     List, TyCtxt,
 };
-use rustc_span::{def_id::CrateNum, sym, Symbol};
+use rustc_span::{def_id::CrateNum, sym, Span, Symbol};
 use rustc_target::abi::{Align, FieldsShape, Size, Variants};
 use rustc_target::spec::abi::Abi;
 
@@ -800,6 +799,43 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
     }
 }
 
+impl<'mir, 'tcx> Evaluator<'mir, 'tcx> {
+    pub fn current_span(&self) -> CurrentSpan<'_, 'mir, 'tcx> {
+        CurrentSpan { span: 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, 'tcx, 'mir> {
+    span: Option<Span>,
+    machine: &'a Evaluator<'tcx, 'mir>,
+}
+
+impl<'a, 'tcx, 'mir> CurrentSpan<'a, 'tcx, 'mir> {
+    pub fn get(&mut self) -> Span {
+        *self.span.get_or_insert_with(|| Self::current_span(&self.machine))
+    }
+
+    #[inline(never)]
+    fn current_span(machine: &Evaluator<'_, '_>) -> Span {
+        machine
+            .threads
+            .active_thread_stack()
+            .into_iter()
+            .rev()
+            .find(|frame| {
+                let def_id = frame.instance.def_id();
+                def_id.is_local() || machine.local_crates.contains(&def_id.krate)
+            })
+            .map(|frame| frame.current_span())
+            .unwrap_or(rustc_span::DUMMY_SP)
+    }
+}
+
 /// Check that the number of args is what we expect.
 pub fn check_arg_count<'a, 'tcx, const N: usize>(
     args: &'a [OpTy<'tcx, Tag>],
@@ -822,7 +858,7 @@ pub fn isolation_abort_error(name: &str) -> InterpResult<'static> {
 
 /// Retrieve the list of local crates that should have been passed by cargo-miri in
 /// MIRI_LOCAL_CRATES and turn them into `CrateNum`s.
-pub fn get_local_crates(tcx: &TyCtxt<'_>) -> Rc<[CrateNum]> {
+pub fn get_local_crates(tcx: &TyCtxt<'_>) -> Vec<CrateNum> {
     // Convert the local crate names from the passed-in config into CrateNums so that they can
     // be looked up quickly during execution
     let local_crate_names = std::env::var("MIRI_LOCAL_CRATES")
@@ -836,7 +872,7 @@ pub fn get_local_crates(tcx: &TyCtxt<'_>) -> Rc<[CrateNum]> {
             local_crates.push(crate_num);
         }
     }
-    Rc::from(local_crates.as_slice())
+    local_crates
 }
 
 /// Formats an AllocRange like [0x1..0x3], for use in diagnostics.

src/lib.rs

@@ -77,7 +77,7 @@ pub use crate::diagnostics::{
 pub use crate::eval::{
     create_ecx, eval_entry, AlignmentCheck, BacktraceStyle, IsolatedOp, MiriConfig, RejectOpWith,
 };
-pub use crate::helpers::EvalContextExt as HelpersEvalContextExt;
+pub use crate::helpers::{CurrentSpan, EvalContextExt as HelpersEvalContextExt};
 pub use crate::machine::{
     AllocExtra, Evaluator, FrameData, MiriEvalContext, MiriEvalContextExt, MiriMemoryKind, Tag,
     NUM_CPUS, PAGE_SIZE, STACK_ADDR, STACK_SIZE,

src/machine.rs

@@ -6,7 +6,6 @@ use std::cell::RefCell;
 use std::collections::HashSet;
 use std::fmt;
 use std::num::NonZeroU64;
-use std::rc::Rc;
 use std::time::Instant;
 
 use rand::rngs::StdRng;
@@ -278,7 +277,7 @@ pub struct Evaluator<'mir, 'tcx> {
     pub(crate) backtrace_style: BacktraceStyle,
 
     /// Crates which are considered local for the purposes of error reporting.
-    pub(crate) local_crates: Rc<[CrateNum]>,
+    pub(crate) local_crates: Vec<CrateNum>,
 
     /// Mapping extern static names to their base pointer.
     extern_statics: FxHashMap<Symbol, Pointer<Tag>>,
@@ -584,8 +583,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'mir, 'tcx> {
                 alloc.size(),
                 stacked_borrows,
                 kind,
-                &ecx.machine.threads,
-                ecx.machine.local_crates.clone(),
+                ecx.machine.current_span(),
             ))
         } else {
             None
@@ -667,7 +665,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'mir, 'tcx> {
                 tag,
                 range,
                 machine.stacked_borrows.as_ref().unwrap(),
-                &machine.threads,
+                machine.current_span(),
             )
         } else {
             Ok(())
@@ -691,7 +689,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'mir, 'tcx> {
                 tag,
                 range,
                 machine.stacked_borrows.as_ref().unwrap(),
-                &machine.threads,
+                machine.current_span(),
             )
         } else {
             Ok(())

src/stacked_borrows.rs

@@ -5,7 +5,6 @@ use log::trace;
 use std::cell::RefCell;
 use std::fmt;
 use std::num::NonZeroU64;
-use std::rc::Rc;
 
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_hir::Mutability;
@@ -14,17 +13,14 @@ use rustc_middle::ty::{
     self,
     layout::{HasParamEnv, LayoutOf},
 };
-use rustc_span::def_id::CrateNum;
 use rustc_span::DUMMY_SP;
 use rustc_target::abi::Size;
 use std::collections::HashSet;
 
 use crate::*;
 
 pub mod diagnostics;
-use diagnostics::AllocHistory;
-
-use diagnostics::TagHistory;
+use diagnostics::{AllocHistory, TagHistory};
 
 pub type PtrId = NonZeroU64;
 pub type CallId = NonZeroU64;
@@ -376,7 +372,7 @@ impl<'tcx> Stack {
         tag: SbTag,
         (alloc_id, alloc_range, offset): (AllocId, AllocRange, Size), // just for debug printing and error messages
         global: &mut GlobalStateInner,
-        threads: &ThreadManager<'_, 'tcx>,
+        current_span: &mut CurrentSpan<'_, '_, 'tcx>,
         alloc_history: &mut AllocHistory,
     ) -> InterpResult<'tcx> {
         // Two main steps: Find granting item, remove incompatible items above.
@@ -400,7 +396,7 @@ impl<'tcx> Stack {
                     global,
                     alloc_history,
                 )?;
-                alloc_history.log_invalidation(item.tag, alloc_range, threads);
+                alloc_history.log_invalidation(item.tag, alloc_range, current_span);
             }
         } else {
             // On a read, *disable* all `Unique` above the granting item.  This ensures U2 for read accesses.
@@ -422,7 +418,7 @@ impl<'tcx> Stack {
                         alloc_history,
                     )?;
                     item.perm = Permission::Disabled;
-                    alloc_history.log_invalidation(item.tag, alloc_range, threads);
+                    alloc_history.log_invalidation(item.tag, alloc_range, current_span);
                 }
             }
         }
@@ -471,7 +467,7 @@ impl<'tcx> Stack {
         new: Item,
         (alloc_id, alloc_range, offset): (AllocId, AllocRange, Size), // just for debug printing and error messages
         global: &mut GlobalStateInner,
-        threads: &ThreadManager<'_, 'tcx>,
+        current_span: &mut CurrentSpan<'_, '_, 'tcx>,
         alloc_history: &mut AllocHistory,
     ) -> InterpResult<'tcx> {
         // Figure out which access `perm` corresponds to.
@@ -505,7 +501,7 @@ impl<'tcx> Stack {
                 derived_from,
                 (alloc_id, alloc_range, offset),
                 global,
-                threads,
+                current_span,
                 alloc_history,
             )?;
 
@@ -533,13 +529,13 @@ impl<'tcx> Stack {
 /// Map per-stack operations to higher-level per-location-range operations.
 impl<'tcx> Stacks {
     /// Creates new stack with initial tag.
-    fn new(size: Size, perm: Permission, tag: SbTag, local_crates: Rc<[CrateNum]>) -> Self {
+    fn new(size: Size, perm: Permission, tag: SbTag) -> Self {
         let item = Item { perm, tag, protector: None };
         let stack = Stack { borrows: vec![item] };
 
         Stacks {
             stacks: RefCell::new(RangeMap::new(size, stack)),
-            history: RefCell::new(AllocHistory::new(local_crates)),
+            history: RefCell::new(AllocHistory::new()),
         }
     }
 
@@ -579,8 +575,7 @@ impl Stacks {
         size: Size,
         state: &GlobalState,
         kind: MemoryKind<MiriMemoryKind>,
-        threads: &ThreadManager<'_, '_>,
-        local_crates: Rc<[CrateNum]>,
+        mut current_span: CurrentSpan<'_, '_, '_>,
     ) -> Self {
         let mut extra = state.borrow_mut();
         let (base_tag, perm) = match kind {
@@ -614,12 +609,12 @@ impl Stacks {
                 (tag, Permission::SharedReadWrite)
             }
         };
-        let stacks = Stacks::new(size, perm, base_tag, local_crates);
+        let stacks = Stacks::new(size, perm, base_tag);
         stacks.history.borrow_mut().log_creation(
             None,
             base_tag,
             alloc_range(Size::ZERO, size),
-            threads,
+            &mut current_span,
         );
         stacks
     }
@@ -631,7 +626,7 @@ impl Stacks {
         tag: SbTag,
         range: AllocRange,
         state: &GlobalState,
-        threads: &ThreadManager<'_, 'tcx>,
+        mut current_span: CurrentSpan<'_, '_, 'tcx>,
     ) -> InterpResult<'tcx> {
         trace!(
             "read access with tag {:?}: {:?}, size {}",
@@ -646,7 +641,7 @@ impl Stacks {
                 tag,
                 (alloc_id, range, offset),
                 &mut state,
-                threads,
+                &mut current_span,
                 history,
             )
         })
@@ -659,7 +654,7 @@ impl Stacks {
         tag: SbTag,
         range: AllocRange,
         state: &GlobalState,
-        threads: &ThreadManager<'_, 'tcx>,
+        mut current_span: CurrentSpan<'_, '_, 'tcx>,
     ) -> InterpResult<'tcx> {
         trace!(
             "write access with tag {:?}: {:?}, size {}",
@@ -674,7 +669,7 @@ impl Stacks {
                 tag,
                 (alloc_id, range, offset),
                 &mut state,
-                threads,
+                &mut current_span,
                 history,
             )
         })
@@ -723,6 +718,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         }
         let (alloc_id, base_offset, orig_tag) = this.ptr_get_alloc_id(place.ptr)?;
 
+        let mut current_span = this.machine.current_span();
         {
             let extra = this.get_alloc_extra(alloc_id)?;
             let stacked_borrows =
@@ -732,10 +728,10 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
                 Some(orig_tag),
                 new_tag,
                 alloc_range(base_offset, size),
-                &this.machine.threads,
+                &mut current_span,
             );
             if protect {
-                alloc_history.log_protector(orig_tag, new_tag, &this.machine.threads);
+                alloc_history.log_protector(orig_tag, new_tag, &mut current_span);
             }
         }
 
@@ -804,7 +800,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
                             item,
                             (alloc_id, range, offset),
                             &mut *global,
-                            &this.machine.threads,
+                            &mut current_span,
                             history,
                         )
                     })
@@ -821,13 +817,14 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         let item = Item { perm, tag: new_tag, protector };
         let range = alloc_range(base_offset, size);
         let mut global = machine.stacked_borrows.as_ref().unwrap().borrow_mut();
+        let mut current_span = machine.current_span();
         stacked_borrows.for_each_mut(range, |offset, stack, history| {
             stack.grant(
                 orig_tag,
                 item,
                 (alloc_id, range, offset),
                 &mut global,
-                &machine.threads,
+                &mut current_span,
                 history,
             )
         })?;