Expand Miri's BorTag GC to a Provenance GC

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -107,6 +107,14 @@ impl<K: Hash + Eq, V> interpret::AllocMap<K, V> for FxIndexMap<K, V> {
         FxIndexMap::contains_key(self, k)
     }
 
+    #[inline(always)]
+    fn contains_key_ref<Q: ?Sized + Hash + Eq>(&self, k: &Q) -> bool
+    where
+        K: Borrow<Q>,
+    {
+        FxIndexMap::contains_key(self, k)
+    }
+
     #[inline(always)]
     fn insert(&mut self, k: K, v: V) -> Option<V> {
         FxIndexMap::insert(self, k, v)

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -49,6 +49,14 @@ pub trait AllocMap<K: Hash + Eq, V> {
     where
         K: Borrow<Q>;
 
+    /// Callers should prefer [`AllocMap::contains_key`] when it is possible to call because it may
+    /// be more efficient. This function exists for callers that only have a shared reference
+    /// (which might make it slightly less efficient than `contains_key`, e.g. if
+    /// the data is stored inside a `RefCell`).
+    fn contains_key_ref<Q: ?Sized + Hash + Eq>(&self, k: &Q) -> bool
+    where
+        K: Borrow<Q>;
+
     /// Inserts a new entry into the map.
     fn insert(&mut self, k: K, v: V) -> Option<V>;
 

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -692,6 +692,15 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         Ok((&mut alloc.extra, machine))
     }
 
+    /// Check whether an allocation is live. This is faster than calling
+    /// [`InterpCx::get_alloc_info`] if all you need to check is whether the kind is
+    /// [`AllocKind::Dead`] because it doesn't have to look up the type and layout of statics.
+    pub fn is_alloc_live(&self, id: AllocId) -> bool {
+        self.tcx.try_get_global_alloc(id).is_some()
+            || self.memory.alloc_map.contains_key_ref(&id)
+            || self.memory.extra_fn_ptr_map.contains_key(&id)
+    }
+
     /// Obtain the size and alignment of an allocation, even if that allocation has
     /// been deallocated.
     pub fn get_alloc_info(&self, id: AllocId) -> (Size, Align, AllocKind) {

compiler/rustc_middle/src/mir/interpret/mod.rs

@@ -525,13 +525,6 @@ impl<'tcx> TyCtxt<'tcx> {
         self.alloc_map.lock().reserve()
     }
 
-    /// Miri's provenance GC needs to see all live allocations. The interpreter manages most
-    /// allocations but some are managed by [`TyCtxt`] and without this method the interpreter
-    /// doesn't know their [`AllocId`]s are in use.
-    pub fn iter_allocs<F: FnMut(AllocId)>(self, func: F) {
-        self.alloc_map.lock().alloc_map.keys().copied().for_each(func)
-    }
-
     /// Reserves a new ID *if* this allocation has not been dedup-reserved before.
     /// Should only be used for "symbolic" allocations (function pointers, vtables, statics), we
     /// don't want to dedup IDs for "real" memory!

src/tools/miri/src/borrow_tracker/mod.rs

@@ -75,8 +75,8 @@ pub struct FrameState {
     protected_tags: SmallVec<[(AllocId, BorTag); 2]>,
 }
 
-impl VisitTags for FrameState {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for FrameState {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
         // `protected_tags` are already recorded by `GlobalStateInner`.
     }
 }
@@ -110,10 +110,10 @@ pub struct GlobalStateInner {
     unique_is_unique: bool,
 }
 
-impl VisitTags for GlobalStateInner {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for GlobalStateInner {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         for &tag in self.protected_tags.keys() {
-            visit(tag);
+            visit(None, Some(tag));
         }
         // The only other candidate is base_ptr_tags, and that does not need visiting since we don't ever
         // GC the bottommost/root tag.
@@ -236,6 +236,10 @@ impl GlobalStateInner {
             tag
         })
     }
+
+    pub fn remove_unreachable_allocs(&mut self, allocs: &LiveAllocs<'_, '_, '_>) {
+        self.base_ptr_tags.retain(|id, _| allocs.is_live(*id));
+    }
 }
 
 /// Which borrow tracking method to use
@@ -503,11 +507,11 @@ impl AllocState {
     }
 }
 
-impl VisitTags for AllocState {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for AllocState {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         match self {
-            AllocState::StackedBorrows(sb) => sb.visit_tags(visit),
-            AllocState::TreeBorrows(tb) => tb.visit_tags(visit),
+            AllocState::StackedBorrows(sb) => sb.visit_provenance(visit),
+            AllocState::TreeBorrows(tb) => tb.visit_provenance(visit),
         }
     }
 }

src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs

@@ -462,10 +462,10 @@ impl Stacks {
     }
 }
 
