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 {
+impl VisitProvenance for FrameState {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    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 {
+impl VisitProvenance for GlobalStateInner {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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 {
+impl VisitProvenance for AllocState {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         match self {
-            AllocState::StackedBorrows(sb) => sb.visit_tags(visit),
+            AllocState::StackedBorrows(sb) => sb.visit_provenance(visit),
-            AllocState::TreeBorrows(tb) => tb.visit_tags(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 {
+impl VisitProvenance for Stacks {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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 {
+impl VisitProvenance for Tree {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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 {
+impl VisitProvenance for VClockAlloc {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
-        // No tags here.
+        // No tags or allocIds here.
     }
 }
 
@@ -1404,8 +1404,8 @@ pub struct GlobalState {
     pub track_outdated_loads: bool,
 }
 
-impl VisitTags for GlobalState {
+impl VisitProvenance for GlobalState {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    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> {
+impl<'mir, 'tcx> VisitProvenance for InitOnce<'mir, 'tcx> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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> {
+impl<'mir, 'tcx> VisitProvenance for SynchronizationState<'mir, 'tcx> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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<'_, '_> {
+impl VisitProvenance for Thread<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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<'_>> {
+impl VisitProvenance for Frame<'_, '_, Provenance, FrameExtra<'_>> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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);
+                    ptr.visit_provenance(visit);
-                    meta.visit_tags(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<'_, '_> {
+impl VisitProvenance for ThreadManager<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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 {
+impl VisitProvenance for StoreBufferAlloc {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    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 {
+impl VisitProvenance for GlobalStateInner {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
-        // Nothing to visit here.
+        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,13 +199,20 @@ 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 {
+        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(())
     }
 

src/tools/miri/src/lib.rs

@@ -77,9 +77,9 @@ mod intptrcast;
 mod machine;
 mod mono_hash_map;
 mod operator;
+mod provenance_gc;
 mod range_map;
 mod shims;
-mod tag_gc;
 
 // Establish a "crate-wide prelude": we often import `crate::*`.
 
@@ -125,8 +125,8 @@ pub use crate::machine::{
 };
 pub use crate::mono_hash_map::MonoHashMap;
 pub use crate::operator::EvalContextExt as _;
+pub use crate::provenance_gc::{EvalContextExt as _, VisitProvenance, VisitWith, LiveAllocs};
 pub use crate::range_map::RangeMap;
-pub use crate::tag_gc::{EvalContextExt as _, VisitTags};
 
 /// Insert rustc arguments at the beginning of the argument list that Miri wants to be
 /// set per default, for maximal validation power.

src/tools/miri/src/machine.rs

@@ -77,12 +77,12 @@ impl<'tcx> std::fmt::Debug for FrameExtra<'tcx> {
     }
 }
 
-impl VisitTags for FrameExtra<'_> {
+impl VisitProvenance for FrameExtra<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let FrameExtra { catch_unwind, borrow_tracker, timing: _, is_user_relevant: _ } = self;
 
-        catch_unwind.visit_tags(visit);
+        catch_unwind.visit_provenance(visit);
-        borrow_tracker.visit_tags(visit);
+        borrow_tracker.visit_provenance(visit);
     }
 }
 
@@ -311,13 +311,13 @@ pub struct AllocExtra<'tcx> {
     pub backtrace: Option<Vec<FrameInfo<'tcx>>>,
 }
 
-impl VisitTags for AllocExtra<'_> {
+impl VisitProvenance for AllocExtra<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let AllocExtra { borrow_tracker, data_race, weak_memory, backtrace: _ } = self;
 
-        borrow_tracker.visit_tags(visit);
+        borrow_tracker.visit_provenance(visit);
-        data_race.visit_tags(visit);
+        data_race.visit_provenance(visit);
-        weak_memory.visit_tags(visit);
+        weak_memory.visit_provenance(visit);
     }
 }
 
@@ -793,8 +793,8 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
     }
 }
 
-impl VisitTags for MiriMachine<'_, '_> {
+impl VisitProvenance for MiriMachine<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         #[rustfmt::skip]
         let MiriMachine {
             threads,
@@ -843,20 +843,20 @@ impl VisitTags for MiriMachine<'_, '_> {
             allocation_spans: _,
         } = self;
 
-        threads.visit_tags(visit);
+        threads.visit_provenance(visit);
-        tls.visit_tags(visit);
+        tls.visit_provenance(visit);
-        env_vars.visit_tags(visit);
+        env_vars.visit_provenance(visit);
-        dir_handler.visit_tags(visit);
+        dir_handler.visit_provenance(visit);
-        file_handler.visit_tags(visit);
+        file_handler.visit_provenance(visit);
-        data_race.visit_tags(visit);
+        data_race.visit_provenance(visit);
-        borrow_tracker.visit_tags(visit);
+        borrow_tracker.visit_provenance(visit);
-        intptrcast.visit_tags(visit);
+        intptrcast.visit_provenance(visit);
-        main_fn_ret_place.visit_tags(visit);
+        main_fn_ret_place.visit_provenance(visit);
-        argc.visit_tags(visit);
+        argc.visit_provenance(visit);
-        argv.visit_tags(visit);
+        argv.visit_provenance(visit);
-        cmd_line.visit_tags(visit);
+        cmd_line.visit_provenance(visit);
         for ptr in extern_statics.values() {
-            ptr.visit_tags(visit);
+            ptr.visit_provenance(visit);
         }
     }
 }
@@ -1380,7 +1380,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         // where it mistakenly removes an important tag become visible.
         if ecx.machine.gc_interval > 0 && ecx.machine.since_gc >= ecx.machine.gc_interval {
             ecx.machine.since_gc = 0;
-            ecx.garbage_collect_tags()?;
+            ecx.run_provenance_gc();
         }
 
         // These are our preemption points.

src/tools/miri/src/mono_hash_map.rs

@@ -46,6 +46,14 @@ impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
         self.0.get_mut().contains_key(k)
     }
 
+    #[inline(always)]
+    fn contains_key_ref<Q: ?Sized + Hash + Eq>(&self, k: &Q) -> bool
+    where
+        K: Borrow<Q>,
+    {
+        self.0.borrow().contains_key(k)
+    }
+
     #[inline(always)]
     fn insert(&mut self, k: K, v: V) -> Option<V> {
         self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)

src/tools/miri/src/provenance_gc.rs

@@ -0,0 +1,209 @@
+use either::Either;
+
+use rustc_data_structures::fx::FxHashSet;
+
+use crate::*;
+
+pub type VisitWith<'a> = dyn FnMut(Option<AllocId>, Option<BorTag>) + 'a;
+
+pub trait VisitProvenance {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>);
+}
+
+impl<T: VisitProvenance> VisitProvenance for Option<T> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        if let Some(x) = self {
+            x.visit_provenance(visit);
+        }
+    }
+}
+
+impl<T: VisitProvenance> VisitProvenance for std::cell::RefCell<T> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        self.borrow().visit_provenance(visit)
+    }
+}
+
+impl VisitProvenance for BorTag {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        visit(None, Some(*self))
+    }
+}
+
+impl VisitProvenance for AllocId {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        visit(Some(*self), None)
+    }
+}
+
+impl VisitProvenance for Provenance {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        if let Provenance::Concrete { alloc_id, tag, .. } = self {
+            visit(Some(*alloc_id), Some(*tag));
+        }
+    }
+}
+
+impl VisitProvenance for Pointer<Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        let (prov, _offset) = self.into_parts();
+        prov.visit_provenance(visit);
+    }
+}
+
+impl VisitProvenance for Pointer<Option<Provenance>> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        let (prov, _offset) = self.into_parts();
+        prov.visit_provenance(visit);
+    }
+}
+
+impl VisitProvenance for Scalar<Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        match self {
+            Scalar::Ptr(ptr, _) => ptr.visit_provenance(visit),
+            Scalar::Int(_) => (),
+        }
+    }
+}
+
+impl VisitProvenance for Immediate<Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        match self {
+            Immediate::Scalar(s) => {
+                s.visit_provenance(visit);
+            }
+            Immediate::ScalarPair(s1, s2) => {
+                s1.visit_provenance(visit);
+                s2.visit_provenance(visit);
+            }
+            Immediate::Uninit => {}
+        }
+    }
+}
+
+impl VisitProvenance for MemPlaceMeta<Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        match self {
+            MemPlaceMeta::Meta(m) => m.visit_provenance(visit),
+            MemPlaceMeta::None => {}
+        }
+    }
+}
+
+impl VisitProvenance for ImmTy<'_, Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        (**self).visit_provenance(visit)
+    }
+}
+
+impl VisitProvenance for MPlaceTy<'_, Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        self.ptr().visit_provenance(visit);
+        self.meta().visit_provenance(visit);
+    }
+}
+
+impl VisitProvenance for PlaceTy<'_, Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        match self.as_mplace_or_local() {
+            Either::Left(mplace) => mplace.visit_provenance(visit),
+            Either::Right(_) => (),
+        }
+    }
+}
+
+impl VisitProvenance for OpTy<'_, Provenance> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        match self.as_mplace_or_imm() {
+            Either::Left(mplace) => mplace.visit_provenance(visit),
+            Either::Right(imm) => imm.visit_provenance(visit),
+        }
+    }
+}
+
+impl VisitProvenance for Allocation<Provenance, AllocExtra<'_>> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        for prov in self.provenance().provenances() {
+            prov.visit_provenance(visit);
+        }
+
+        self.extra.visit_provenance(visit);
+    }
+}
+
+impl VisitProvenance for crate::MiriInterpCx<'_, '_> {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        // Visit the contents of the allocations and the IDs themselves, to account for all
+        // live allocation IDs and all provenance in the allocation bytes, even if they are leaked.
+        // We do *not* visit all the `AllocId` of the live allocations; we tried that and adding
+        // them all to the live set is too expensive. Instead we later do liveness check by
+        // checking both "is this alloc id live" and "is it mentioned anywhere else in
+        // the interpreter state".
+        self.memory.alloc_map().iter(|it| {
+            for (_id, (_kind, alloc)) in it {
+                alloc.visit_provenance(visit);
+            }
+        });
+        // And all the other machine values.
+        self.machine.visit_provenance(visit);
+    }
+}
+
+pub struct LiveAllocs<'a, 'mir, 'tcx> {
+    collected: FxHashSet<AllocId>,
+    ecx: &'a MiriInterpCx<'mir, 'tcx>,
+}
+
+impl LiveAllocs<'_, '_, '_> {
+    pub fn is_live(&self, id: AllocId) -> bool {
+        self.collected.contains(&id) ||
+        self.ecx.is_alloc_live(id)
+    }
+}
+
+impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
+pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
+    fn run_provenance_gc(&mut self) {
+
+        // We collect all tags from various parts of the interpreter, but also
+        let this = self.eval_context_mut();
+
+        let mut tags = FxHashSet::default();
+        let mut alloc_ids = FxHashSet::default();
+        this.visit_provenance(&mut |id, tag| {
+            if let Some(id) = id {
+                alloc_ids.insert(id);
+            }
+            if let Some(tag) = tag {
+                tags.insert(tag);
+            }
+        });
+        self.remove_unreachable_tags(tags);
+        self.remove_unreachable_allocs(alloc_ids);
+    }
+
+    fn remove_unreachable_tags(&mut self, tags: FxHashSet<BorTag>) {
+        let this = self.eval_context_mut();
+        this.memory.alloc_map().iter(|it| {
+            for (_id, (_kind, alloc)) in it {
+                if let Some(bt) = &alloc.extra.borrow_tracker {
+                    bt.remove_unreachable_tags(&tags);
+                }
+            }
+        });
+    }
+
+    fn remove_unreachable_allocs(&mut self, allocs: FxHashSet<AllocId>) {
+        let this = self.eval_context_ref();
+        let allocs = LiveAllocs {
+            ecx: this,
+            collected: allocs,
+        };
+        this.machine.allocation_spans.borrow_mut().retain(|id, _| allocs.is_live(*id));
+        this.machine.intptrcast.borrow_mut().remove_unreachable_allocs(&allocs);
+        if let Some(borrow_tracker) = &this.machine.borrow_tracker {
+            borrow_tracker.borrow_mut().remove_unreachable_allocs(&allocs);
+        }
+    }
+}

