Extend the BorTag GC to AllocIds

Author: saethlin

src/borrow_tracker/mod.rs

@@ -75,8 +75,8 @@ pub struct FrameState {
     pub 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_tags(&self, _visit: &mut VisitWith<'_>) {
         // `protected_tags` are already recorded by `GlobalStateInner`.
     }
 }
@@ -110,10 +110,10 @@ pub struct GlobalStateInner {
     pub unique_is_unique: bool,
 }
 
-impl VisitTags for GlobalStateInner {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for GlobalStateInner {
+    fn visit_tags(&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.
@@ -471,8 +471,8 @@ impl AllocState {
     }
 }
 
-impl VisitTags for AllocState {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for AllocState {
+    fn visit_tags(&self, visit: &mut VisitWith<'_>) {
         match self {
             AllocState::StackedBorrows(sb) => sb.visit_tags(visit),
             AllocState::TreeBorrows(tb) => tb.visit_tags(visit),

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_tags(&self, visit: &mut VisitWith<'_>) {
         for tag in self.exposed_tags.iter().copied() {
-            visit(tag);
+            visit(None, Some(tag));
         }
     }
 }

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_tags(&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/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_tags(&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_tags(&self, _visit: &mut VisitWith<'_>) {
         // We don't have any tags.
     }
 }

src/concurrency/init_once.rs

@@ -45,8 +45,8 @@ 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_tags(&self, visit: &mut VisitWith<'_>) {
         for waiter in self.waiters.iter() {
             waiter.callback.visit_tags(visit);
         }

src/concurrency/sync.rs

@@ -181,8 +181,8 @@ 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_tags(&self, visit: &mut VisitWith<'_>) {
         for init_once in self.init_onces.iter() {
             init_once.visit_tags(visit);
         }

src/concurrency/thread.rs

@@ -34,7 +34,7 @@ enum SchedulingAction {
 }
 
 /// 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>;
 }
 
@@ -219,8 +219,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_tags(&self, visit: &mut VisitWith<'_>) {
         let Thread {
             panic_payloads: panic_payload,
             last_error,
@@ -242,8 +242,8 @@ impl VisitTags for Thread<'_, '_> {
     }
 }
 
-impl VisitTags for Frame<'_, '_, Provenance, FrameExtra<'_>> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for Frame<'_, '_, Provenance, FrameExtra<'_>> {
+    fn visit_tags(&self, visit: &mut VisitWith<'_>) {
         let Frame {
             return_place,
             locals,
@@ -332,8 +332,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_tags(&self, visit: &mut VisitWith<'_>) {
         let ThreadManager {
             threads,
             thread_local_alloc_ids,

src/concurrency/weak_memory.rs

@@ -108,8 +108,8 @@ 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_tags(&self, visit: &mut VisitWith<'_>) {
         let Self { store_buffers } = self;
         for val in store_buffers
             .borrow()

src/intptrcast.rs

@@ -46,9 +46,23 @@ 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_tags(&self, visit: &mut VisitWith<'_>) {
+        let GlobalStateInner {
+            int_to_ptr_map: _,
+            base_addr: _,
+            exposed,
+            next_base_addr: _,
+            provenance_mode: _,
+        } = self;
+        // Though base_addr and int_to_ptr_map contain AllocIds, we do not want to visit them.
+        // We're trying to remove unused elements from base_addr, so visiting it would prevent
+        // removing anything. int_to_ptr_map is managed by free_alloc_id, and entries in it do not
+        // actually make allocation base addresses reachable so we don't need to visit it.
+        // But exposed tags do make base addresses reachable.
+        for id in exposed {
+            id.visit_tags(visit)
+        }
     }
 }
 
@@ -62,6 +76,10 @@ impl GlobalStateInner {
             provenance_mode: config.provenance_mode,
         }
     }
+
+    pub fn remove_unreachable_allocs(&mut self, reachable_allocs: &FxHashSet<AllocId>) {
+        self.base_addr.retain(|id, _| reachable_allocs.contains(id));
+    }
 }
 
 /// Shifts `addr` to make it aligned with `align` by rounding `addr` to the smallest multiple

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};
 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/machine.rs

@@ -77,8 +77,8 @@ impl<'tcx> std::fmt::Debug for FrameExtra<'tcx> {
     }
 }
 
-impl VisitTags for FrameExtra<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for FrameExtra<'_> {
+    fn visit_tags(&self, visit: &mut VisitWith<'_>) {
         let FrameExtra { catch_unwind, borrow_tracker, timing: _, is_user_relevant: _ } = self;
 
         catch_unwind.visit_tags(visit);
@@ -311,8 +311,8 @@ pub struct AllocExtra<'tcx> {
     pub backtrace: Option<Vec<FrameInfo<'tcx>>>,
 }
 
-impl VisitTags for AllocExtra<'_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for AllocExtra<'_> {
+    fn visit_tags(&self, visit: &mut VisitWith<'_>) {
         let AllocExtra { borrow_tracker, data_race, weak_memory, backtrace: _ } = self;
 
         borrow_tracker.visit_tags(visit);
@@ -793,8 +793,8 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
     }
 }
 
-impl VisitTags for MiriMachine<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+impl VisitProvenance for MiriMachine<'_, '_> {
+    fn visit_tags(&self, visit: &mut VisitWith<'_>) {
         #[rustfmt::skip]
         let MiriMachine {
             threads,
@@ -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.