Auto merge of rust-lang#2559 - RalfJung:gc-refactor, r=saethlin

GC: factor out visiting all machine values

`@saethlin` that is roughly what I had in mind.

I think some parts of the state are skipped by the visitor. I listed the ones that I found in FIXMEs but I am not sure if that list is complete.

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

@@ -290,10 +290,6 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
         &mut self.threads[self.active_thread].stack
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = &Thread<'mir, 'tcx>> {
-        self.threads.iter()
-    }
-
     pub fn all_stacks(
         &self,
     ) -> impl Iterator<Item = &[Frame<'mir, 'tcx, Provenance, FrameData<'tcx>>]> {
@@ -628,6 +624,33 @@ impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
     }
 }
 
+impl VisitMachineValues for ThreadManager<'_, '_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        // FIXME some other fields also contain machine values
+        let ThreadManager { threads, .. } = self;
+
+        for thread in threads {
+            // FIXME: implement VisitMachineValues for `Thread` and `Frame` instead.
+            // In particular we need to visit the `last_error` and `catch_unwind` fields.
+            if let Some(payload) = thread.panic_payload {
+                visit(&Operand::Immediate(Immediate::Scalar(payload)))
+            }
+            for frame in &thread.stack {
+                // Return place.
+                if let Place::Ptr(mplace) = *frame.return_place {
+                    visit(&Operand::Indirect(mplace));
+                }
+                // Locals.
+                for local in frame.locals.iter() {
+                    if let LocalValue::Live(value) = &local.value {
+                        visit(value);
+                    }
+                }
+            }
+        }
+    }
+}
+
 // Public interface to thread management.
 impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
 pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {

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 _;
+pub use crate::tag_gc::{EvalContextExt as _, VisitMachineValues};
 
 /// 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

@@ -291,6 +291,9 @@ impl<'mir, 'tcx: 'mir> PrimitiveLayouts<'tcx> {
 }
 
 /// The machine itself.
+///
+/// If you add anything here that stores machine values, remember to update
+/// `visit_all_machine_values`!
 pub struct MiriMachine<'mir, 'tcx> {
     // We carry a copy of the global `TyCtxt` for convenience, so methods taking just `&Evaluator` have `tcx` access.
     pub tcx: TyCtxt<'tcx>,
@@ -586,6 +589,17 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
     }
 }
 
+impl VisitMachineValues for MiriMachine<'_, '_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        // FIXME: visit the missing fields: env vars, weak mem, the MemPlace fields in the machine,
+        // DirHandler, extern_statics, the Stacked Borrows base pointers; maybe more.
+        let MiriMachine { threads, tls, .. } = self;
+
+        threads.visit_machine_values(visit);
+        tls.visit_machine_values(visit);
+    }
+}
+
 /// A rustc InterpCx for Miri.
 pub type MiriInterpCx<'mir, 'tcx> = InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>;
 

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

@@ -233,10 +233,17 @@ impl<'tcx> TlsData<'tcx> {
             data.remove(&thread_id);
         }
     }
+}
+
+impl VisitMachineValues for TlsData<'_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        let TlsData { keys, macos_thread_dtors, next_key: _, dtors_running: _ } = self;
 
-    pub fn iter(&self, mut visitor: impl FnMut(&Scalar<Provenance>)) {
-        for scalar in self.keys.values().flat_map(|v| v.data.values()) {
-            visitor(scalar);
+        for scalar in keys.values().flat_map(|v| v.data.values()) {
+            visit(&Operand::Immediate(Immediate::Scalar(*scalar)));
+        }
+        for (_, scalar) in macos_thread_dtors.values() {
+            visit(&Operand::Immediate(Immediate::Scalar(*scalar)));
         }
     }
 }

src/tools/miri/src/tag_gc.rs

@@ -1,8 +1,34 @@
-use crate::*;
 use rustc_data_structures::fx::FxHashSet;
 
+use crate::*;
+
+pub trait VisitMachineValues {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>));
+}
+
 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
+    /// accumulate everything.
+    fn visit_all_machine_values<T>(
+        &self,
+        acc: &mut T,
+        mut visit_operand: impl FnMut(&mut T, &Operand<Provenance>),
+        mut visit_alloc: impl FnMut(&mut T, &Allocation<Provenance, AllocExtra>),
+    ) {
+        let this = self.eval_context_ref();
+
+        // Memory.
+        this.memory.alloc_map().iter(|it| {
+            for (_id, (_kind, alloc)) in it {
+                visit_alloc(acc, alloc);
+            }
+        });
+
+        // And all the other machine values.
+        this.machine.visit_machine_values(&mut |op| visit_operand(acc, op));
+    }
+
     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.
