Expand VisitMachineValues to cover more pointers in the interpreter

src/concurrency/data_race.rs

@@ -696,6 +696,12 @@ pub struct VClockAlloc {
     alloc_ranges: RefCell<RangeMap<MemoryCellClocks>>,
 }
 
+impl VisitTags for VClockAlloc {
+    fn visit_tags(&self, _visit: &mut dyn FnMut(SbTag)) {
+        // No tags here.
+    }
+}
+
 impl VClockAlloc {
     /// Create a new data-race detector for newly allocated memory.
     pub fn new_allocation(
@@ -1239,6 +1245,12 @@ pub struct GlobalState {
     pub track_outdated_loads: bool,
 }
 
+impl VisitTags for GlobalState {
+    fn visit_tags(&self, _visit: &mut dyn FnMut(SbTag)) {
+        // We don't have any tags.
+    }
+}
+
 impl GlobalState {
     /// Create a new global state, setup with just thread-id=0
     /// advanced to timestamp = 1.

src/concurrency/range_object_map.rs

@@ -132,6 +132,10 @@ impl<T> RangeObjectMap<T> {
     pub fn remove_from_pos(&mut self, pos: Position) {
         self.v.remove(pos);
     }
+
+    pub fn iter(&self) -> impl Iterator<Item = &T> {
+        self.v.iter().map(|e| &e.data)
+    }
 }
 
 impl<T> Index<Position> for RangeObjectMap<T> {

src/concurrency/thread.rs

@@ -32,9 +32,11 @@ pub enum SchedulingAction {
 
 /// Timeout callbacks can be created by synchronization primitives to tell the
 /// scheduler that they should be called once some period of time passes.
-type TimeoutCallback<'mir, 'tcx> = Box<
-    dyn FnOnce(&mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>) -> InterpResult<'tcx> + 'tcx,
->;
+pub trait MachineCallback<'mir, 'tcx>: VisitTags {
+    fn call(&self, ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>) -> InterpResult<'tcx>;
+}
+
+type TimeoutCallback<'mir, 'tcx> = Box<dyn MachineCallback<'mir, 'tcx> + 'tcx>;
 
 /// A thread identifier.
 #[derive(Clone, Copy, Debug, PartialOrd, Ord, PartialEq, Eq, Hash)]
@@ -181,6 +183,46 @@ impl<'mir, 'tcx> Thread<'mir, 'tcx> {
     }
 }
 
+impl VisitTags for Thread<'_, '_> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let Thread { panic_payload, last_error, stack, state: _, thread_name: _, join_status: _ } =
+            self;
+
+        panic_payload.visit_tags(visit);
+        last_error.visit_tags(visit);
+        for frame in stack {
+            frame.visit_tags(visit)
+        }
+    }
+}
+
+impl VisitTags for Frame<'_, '_, Provenance, FrameData<'_>> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let Frame {
+            return_place,
+            locals,
+            extra,
+            body: _,
+            instance: _,
+            return_to_block: _,
+            loc: _,
+            // There are some private fields we cannot access; they contain no tags.
+            ..
+        } = self;
+
+        // Return place.
+        return_place.visit_tags(visit);
+        // Locals.
+        for local in locals.iter() {
+            if let LocalValue::Live(value) = &local.value {
+                value.visit_tags(visit);
+            }
+        }
+
+        extra.visit_tags(visit);
+    }
+}
+
 /// A specific moment in time.
 #[derive(Debug)]
 pub enum Time {
@@ -253,6 +295,29 @@ impl<'mir, 'tcx> Default for ThreadManager<'mir, 'tcx> {
     }
 }
 