-impl VisitTags for Stacks {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for Stacks {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         for tag in self.exposed_tags.iter().copied() {
-            visit(tag);
+            visit(None, Some(tag));
         }
     }
 }

src/tools/miri/src/borrow_tracker/tree_borrows/tree.rs

@@ -742,11 +742,11 @@ impl Tree {
     }
 }
 
-impl VisitTags for Tree {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for Tree {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         // To ensure that the root never gets removed, we visit it
         // (the `root` node of `Tree` is not an `Option<_>`)
-        visit(self.nodes.get(self.root).unwrap().tag)
+        visit(None, Some(self.nodes.get(self.root).unwrap().tag))
     }
 }
 

src/tools/miri/src/concurrency/data_race.rs

@@ -790,9 +790,9 @@ pub struct VClockAlloc {
     alloc_ranges: RefCell<RangeMap<MemoryCellClocks>>,
 }
 
-impl VisitTags for VClockAlloc {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
-        // No tags here.
+impl VisitProvenance for VClockAlloc {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
+        // No tags or allocIds here.
     }
 }
 
@@ -1404,8 +1404,8 @@ pub struct GlobalState {
     pub track_outdated_loads: bool,
 }
 
-impl VisitTags for GlobalState {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for GlobalState {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
         // We don't have any tags.
     }
 }

src/tools/miri/src/concurrency/init_once.rs

@@ -45,10 +45,10 @@ pub(super) struct InitOnce<'mir, 'tcx> {
     data_race: VClock,
 }
 
-impl<'mir, 'tcx> VisitTags for InitOnce<'mir, 'tcx> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl<'mir, 'tcx> VisitProvenance for InitOnce<'mir, 'tcx> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         for waiter in self.waiters.iter() {
-            waiter.callback.visit_tags(visit);
+            waiter.callback.visit_provenance(visit);
         }
     }
 }

src/tools/miri/src/concurrency/sync.rs

@@ -181,10 +181,10 @@ pub(crate) struct SynchronizationState<'mir, 'tcx> {
     pub(super) init_onces: IndexVec<InitOnceId, InitOnce<'mir, 'tcx>>,
 }
 
-impl<'mir, 'tcx> VisitTags for SynchronizationState<'mir, 'tcx> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl<'mir, 'tcx> VisitProvenance for SynchronizationState<'mir, 'tcx> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         for init_once in self.init_onces.iter() {
-            init_once.visit_tags(visit);
+            init_once.visit_provenance(visit);
         }
     }
 }

src/tools/miri/src/concurrency/thread.rs

@@ -43,7 +43,7 @@ pub enum TlsAllocAction {
 }
 
 /// Trait for callbacks that can be executed when some event happens, such as after a timeout.
-pub trait MachineCallback<'mir, 'tcx>: VisitTags {
+pub trait MachineCallback<'mir, 'tcx>: VisitProvenance {
     fn call(&self, ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>) -> InterpResult<'tcx>;
 }
 
@@ -228,8 +228,8 @@ impl<'mir, 'tcx> Thread<'mir, 'tcx> {
     }
 }
 
-impl VisitTags for Thread<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for Thread<'_, '_> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let Thread {
             panic_payloads: panic_payload,
             last_error,
@@ -242,17 +242,17 @@ impl VisitTags for Thread<'_, '_> {
         } = self;
 
         for payload in panic_payload {
-            payload.visit_tags(visit);
+            payload.visit_provenance(visit);
         }
-        last_error.visit_tags(visit);
+        last_error.visit_provenance(visit);
         for frame in stack {
-            frame.visit_tags(visit)
+            frame.visit_provenance(visit)
         }
     }
 }
 
-impl VisitTags for Frame<'_, '_, Provenance, FrameExtra<'_>> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for Frame<'_, '_, Provenance, FrameExtra<'_>> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let Frame {
             return_place,
             locals,
@@ -266,22 +266,22 @@ impl VisitTags for Frame<'_, '_, Provenance, FrameExtra<'_>> {
         } = self;
 
         // Return place.
-        return_place.visit_tags(visit);
+        return_place.visit_provenance(visit);
         // Locals.
         for local in locals.iter() {
             match local.as_mplace_or_imm() {
                 None => {}
                 Some(Either::Left((ptr, meta))) => {
-                    ptr.visit_tags(visit);
-                    meta.visit_tags(visit);
+                    ptr.visit_provenance(visit);
+                    meta.visit_provenance(visit);
                 }
                 Some(Either::Right(imm)) => {
-                    imm.visit_tags(visit);
+                    imm.visit_provenance(visit);
                 }
             }
         }
 
-        extra.visit_tags(visit);
+        extra.visit_provenance(visit);
     }
 }
 