src/tools/miri/src/shims/env.rs

@@ -37,13 +37,13 @@ pub struct EnvVars<'tcx> {
     pub(crate) environ: Option<MPlaceTy<'tcx, Provenance>>,
 }
 
-impl VisitTags for EnvVars<'_> {
+impl VisitProvenance for EnvVars<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let EnvVars { map, environ } = self;
 
-        environ.visit_tags(visit);
+        environ.visit_provenance(visit);
         for ptr in map.values() {
-            ptr.visit_tags(visit);
+            ptr.visit_provenance(visit);
         }
     }
 }

src/tools/miri/src/shims/foreign_items.rs

@@ -459,6 +459,10 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // shim, add it to the corresponding submodule.
         match link_name.as_str() {
             // Miri-specific extern functions
+            "miri_run_provenance_gc" => {
+                let [] = this.check_shim(abi, Abi::Rust, link_name, args)?;
+                this.run_provenance_gc();
+            }
             "miri_get_alloc_id" => {
                 let [ptr] = this.check_shim(abi, Abi::Rust, link_name, args)?;
                 let ptr = this.read_pointer(ptr)?;

src/tools/miri/src/shims/panic.rs

@@ -35,12 +35,12 @@ pub struct CatchUnwindData<'tcx> {
     ret: mir::BasicBlock,
 }
 
-impl VisitTags for CatchUnwindData<'_> {
+impl VisitProvenance for CatchUnwindData<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let CatchUnwindData { catch_fn, data, dest, ret: _ } = self;
-        catch_fn.visit_tags(visit);
+        catch_fn.visit_provenance(visit);
-        data.visit_tags(visit);
+        data.visit_provenance(visit);
-        dest.visit_tags(visit);
+        dest.visit_provenance(visit);
     }
 }
 