+impl VisitTags for ThreadManager<'_, '_> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let ThreadManager {
+            threads,
+            thread_local_alloc_ids,
+            timeout_callbacks,
+            active_thread: _,
+            yield_active_thread: _,
+            sync: _,
+        } = self;
+
+        for thread in threads {
+            thread.visit_tags(visit);
+        }
+        for ptr in thread_local_alloc_ids.borrow().values() {
+            ptr.visit_tags(visit);
+        }
+        for callback in timeout_callbacks.values() {
+            callback.callback.visit_tags(visit);
+        }
+    }
+}
+
 impl<'mir, 'tcx: 'mir> ThreadManager<'mir, 'tcx> {
     pub(crate) fn init(ecx: &mut MiriInterpCx<'mir, 'tcx>) {
         if ecx.tcx.sess.target.os.as_ref() != "windows" {
@@ -625,33 +690,6 @@ 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> {
@@ -930,7 +968,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // 2. Make the scheduler the only place that can change the active
         //    thread.
         let old_thread = this.set_active_thread(thread);
-        callback(this)?;
+        callback.call(this)?;
         this.set_active_thread(old_thread);
         Ok(())
     }

src/concurrency/weak_memory.rs

@@ -108,6 +108,19 @@ pub struct StoreBufferAlloc {
     store_buffers: RefCell<RangeObjectMap<StoreBuffer>>,
 }
 
+impl VisitTags for StoreBufferAlloc {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        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);
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) struct StoreBuffer {
     // Stores to this location in modification order

src/intptrcast.rs

@@ -44,6 +44,12 @@ pub struct GlobalStateInner {
     provenance_mode: ProvenanceMode,
 }
 
+impl VisitTags for GlobalStateInner {
+    fn visit_tags(&self, _visit: &mut dyn FnMut(SbTag)) {
+        // Nothing to visit here.
+    }
+}
+
 impl GlobalStateInner {
     pub fn new(config: &MiriConfig) -> Self {
         GlobalStateInner {

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 _, 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

@@ -63,6 +63,15 @@ impl<'tcx> std::fmt::Debug for FrameData<'tcx> {
     }
 }
 
+impl VisitTags for FrameData<'_> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let FrameData { catch_unwind, stacked_borrows, timing: _ } = self;
+
+        catch_unwind.visit_tags(visit);
+        stacked_borrows.visit_tags(visit);
+    }
+}
+
 /// Extra memory kinds
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum MiriMemoryKind {
@@ -251,6 +260,16 @@ pub struct AllocExtra {
     pub weak_memory: Option<weak_memory::AllocExtra>,
 }
 
+impl VisitTags for AllocExtra {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let AllocExtra { stacked_borrows, data_race, weak_memory } = self;
+
+        stacked_borrows.visit_tags(visit);
+        data_race.visit_tags(visit);
+        weak_memory.visit_tags(visit);
+    }
+}
+
 /// Precomputed layouts of primitive types
 pub struct PrimitiveLayouts<'tcx> {
     pub unit: TyAndLayout<'tcx>,
@@ -591,14 +610,65 @@ 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);
+impl VisitTags for MiriMachine<'_, '_> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        #[rustfmt::skip]
+        let MiriMachine {
+            threads,
+            tls,
+            env_vars,
+            argc,
+            argv,
+            cmd_line,
+            extern_statics,
+            dir_handler,
+            stacked_borrows,
+            data_race,
+            intptrcast,
+            file_handler,
+            tcx: _,
+            isolated_op: _,
+            validate: _,
+            enforce_abi: _,
+            clock: _,
+            layouts: _,
+            static_roots: _,
+            profiler: _,
+            string_cache: _,
+            exported_symbols_cache: _,
+            panic_on_unsupported: _,
+            backtrace_style: _,
+            local_crates: _,
+            rng: _,
+            tracked_alloc_ids: _,
+            check_alignment: _,
+            cmpxchg_weak_failure_rate: _,
+            mute_stdout_stderr: _,
+            weak_memory: _,
+            preemption_rate: _,
+            report_progress: _,
+            basic_block_count: _,
+            #[cfg(unix)]
+            external_so_lib: _,
+            gc_interval: _,
+            since_gc: _,
+            num_cpus: _,
+        } = self;
+
+        threads.visit_tags(visit);
+        tls.visit_tags(visit);
+        env_vars.visit_tags(visit);
+        dir_handler.visit_tags(visit);
+        file_handler.visit_tags(visit);
+        data_race.visit_tags(visit);
+        stacked_borrows.visit_tags(visit);
+        intptrcast.visit_tags(visit);
+        argc.visit_tags(visit);
+        argv.visit_tags(visit);
+        cmd_line.visit_tags(visit);
+        for ptr in extern_statics.values() {
+            ptr.visit_tags(visit);
+        }
     }
 }
 

src/shims/env.rs

@@ -36,6 +36,17 @@ pub struct EnvVars<'tcx> {
     pub(crate) environ: Option<MPlaceTy<'tcx, Provenance>>,
 }
 
+impl VisitTags for EnvVars<'_> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let EnvVars { map, environ } = self;
+
+        environ.visit_tags(visit);
+        for ptr in map.values() {
+            ptr.visit_tags(visit);
+        }
+    }
+}
+
 impl<'tcx> EnvVars<'tcx> {
     pub(crate) fn init<'mir>(
         ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>,

src/shims/panic.rs

@@ -35,6 +35,15 @@ pub struct CatchUnwindData<'tcx> {
     ret: mir::BasicBlock,
 }
 
+impl VisitTags for CatchUnwindData<'_> {
+    fn visit_tags(&self, visit: &mut dyn FnMut(SbTag)) {
+        let CatchUnwindData { catch_fn, data, dest, ret: _ } = self;
+        catch_fn.visit_tags(visit);
+        data.visit_tags(visit);
+        dest.visit_tags(visit);
+    }
+}
+
 impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriInterpCx<'mir, 'tcx> {}
 pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     /// Handles the special `miri_start_panic` intrinsic, which is called