Expand VisitMachineValues to cover more pointers in the interpreter

src/tools/miri/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/tools/miri/src/concurrency/thread.rs

@@ -181,6 +181,41 @@ impl<'mir, 'tcx> Thread<'mir, 'tcx> {
     }
 }
 
+impl VisitMachineValues for Thread<'_, '_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        let Thread { panic_payload, last_error, stack, .. } = self;
+
+        if let Some(payload) = panic_payload {
+            visit(&Operand::Immediate(Immediate::Scalar(*payload)))
+        }
+        if let Some(error) = last_error {
+            visit(&Operand::Indirect(**error))
+        }
+        for frame in stack {
+            frame.visit_machine_values(visit)
+        }
+    }
+}
+
+impl VisitMachineValues for Frame<'_, '_, Provenance, FrameData<'_>> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        let Frame { return_place, locals, extra, .. } = self;
+
+        // Return place.
+        if let Place::Ptr(mplace) = **return_place {
+            visit(&Operand::Indirect(mplace));
+        }
+        // Locals.
+        for local in locals.iter() {
+            if let LocalValue::Live(value) = &local.value {
+                visit(value);
+            }
+        }
+
+        extra.visit_machine_values(visit);
+    }
+}
+
 /// A specific moment in time.
 #[derive(Debug)]
 pub enum Time {
@@ -253,6 +288,22 @@ impl<'mir, 'tcx> Default for ThreadManager<'mir, 'tcx> {
     }
 }
 
+impl VisitMachineValues for ThreadManager<'_, '_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        let ThreadManager { threads, thread_local_alloc_ids, .. } = self;
+
+        for thread in threads {
+            thread.visit_machine_values(visit);
+        }
+        for ptr in thread_local_alloc_ids.borrow().values().copied() {
+            let ptr: Pointer<Option<Provenance>> = ptr.into();
+            visit(&Operand::Indirect(MemPlace::from_ptr(ptr)));
+        }
+        // FIXME: Do we need to do something for TimeoutCallback? That's a Box<dyn>, not sure what
+        // to do.
+    }
+}
+
 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 +676,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> {

src/tools/miri/src/concurrency/weak_memory.rs

@@ -108,6 +108,19 @@ pub struct StoreBufferAlloc {
     store_buffers: RefCell<RangeObjectMap<StoreBuffer>>,
 }
 
+impl StoreBufferAlloc {
+    pub fn iter(&self, mut visitor: impl FnMut(&Scalar<Provenance>)) {
+        for val in self
+            .store_buffers
+            .borrow()
+            .iter()
+            .flat_map(|buf| buf.buffer.iter().map(|element| &element.val))
+        {
+            visitor(val)
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(super) struct StoreBuffer {
     // Stores to this location in modification order

src/tools/miri/src/machine.rs

@@ -63,6 +63,16 @@ impl<'tcx> std::fmt::Debug for FrameData<'tcx> {
     }
 }
 
+impl VisitMachineValues for FrameData<'_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        let FrameData { catch_unwind, .. } = self;
+
+        if let Some(catch_unwind) = catch_unwind {
+            catch_unwind.visit_machine_values(visit);
+        }
+    }
+}
+
 /// Extra memory kinds
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum MiriMemoryKind {
@@ -593,12 +603,36 @@ 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;
+        let MiriMachine {
+            threads,
+            tls,
+            env_vars,
+            argc,
+            argv,
+            cmd_line,
+            extern_statics,
+            dir_handler,
+            ..
+        } = self;
 
         threads.visit_machine_values(visit);
         tls.visit_machine_values(visit);
+        env_vars.visit_machine_values(visit);
+        dir_handler.visit_machine_values(visit);
+
+        if let Some(argc) = argc {
+            visit(&Operand::Indirect(*argc));
+        }
+        if let Some(argv) = argv {
+            visit(&Operand::Indirect(*argv));
+        }
+        if let Some(cmd_line) = cmd_line {
+            visit(&Operand::Indirect(*cmd_line));
+        }
+        for ptr in extern_statics.values().copied() {
+            let ptr: Pointer<Option<Provenance>> = ptr.into();
+            visit(&Operand::Indirect(MemPlace::from_ptr(ptr)));
+        }
     }
 }
 

src/tools/miri/src/shims/env.rs

@@ -36,6 +36,20 @@ pub struct EnvVars<'tcx> {
     pub(crate) environ: Option<MPlaceTy<'tcx, Provenance>>,
 }
 
+impl VisitMachineValues for EnvVars<'_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        let EnvVars { map, environ } = self;
+
+        for ptr in map.values() {
+            visit(&Operand::Indirect(MemPlace::from_ptr(*ptr)));
+        }
+
+        if let Some(env) = environ {
+            visit(&Operand::Indirect(**env));
+        }
+    }
+}
+
 impl<'tcx> EnvVars<'tcx> {
     pub(crate) fn init<'mir>(
         ecx: &mut InterpCx<'mir, 'tcx, MiriMachine<'mir, 'tcx>>,

src/tools/miri/src/shims/panic.rs

@@ -35,6 +35,13 @@ pub struct CatchUnwindData<'tcx> {
     ret: mir::BasicBlock,
 }
 
+impl VisitMachineValues for CatchUnwindData<'_> {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        visit(&Operand::Indirect(MemPlace::from_ptr(self.catch_fn)));
+        visit(&Operand::Immediate(Immediate::Scalar(self.data)));
+    }
+}
+
 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

src/tools/miri/src/shims/unix/fs.rs

@@ -462,6 +462,14 @@ impl Default for DirHandler {
     }
 }
 
+impl VisitMachineValues for DirHandler {
+    fn visit_machine_values(&self, visit: &mut impl FnMut(&Operand<Provenance>)) {
+        for dir in self.streams.values() {
+            visit(&Operand::Indirect(MemPlace::from_ptr(dir.entry)));
+        }
+    }
+}
+
 fn maybe_sync_file(
     file: &File,
     writable: bool,

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.
-    exposed_tags: FxHashSet<SbTag>,
+    pub exposed_tags: FxHashSet<SbTag>,
     /// Whether this memory has been modified since the last time the tag GC ran
     modified_since_last_gc: bool,
 }

src/tools/miri/src/stacked_borrows/stack.rs

@@ -43,8 +43,11 @@ impl Stack {
     pub fn retain(&mut self, tags: &FxHashSet<SbTag>) {
         let mut first_removed = None;
 
-        let mut read_idx = 1;
-        let mut write_idx = 1;
+        // For stacks with a known bottom, we never consider removing the bottom-most tag, because
+        // that is the base tag which exists whether or not there are any pointers to the
+        // allocation.
+        let mut read_idx = usize::from(self.unknown_bottom.is_none());
+        let mut write_idx = read_idx;
         while read_idx < self.borrows.len() {
             let left = self.borrows[read_idx - 1];
             let this = self.borrows[read_idx];

src/tools/miri/src/tag_gc.rs

@@ -71,6 +71,22 @@ 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);
+
+                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);
+                            }
+                        }
+                    });
+                }
             },
         );