src/tools/miri/src/shims/time.rs

@@ -274,8 +274,8 @@ struct UnblockCallback {
     thread_to_unblock: ThreadId,
 }
 
-impl VisitTags for UnblockCallback {
+impl VisitProvenance for UnblockCallback {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {}
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {}
 }
 
 impl<'mir, 'tcx: 'mir> MachineCallback<'mir, 'tcx> for UnblockCallback {

src/tools/miri/src/shims/tls.rs

@@ -207,15 +207,15 @@ impl<'tcx> TlsData<'tcx> {
     }
 }
 
-impl VisitTags for TlsData<'_> {
+impl VisitProvenance for TlsData<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let TlsData { keys, macos_thread_dtors, next_key: _ } = self;
 
         for scalar in keys.values().flat_map(|v| v.data.values()) {
-            scalar.visit_tags(visit);
+            scalar.visit_provenance(visit);
         }
         for (_, scalar) in macos_thread_dtors.values() {
-            scalar.visit_tags(visit);
+            scalar.visit_provenance(visit);
         }
     }
 }

src/tools/miri/src/shims/unix/fs.rs

@@ -288,8 +288,8 @@ pub struct FileHandler {
     pub handles: BTreeMap<i32, Box<dyn FileDescriptor>>,
 }
 
-impl VisitTags for FileHandler {
+impl VisitProvenance for FileHandler {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
         // All our FileDescriptor do not have any tags.
     }
 }