@@ -12,90 +38,43 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
 
         let mut tags = FxHashSet::default();
 
-        for thread in this.machine.threads.iter() {
-            if let Some(Scalar::Ptr(
-                Pointer { provenance: Provenance::Concrete { sb, .. }, .. },
-                _,
-            )) = thread.panic_payload
-            {
-                tags.insert(sb);
-            }
-        }
-
-        self.find_tags_in_tls(&mut tags);
-        self.find_tags_in_memory(&mut tags);
-        self.find_tags_in_locals(&mut tags)?;
-
-        self.remove_unreachable_tags(tags);
-
-        Ok(())
-    }
-
-    fn find_tags_in_tls(&mut self, tags: &mut FxHashSet<SbTag>) {
-        let this = self.eval_context_mut();
-        this.machine.tls.iter(|scalar| {
-            if let Scalar::Ptr(Pointer { provenance: Provenance::Concrete { sb, .. }, .. }, _) =
-                scalar
-            {
-                tags.insert(*sb);
-            }
-        });
-    }
-
-    fn find_tags_in_memory(&mut self, tags: &mut FxHashSet<SbTag>) {
-        let this = self.eval_context_mut();
-        this.memory.alloc_map().iter(|it| {
-            for (_id, (_kind, alloc)) in it {
-                for (_size, prov) in alloc.provenance().iter() {
-                    if let Provenance::Concrete { sb, .. } = prov {
-                        tags.insert(*sb);
-                    }
-                }
-            }
-        });
-    }
-
-    fn find_tags_in_locals(&mut self, tags: &mut FxHashSet<SbTag>) -> InterpResult<'tcx> {
-        let this = self.eval_context_mut();
-        for frame in this.machine.threads.all_stacks().flatten() {
-            // Handle the return place of each frame
-            if let Ok(return_place) = frame.return_place.try_as_mplace() {
-                if let Some(Provenance::Concrete { sb, .. }) = return_place.ptr.provenance {
+        let visit_scalar = |tags: &mut FxHashSet<SbTag>, s: &Scalar<Provenance>| {
+            if let Scalar::Ptr(ptr, _) = s {
+                if let Provenance::Concrete { sb, .. } = ptr.provenance {
                     tags.insert(sb);
                 }
             }
+        };
 
-            for local in frame.locals.iter() {
-                let LocalValue::Live(value) = local.value else {
-                continue;
-            };
-                match value {
-                    Operand::Immediate(Immediate::Scalar(Scalar::Ptr(ptr, _))) =>
-                        if let Provenance::Concrete { sb, .. } = ptr.provenance {
-                            tags.insert(sb);
-                        },
+        this.visit_all_machine_values(
+            &mut tags,
+            |tags, op| {
+                match op {
+                    Operand::Immediate(Immediate::Scalar(s)) => {
+                        visit_scalar(tags, s);
+                    }
                     Operand::Immediate(Immediate::ScalarPair(s1, s2)) => {
-                        if let Scalar::Ptr(ptr, _) = s1 {
-                            if let Provenance::Concrete { sb, .. } = ptr.provenance {
-                                tags.insert(sb);
-                            }
-                        }
-                        if let Scalar::Ptr(ptr, _) = s2 {
-                            if let Provenance::Concrete { sb, .. } = ptr.provenance {
-                                tags.insert(sb);
-                            }
-                        }
+                        visit_scalar(tags, s1);
+                        visit_scalar(tags, s2);
                     }
+                    Operand::Immediate(Immediate::Uninit) => {}
                     Operand::Indirect(MemPlace { ptr, .. }) => {
                         if let Some(Provenance::Concrete { sb, .. }) = ptr.provenance {
                             tags.insert(sb);
                         }
                     }
-                    Operand::Immediate(Immediate::Uninit)
-                    | Operand::Immediate(Immediate::Scalar(Scalar::Int(_))) => {}
                 }
-            }
-        }
+            },
+            |tags, alloc| {
+                for (_size, prov) in alloc.provenance().iter() {
+                    if let Provenance::Concrete { sb, .. } = prov {
+                        tags.insert(*sb);
+                    }
+                }
+            },
+        );
+
+        self.remove_unreachable_tags(tags);
 
         Ok(())
     }