Use VisitProvenance to factor allocation visiting better

4 files changed, 32 insertions, 18 deletions

diff --git a/src/tools/miri/src/concurrency/weak_memory.rs b/src/tools/miri/src/concurrency/weak_memory.rs
index aa4e5add503..becd61f4fea 100644
--- a/src/tools/miri/src/concurrency/weak_memory.rs
+++ b/src/tools/miri/src/concurrency/weak_memory.rs
@@ -108,15 +108,19 @@ pub struct StoreBufferAlloc {
     store_buffers: RefCell<RangeObjectMap<StoreBuffer>>,
 }
 
-impl StoreBufferAlloc {
-    pub fn iter(&self, mut visitor: impl FnMut(&Scalar<Provenance>)) {
+impl VisitProvenance for StoreBufferAlloc {
+    fn visit_provenance(&self, visitor: &mut impl FnMut(SbTag)) {
         for val in self
             .store_buffers
             .borrow()
             .iter()
             .flat_map(|buf| buf.buffer.iter().map(|element| &element.val))
         {
-            visitor(val)
+            if let Scalar::Ptr(ptr, _) = val {
+                if let Provenance::Concrete { sb, .. } = ptr.provenance {
+                    visitor(sb);
+                }
+            }
         }
     }
 }
diff --git a/src/tools/miri/src/lib.rs b/src/tools/miri/src/lib.rs
index bf4e21319b5..e60e1f15b5f 100644
--- a/src/tools/miri/src/lib.rs
+++ b/src/tools/miri/src/lib.rs
@@ -112,7 +112,7 @@ pub use crate::range_map::RangeMap;
 pub use crate::stacked_borrows::{
     CallId, EvalContextExt as StackedBorEvalContextExt, Item, Permission, SbTag, Stack, Stacks,
 };
-pub use crate::tag_gc::{EvalContextExt as _, VisitMachineValues};
+pub use crate::tag_gc::{EvalContextExt as _, VisitMachineValues, VisitProvenance};
 
 /// Insert rustc arguments at the beginning of the argument list that Miri wants to be
 /// set per default, for maximal validation power.
diff --git a/src/tools/miri/src/stacked_borrows/mod.rs b/src/tools/miri/src/stacked_borrows/mod.rs
index 956f936dbca..ab90e358449 100644
--- a/src/tools/miri/src/stacked_borrows/mod.rs
+++ b/src/tools/miri/src/stacked_borrows/mod.rs
@@ -79,7 +79,7 @@ pub struct Stacks {
     /// Stores past operations on this allocation
     history: AllocHistory,
     /// The set of tags that have been exposed inside this allocation.
-    pub exposed_tags: FxHashSet<SbTag>,
+    exposed_tags: FxHashSet<SbTag>,
     /// Whether this memory has been modified since the last time the tag GC ran
     modified_since_last_gc: bool,
 }
@@ -513,6 +513,14 @@ impl Stacks {
     }
 }
 
+impl VisitProvenance for Stacks {
+    fn visit_provenance(&self, visit: &mut impl FnMut(SbTag)) {
+        for tag in self.exposed_tags.iter().copied() {
+            visit(tag);
+        }
+    }
+}
+
 /// Map per-stack operations to higher-level per-location-range operations.
 impl<'tcx> Stacks {
     /// Creates a new stack with an initial tag. For diagnostic purposes, we also need to know
diff --git a/src/tools/miri/src/tag_gc.rs b/src/tools/miri/src/tag_gc.rs
index 9378cc1e920..0a8d5d00cfb 100644
--- a/src/tools/miri/src/tag_gc.rs
+++ b/src/tools/miri/src/tag_gc.rs
@@ -6,6 +6,10 @@ pub trait VisitMachineValues {
     fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>));
 }
 
+pub trait VisitProvenance {
+    fn visit_provenance(&self, visit: &mut impl FnMut(SbTag));
+}
+
 impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
 pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
     /// Generic GC helper to visit everything that can store a value. The `acc` offers some chance to
@@ -46,6 +50,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
             }
         };
 
+        let visit_provenance = |tags: &mut FxHashSet<SbTag>, tag: SbTag| {
+            tags.insert(tag);
+        };
+
         this.visit_all_machine_values(
             &mut tags,
             |tags, op| {
@@ -71,21 +79,15 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
                         tags.insert(*sb);
                     }
                 }
-                let stacks = alloc
-                    .extra
-                    .stacked_borrows
-                    .as_ref()
-                    .expect("we should not even enter the GC if Stacked Borrows is disabled");
-                tags.extend(&stacks.borrow().exposed_tags);
+
+                let stacks =
+                    alloc.extra.stacked_borrows.as_ref().expect(
+                        "we should not even enter the tag GC if Stacked Borrows is disabled",
+                    );
+                stacks.borrow().visit_provenance(&mut |tag| visit_provenance(tags, tag));
 
                 if let Some(store_buffers) = alloc.extra.weak_memory.as_ref() {
-                    store_buffers.iter(|val| {
-                        if let Scalar::Ptr(ptr, _) = val {
-                            if let Provenance::Concrete { sb, .. } = ptr.provenance {
-                                tags.insert(sb);
-                            }
-                        }
-                    });
+                    store_buffers.visit_provenance(&mut |tag| visit_provenance(tags, tag));
                 }
             },
         );