@@ -490,12 +490,12 @@ impl Default for DirHandler {
     }
 }
 
-impl VisitTags for DirHandler {
+impl VisitProvenance for DirHandler {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
         let DirHandler { streams, next_id: _ } = self;
 
         for dir in streams.values() {
-            dir.entry.visit_tags(visit);
+            dir.entry.visit_provenance(visit);
         }
     }
 }

src/tools/miri/src/shims/unix/linux/sync.rs

@@ -182,10 +182,10 @@ pub fn futex<'tcx>(
                         dest: PlaceTy<'tcx, Provenance>,
                     }
 
-                    impl<'tcx> VisitTags for Callback<'tcx> {
+                    impl<'tcx> VisitProvenance for Callback<'tcx> {
-                        fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+                        fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
                             let Callback { thread: _, addr_usize: _, dest } = self;
-                            dest.visit_tags(visit);
+                            dest.visit_provenance(visit);
                         }
                     }
 

src/tools/miri/src/shims/unix/sync.rs

@@ -886,10 +886,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             dest: PlaceTy<'tcx, Provenance>,
         }
 
-        impl<'tcx> VisitTags for Callback<'tcx> {
+        impl<'tcx> VisitProvenance for Callback<'tcx> {
-            fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+            fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
                 let Callback { active_thread: _, mutex_id: _, id: _, dest } = self;
-                dest.visit_tags(visit);
+                dest.visit_provenance(visit);
             }
         }
 

src/tools/miri/src/shims/windows/sync.rs

@@ -204,10 +204,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     pending_place: PlaceTy<'tcx, Provenance>,
                 }
 
-                impl<'tcx> VisitTags for Callback<'tcx> {
+                impl<'tcx> VisitProvenance for Callback<'tcx> {
-                    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+                    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
                         let Callback { init_once_id: _, pending_place } = self;
-                        pending_place.visit_tags(visit);
+                        pending_place.visit_provenance(visit);
                     }
                 }
 
@@ -337,10 +337,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     dest: PlaceTy<'tcx, Provenance>,
                 }
 
-                impl<'tcx> VisitTags for Callback<'tcx> {
+                impl<'tcx> VisitProvenance for Callback<'tcx> {
-                    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+                    fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
                         let Callback { thread: _, addr: _, dest } = self;
-                        dest.visit_tags(visit);
+                        dest.visit_provenance(visit);
                     }
                 }
 
@@ -441,10 +441,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 dest: PlaceTy<'tcx, Provenance>,
             }
 
