Extend the BorTag GC to AllocIds

Author: saethlin

src/borrow_tracker/mod.rs

@@ -76,7 +76,7 @@ pub struct FrameState {
 }
 
 impl VisitTags for FrameState {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, _visit: &mut TagVisitor<'_>) {
         // `protected_tags` are already recorded by `GlobalStateInner`.
     }
 }
@@ -111,9 +111,9 @@ pub struct GlobalStateInner {
 }
 
 impl VisitTags for GlobalStateInner {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         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.
@@ -472,7 +472,7 @@ impl AllocState {
 }
 
 impl VisitTags for AllocState {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         match self {
             AllocState::StackedBorrows(sb) => sb.visit_tags(visit),
             AllocState::TreeBorrows(tb) => tb.visit_tags(visit),

src/borrow_tracker/stacked_borrows/mod.rs

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

src/borrow_tracker/tree_borrows/tree.rs

@@ -745,10 +745,10 @@ impl Tree {
 }
 
 impl VisitTags for Tree {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         // 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

@@ -694,8 +694,8 @@ pub struct VClockAlloc {
 }
 
 impl VisitTags for VClockAlloc {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
-        // No tags here.
+    fn visit_tags(&self, _visit: &mut TagVisitor<'_>) {
+        // No tags or allocIds here.
     }
 }
 
@@ -1275,7 +1275,7 @@ pub struct GlobalState {
 }
 
 impl VisitTags for GlobalState {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, _visit: &mut TagVisitor<'_>) {
         // We don't have any tags.
     }
 }

src/concurrency/init_once.rs

@@ -46,7 +46,7 @@ pub(super) struct InitOnce<'mir, 'tcx> {
 }
 
 impl<'mir, 'tcx> VisitTags for InitOnce<'mir, 'tcx> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         for waiter in self.waiters.iter() {
             waiter.callback.visit_tags(visit);
         }

src/concurrency/sync.rs

@@ -182,7 +182,7 @@ pub(crate) struct SynchronizationState<'mir, 'tcx> {
 }
 
 impl<'mir, 'tcx> VisitTags for SynchronizationState<'mir, 'tcx> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         for init_once in self.init_onces.iter() {
             init_once.visit_tags(visit);
         }

src/concurrency/thread.rs

@@ -220,7 +220,7 @@ impl<'mir, 'tcx> Thread<'mir, 'tcx> {
 }
 
 impl VisitTags for Thread<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         let Thread {
             panic_payloads: panic_payload,
             last_error,
@@ -243,7 +243,7 @@ impl VisitTags for Thread<'_, '_> {
 }
 
 impl VisitTags for Frame<'_, '_, Provenance, FrameExtra<'_>> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         let Frame {
             return_place,
             locals,
@@ -333,7 +333,7 @@ pub struct ThreadManager<'mir, 'tcx> {
 }
 
 impl VisitTags for ThreadManager<'_, '_> {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         let ThreadManager {
             threads,
             thread_local_alloc_ids,

src/concurrency/weak_memory.rs

@@ -109,7 +109,7 @@ pub struct StoreBufferAlloc {
 }
 
 impl VisitTags for StoreBufferAlloc {
-    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
         let Self { store_buffers } = self;
         for val in store_buffers
             .borrow()

src/intptrcast.rs

@@ -45,8 +45,10 @@ pub struct GlobalStateInner {
 }
 
 impl VisitTags for GlobalStateInner {
-    fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
-        // Nothing to visit here.
+    fn visit_tags(&self, visit: &mut TagVisitor<'_>) {
+        for id in &self.exposed {
+            id.visit_tags(visit)
+        }
     }
 }
 
@@ -63,6 +65,11 @@ impl GlobalStateInner {
 }
 
 impl<'mir, 'tcx> GlobalStateInner {
+    pub fn remove_unreachable_allocs(&mut self, live_allocs: &FxHashSet<AllocId>) {
+        self.int_to_ptr_map.retain(|(_, id)| live_allocs.contains(id));
+        self.base_addr.retain(|id, _| live_allocs.contains(id));
+    }
+
     // Returns the exposed `AllocId` that corresponds to the specified addr,
     // or `None` if the addr is out of bounds
     fn alloc_id_from_addr(ecx: &MiriInterpCx<'mir, 'tcx>, addr: u64) -> Option<AllocId> {

src/lib.rs

@@ -125,7 +125,7 @@ pub use crate::machine::{
 pub use crate::mono_hash_map::MonoHashMap;
 pub use crate::operator::EvalContextExt as _;
 pub use crate::range_map::RangeMap;
-pub use crate::tag_gc::{EvalContextExt as _, VisitTags};
+pub use crate::tag_gc::{EvalContextExt as _, TagVisitor, VisitTags};
 
 /// Insert rustc arguments at the beginning of the argument list that Miri wants to be
 /// set per default, for maximal validation power.