@@ -341,8 +341,8 @@ pub struct ThreadManager<'mir, 'tcx> {
     timeout_callbacks: FxHashMap<ThreadId, TimeoutCallbackInfo<'mir, 'tcx>>,
 }
 
-impl VisitTags for ThreadManager<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for ThreadManager<'_, '_> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let ThreadManager {
             threads,
             thread_local_alloc_ids,
@@ -353,15 +353,15 @@ impl VisitTags for ThreadManager<'_, '_> {
         } = self;
 
         for thread in threads {
-            thread.visit_tags(visit);
+            thread.visit_provenance(visit);
         }
         for ptr in thread_local_alloc_ids.borrow().values() {
-            ptr.visit_tags(visit);
+            ptr.visit_provenance(visit);
         }
         for callback in timeout_callbacks.values() {
-            callback.callback.visit_tags(visit);
+            callback.callback.visit_provenance(visit);
         }
-        sync.visit_tags(visit);
+        sync.visit_provenance(visit);
     }
 }
 

src/tools/miri/src/concurrency/weak_memory.rs

@@ -108,15 +108,15 @@ pub struct StoreBufferAlloc {
     store_buffers: RefCell<RangeObjectMap<StoreBuffer>>,
 }
 
-impl VisitTags for StoreBufferAlloc {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for StoreBufferAlloc {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let Self { store_buffers } = self;
         for val in store_buffers
             .borrow()
             .iter()
             .flat_map(|buf| buf.buffer.iter().map(|element| &element.val))
         {
-            val.visit_tags(visit);
+            val.visit_provenance(visit);
         }
     }
 }

src/tools/miri/src/intptrcast.rs

@@ -46,9 +46,21 @@ pub struct GlobalStateInner {
     provenance_mode: ProvenanceMode,
 }
 
-impl VisitTags for GlobalStateInner {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
-        // Nothing to visit here.
+impl VisitProvenance for GlobalStateInner {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
+        let GlobalStateInner {
+            int_to_ptr_map: _,
+            base_addr: _,
+            exposed: _,
+            next_base_addr: _,
+            provenance_mode: _,
+        } = self;
+        // Though base_addr, int_to_ptr_map, and exposed contain AllocIds, we do not want to visit them.
+        // int_to_ptr_map and exposed must contain only live allocations, and those
+        // are never garbage collected.
+        // base_addr is only relevant if we have a pointer to an AllocId and need to look up its
+        // base address; so if an AllocId is not reachable from somewhere else we can remove it
+        // here.
     }
 }
 
@@ -62,6 +74,12 @@ impl GlobalStateInner {
             provenance_mode: config.provenance_mode,
         }
     }
+
+    pub fn remove_unreachable_allocs(&mut self, allocs: &LiveAllocs<'_, '_, '_>) {
+        // `exposed` and `int_to_ptr_map` are cleared immediately when an allocation
+        // is freed, so `base_addr` is the only one we have to clean up based on the GC.
+        self.base_addr.retain(|id, _| allocs.is_live(*id));
+    }
 }
 
 /// Shifts `addr` to make it aligned with `align` by rounding `addr` to the smallest multiple
@@ -107,7 +125,7 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // We only use this provenance if it has been exposed.
         if global_state.exposed.contains(&alloc_id) {
             // This must still be live, since we remove allocations from `int_to_ptr_map` when they get freed.
-            debug_assert!(!matches!(ecx.get_alloc_info(alloc_id).2, AllocKind::Dead));
+            debug_assert!(ecx.is_alloc_live(alloc_id));
             Some(alloc_id)
         } else {
             None
@@ -181,12 +199,19 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         let ecx = self.eval_context_mut();
         let global_state = ecx.machine.intptrcast.get_mut();
         // In strict mode, we don't need this, so we can save some cycles by not tracking it.
-        if global_state.provenance_mode != ProvenanceMode::Strict {
-            trace!("Exposing allocation id {alloc_id:?}");
-            global_state.exposed.insert(alloc_id);
-            if ecx.machine.borrow_tracker.is_some() {
-                ecx.expose_tag(alloc_id, tag)?;
-            }
+        if global_state.provenance_mode == ProvenanceMode::Strict {
+            return Ok(());
+        }
+        // Exposing a dead alloc is a no-op, because it's not possible to get a dead allocation
+        // via int2ptr.
+        if !ecx.is_alloc_live(alloc_id) {
+            return Ok(());
+        }
+        trace!("Exposing allocation id {alloc_id:?}");
+        let global_state = ecx.machine.intptrcast.get_mut();
+        global_state.exposed.insert(alloc_id);
+        if ecx.machine.borrow_tracker.is_some() {
+            ecx.expose_tag(alloc_id, tag)?;
         }
         Ok(())
     }