-            impl<'tcx> VisitTags for Callback<'tcx> {
+            impl<'tcx> VisitProvenance for Callback<'tcx> {
-                fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+                fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
                     let Callback { thread: _, condvar_id: _, lock_id: _, mode: _, dest } = self;
-                    dest.visit_tags(visit);
+                    dest.visit_provenance(visit);
                 }
             }
 

src/tools/miri/src/tag_gc.rs

@@ -1,169 +0,0 @@
-use either::Either;
-
-use rustc_data_structures::fx::FxHashSet;
-
-use crate::*;
-
-pub trait VisitTags {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag));
-}
-
-impl<T: VisitTags> VisitTags for Option<T> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        if let Some(x) = self {
-            x.visit_tags(visit);
-        }
-    }
-}
-
-impl<T: VisitTags> VisitTags for std::cell::RefCell<T> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        self.borrow().visit_tags(visit)
-    }
-}
-
-impl VisitTags for BorTag {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        visit(*self)
-    }
-}
-
-impl VisitTags for Provenance {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        if let Provenance::Concrete { tag, .. } = self {
-            visit(*tag);
-        }
-    }
-}
-
-impl VisitTags for Pointer<Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        let (prov, _offset) = self.into_parts();
-        prov.visit_tags(visit);
-    }
-}
-
-impl VisitTags for Pointer<Option<Provenance>> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        let (prov, _offset) = self.into_parts();
-        prov.visit_tags(visit);
-    }
-}
-
-impl VisitTags for Scalar<Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        match self {
-            Scalar::Ptr(ptr, _) => ptr.visit_tags(visit),
-            Scalar::Int(_) => (),
-        }
-    }
-}
-
-impl VisitTags for Immediate<Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        match self {
-            Immediate::Scalar(s) => {
-                s.visit_tags(visit);
-            }
-            Immediate::ScalarPair(s1, s2) => {
-                s1.visit_tags(visit);
-                s2.visit_tags(visit);
-            }
-            Immediate::Uninit => {}
-        }
-    }
-}
-
-impl VisitTags for MemPlaceMeta<Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        match self {
-            MemPlaceMeta::Meta(m) => m.visit_tags(visit),
-            MemPlaceMeta::None => {}
-        }
-    }
-}
-
-impl VisitTags for ImmTy<'_, Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        (**self).visit_tags(visit)
-    }
-}
-
-impl VisitTags for MPlaceTy<'_, Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        self.ptr().visit_tags(visit);
-        self.meta().visit_tags(visit);
-    }
-}
-
-impl VisitTags for PlaceTy<'_, Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        match self.as_mplace_or_local() {
-            Either::Left(mplace) => mplace.visit_tags(visit),
-            Either::Right(_) => (),
-        }
-    }
-}
-
-impl VisitTags for OpTy<'_, Provenance> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        match self.as_mplace_or_imm() {
-            Either::Left(mplace) => mplace.visit_tags(visit),
-            Either::Right(imm) => imm.visit_tags(visit),
-        }
-    }
-}
-
-impl VisitTags for Allocation<Provenance, AllocExtra<'_>> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        for prov in self.provenance().provenances() {
-            prov.visit_tags(visit);
-        }
-
-        self.extra.visit_tags(visit);
-    }
-}
-
-impl VisitTags for crate::MiriInterpCx<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
-        // Memory.
-        self.memory.alloc_map().iter(|it| {
-            for (_id, (_kind, alloc)) in it {
-                alloc.visit_tags(visit);
-            }
-        });
-
-        // And all the other machine values.
-        self.machine.visit_tags(visit);
-    }
-}
-
-impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
-pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
-    fn garbage_collect_tags(&mut self) -> InterpResult<'tcx> {
-        let this = self.eval_context_mut();
-        // No reason to do anything at all if stacked borrows is off.
-        if this.machine.borrow_tracker.is_none() {
-            return Ok(());
-        }
-
-        let mut tags = FxHashSet::default();
-        this.visit_tags(&mut |tag| {
-            tags.insert(tag);
-        });
-        self.remove_unreachable_tags(tags);
-
-        Ok(())
-    }
-
-    fn remove_unreachable_tags(&mut self, tags: FxHashSet<BorTag>) {
-        let this = self.eval_context_mut();
-        this.memory.alloc_map().iter(|it| {
-            for (_id, (_kind, alloc)) in it {
-                if let Some(bt) = &alloc.extra.borrow_tracker {
-                    bt.remove_unreachable_tags(&tags);
-                }
-            }
-        });
-    }
-}