Implement -Zmiri-tag-gc a garbage collector for tags

Author: saethlin

miri

@@ -131,7 +131,7 @@ fi
 # Enable rustc-specific lints (ignored without `-Zunstable-options`).
 export RUSTFLAGS="-Zunstable-options -Wrustc::internal $RUSTFLAGS"
 # We set the rpath so that Miri finds the private rustc libraries it needs.
-export RUSTFLAGS="-C link-args=-Wl,-rpath,$LIBDIR $RUSTFLAGS"
+export RUSTFLAGS="-C link-args=-Wl,-rpath,$LIBDIR $RUSTFLAGS -Cdebuginfo=1"
 
 ## Helper functions
 

src/bin/miri.rs

@@ -521,6 +521,14 @@ fn main() {
                 Err(err) => show_error!("-Zmiri-report-progress requires a `u32`: {}", err),
             };
             miri_config.report_progress = Some(interval);
+        } else if arg == "-Zmiri-tag-gc" {
+            miri_config.gc_interval = Some(10_000);
+        } else if let Some(param) = arg.strip_prefix("-Zmiri-tag-gc=") {
+            let interval = match param.parse::<u32>() {
+                Ok(i) => i,
+                Err(err) => show_error!("-Zmiri-tag-gc requires a `u32`: {}", err),
+            };
+            miri_config.gc_interval = Some(interval);
         } else if let Some(param) = arg.strip_prefix("-Zmiri-measureme=") {
             miri_config.measureme_out = Some(param.to_string());
         } else if let Some(param) = arg.strip_prefix("-Zmiri-backtrace=") {

src/eval.rs

@@ -128,6 +128,8 @@ pub struct MiriConfig {
     pub report_progress: Option<u32>,
     /// Whether Stacked Borrows retagging should recurse into fields of datatypes.
     pub retag_fields: bool,
+    /// Run a garbage collector for SbTags every N basic blocks.
+    pub gc_interval: Option<u32>,
 }
 
 impl Default for MiriConfig {
@@ -159,6 +161,7 @@ impl Default for MiriConfig {
             preemption_rate: 0.01, // 1%
             report_progress: None,
             retag_fields: false,
+            gc_interval: None,
         }
     }
 }
@@ -342,6 +345,7 @@ pub fn eval_entry<'tcx>(
             match ecx.schedule()? {
                 SchedulingAction::ExecuteStep => {
                     assert!(ecx.step()?, "a terminated thread was scheduled for execution");
+                    //ecx.garbage_collect_tags()?;
                 }
                 SchedulingAction::ExecuteTimeoutCallback => {
                     assert!(

src/intptrcast.rs

@@ -36,7 +36,7 @@ pub struct GlobalStateInner {
     base_addr: FxHashMap<AllocId, u64>,
     /// Whether an allocation has been exposed or not. This cannot be put
     /// into `AllocExtra` for the same reason as `base_addr`.
-    exposed: FxHashSet<AllocId>,
+    pub exposed: FxHashSet<AllocId>,
     /// This is used as a memory address when a new pointer is casted to an integer. It
     /// is always larger than any address that was previously made part of a block.
     next_base_addr: u64,
@@ -246,6 +246,11 @@ impl<'mir, 'tcx> GlobalStateInner {
             rem => addr.checked_add(align).unwrap() - rem,
         }
     }
+
+    /// Returns whether the specified `AllocId` has been exposed.
+    pub fn is_exposed(&self, id: AllocId) -> bool {
+        self.exposed.contains(&id)
+    }
 }
 
 #[cfg(test)]

src/lib.rs

@@ -62,6 +62,7 @@ mod operator;
 mod range_map;
 mod shims;
 mod stacked_borrows;
+mod tag_gc;
 pub mod thread;
 
 // Establish a "crate-wide prelude": we often import `crate::*`.
@@ -104,6 +105,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::thread::{
     EvalContextExt as ThreadsEvalContextExt, SchedulingAction, ThreadId, ThreadManager, ThreadState,
 };

src/machine.rs

@@ -358,6 +358,12 @@ pub struct Evaluator<'mir, 'tcx> {
     pub(crate) report_progress: Option<u32>,
     /// The number of blocks that passed since the last progress report.
     pub(crate) since_progress_report: u32,
+
+    /// If `Some`, we will attempt to run a garbage collector on SbTags every N basic blocks.
+    pub(crate) gc_interval: Option<u32>,
+    /// The number of blocks that passed since the last SbTag GC pass.
+    pub(crate) since_gc: u32,
+    //pub(crate) hot_allocations: FxHashSet<AllocId>,
 }
 
 impl<'mir, 'tcx> Evaluator<'mir, 'tcx> {
@@ -412,6 +418,8 @@ impl<'mir, 'tcx> Evaluator<'mir, 'tcx> {
             preemption_rate: config.preemption_rate,
             report_progress: config.report_progress,
             since_progress_report: 0,
+            gc_interval: config.gc_interval,
+            since_gc: 0,
         }
     }
 
@@ -959,6 +967,17 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'mir, 'tcx> {
             // Cannot overflow, since it is strictly less than `report_progress`.
             ecx.machine.since_progress_report += 1;
         }
+
+        // Search all memory for SbTags, then use Stack::retain to drop all other tags.
+        if let Some(gc_interval) = ecx.machine.gc_interval {
+            if ecx.machine.since_gc >= gc_interval {
+                ecx.machine.since_gc = 0;
+                ecx.garbage_collect_tags()?;
+            } else {
+                ecx.machine.since_gc += 1;
+            }
+        }
+
         // These are our preemption points.
         ecx.maybe_preempt_active_thread();
         Ok(())

src/shims/panic.rs

@@ -57,6 +57,13 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
         assert!(thread.panic_payload.is_none(), "the panic runtime should avoid double-panics");
         thread.panic_payload = Some(payload);
 
+        // Expose the allocation so we don't GC it
+        if let Scalar::Ptr(Pointer { provenance: Provenance::Concrete { alloc_id, .. }, .. }, _) =
+            payload
+        {
+            this.machine.intptrcast.borrow_mut().exposed.insert(alloc_id);
+        }
+
         // Jump to the unwind block to begin unwinding.
         this.unwind_to_block(unwind)?;
         Ok(())

src/shims/tls.rs

@@ -233,6 +233,12 @@ impl<'tcx> TlsData<'tcx> {
             data.remove(&thread_id);
         }
     }
+
+    pub fn iter(&self, mut visitor: impl FnMut(&Scalar<Provenance>)) {
+        for scalar in self.keys.values().flat_map(|v| v.data.values()) {
+            visitor(scalar);
+        }
+    }
 }
 
 impl<'mir, 'tcx: 'mir> EvalContextPrivExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}

src/stacked_borrows/mod.rs

@@ -80,6 +80,7 @@ pub struct Stacks {
     history: AllocHistory,
     /// The set of tags that have been exposed inside this allocation.
     exposed_tags: FxHashSet<SbTag>,
+    dirty: bool,
 }
 
 /// Extra global state, available to the memory access hooks.
@@ -422,6 +423,7 @@ impl<'tcx> Stack {
             let item = self.get(idx).unwrap();
             Stack::item_popped(&item, global, dcx)?;
         }
+
         Ok(())
     }
 
@@ -496,6 +498,20 @@ impl<'tcx> Stack {
 }
 // # Stacked Borrows Core End
 
+/// Integration with the SbTag garbage collector
+impl Stacks {
+    pub fn collect(&mut self, live_tags: &FxHashSet<SbTag>) {
+        if self.dirty {
+            for stack in self.stacks.iter_mut_all() {
+                if stack.len() > 64 {
+                    stack.retain(live_tags);
+                }
+            }
+            self.dirty = false;
+        }
+    }
+}
+
 /// 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
@@ -508,6 +524,7 @@ impl<'tcx> Stacks {
             stacks: RangeMap::new(size, stack),
             history: AllocHistory::new(id),
             exposed_tags: FxHashSet::default(),
+            dirty: false,
         }
     }
 
@@ -522,6 +539,7 @@ impl<'tcx> Stacks {
             &mut FxHashSet<SbTag>,
         ) -> InterpResult<'tcx>,
     ) -> InterpResult<'tcx> {
+        self.dirty = true;
         for (offset, stack) in self.stacks.iter_mut(range.start, range.size) {
             let mut dcx = dcx_builder.build(&mut self.history, offset);
             f(stack, &mut dcx, &mut self.exposed_tags)?;
@@ -614,7 +632,7 @@ impl Stacks {
     ) -> InterpResult<'tcx> {
         trace!("deallocation with tag {:?}: {:?}, size {}", tag, alloc_id, range.size.bytes());
         let dcx = DiagnosticCxBuilder::dealloc(&mut current_span, threads, tag);
-        let state = state.borrow();
+        let state = state.borrow_mut();
         self.for_each(range, dcx, |stack, dcx, exposed_tags| {
             stack.dealloc(tag, &state, dcx, exposed_tags)
         })?;

src/stacked_borrows/stack.rs

@@ -39,6 +39,50 @@ pub struct Stack {
     unique_range: Range<usize>,
 }
 
+impl Stack {
+    pub fn retain(&mut self, tags: &FxHashSet<SbTag>) {
+        let prev_len = self.borrows.len();
+        let mut read_idx = 1;
+        let mut write_idx = 1;
+        while read_idx < self.borrows.len() {
+            let left = self.borrows[read_idx - 1];
+            let this = self.borrows[read_idx];
+            let should_keep = match this.perm() {
+                // Always destroy adjacent Disabled tags, they only exist to separate SRW blocks so
+                // two is just as good as one.
+                Permission::Disabled => left.perm() != Permission::Disabled,
+                // SRW and SRO don't do anything special, so we keep them if they are in use
+                Permission::SharedReadWrite | Permission::SharedReadOnly =>
+                    tags.contains(&this.tag()),
+                Permission::Unique =>
+                    left.perm() != Permission::Unique || tags.contains(&this.tag()),
+            };
+
+            if should_keep {
+                if read_idx != write_idx {
+                    self.borrows[write_idx] = self.borrows[read_idx];
+                }
+                write_idx += 1;
+            }
+
+            read_idx += 1;
+        }
+        self.borrows.truncate(write_idx);
+
+        // Reset the cache if anything was changed
+        #[cfg(feature = "stack-cache")]
+        if self.borrows.len() != prev_len {
+            if !self.unique_range.is_empty() {
+                self.unique_range = 0..self.len();
+            }
+            for i in 0..CACHE_LEN {
+                self.cache.items[i] = self.borrows[0];
+                self.cache.idx[i] = 0;
+            }
+        }
+    }
+}
+
 /// A very small cache of searches of a borrow stack, mapping `Item`s to their position in said stack.
 ///
 /// It may seem like maintaining this cache is a waste for small stacks, but
@@ -105,12 +149,14 @@ impl<'tcx> Stack {
 
         // Check that the unique_range is a valid index into the borrow stack.
         // This asserts that the unique_range's start <= end.
-        let uniques = &self.borrows[self.unique_range.clone()];
+        let _uniques = &self.borrows[self.unique_range.clone()];
 
         // Check that the start of the unique_range is precise.
+        /*
         if let Some(first_unique) = uniques.first() {
             assert_eq!(first_unique.perm(), Permission::Unique);
         }
+        */
         // We cannot assert that the unique range is exact on the upper end.
         // When we pop items within the unique range, setting the end of the range precisely
         // requires doing a linear search of the borrow stack, which is exactly the kind of

src/tag_gc.rs

@@ -0,0 +1,111 @@
+use crate::*;
+use rustc_data_structures::fx::FxHashSet;
+
+impl<'mir, 'tcx: 'mir> EvalContextExt<'mir, 'tcx> for crate::MiriEvalContext<'mir, 'tcx> {}
+pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriEvalContextExt<'mir, 'tcx> {
+    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.
+        if this.machine.stacked_borrows.is_none() {
+            return Ok(());
+        }
+
+        let mut tags = FxHashSet::default();
+
+        self.find_pointers_in_tls(&mut tags);
+        self.find_pointers_in_memory(&mut tags);
+        self.find_pointers_in_locals(&mut tags)?;
+
+        self.remove_extraneous_tags(tags);
+
+        Ok(())
+    }
+
+    fn find_pointers_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_pointers_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.relocations().iter() {
+                    if let Provenance::Concrete { sb, .. } = prov {
+                        tags.insert(*sb);
+                    }
+                }
+            }
+        });
+    }
+
+    fn find_pointers_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 {
+                    tags.insert(sb);
+                }
+            }
+
+            for local in frame.locals.iter() {
+                let LocalValue::Live(value) = local.value else {
+                continue;
+            };
+                match value {
+                    Operand::Immediate(Immediate::Scalar(ScalarMaybeUninit::Scalar(
+                        Scalar::Ptr(ptr, _),
+                    ))) =>
+                        if let Provenance::Concrete { sb, .. } = ptr.provenance {
+                            tags.insert(sb);
+                        },
+                    Operand::Immediate(Immediate::ScalarPair(s1, s2)) => {
+                        if let ScalarMaybeUninit::Scalar(Scalar::Ptr(ptr, _)) = s1 {
+                            if let Provenance::Concrete { sb, .. } = ptr.provenance {
+                                tags.insert(sb);
+                            }
+                        }
+                        if let ScalarMaybeUninit::Scalar(Scalar::Ptr(ptr, _)) = s2 {
+                            if let Provenance::Concrete { sb, .. } = ptr.provenance {
+                                tags.insert(sb);
+                            }
+                        }
+                    }
+                    Operand::Indirect(MemPlace { ptr, .. }) => {
+                        if let Some(Provenance::Concrete { sb, .. }) = ptr.provenance {
+                            tags.insert(sb);
+                        }
+                    }
+                    Operand::Immediate(Immediate::Uninit)
+                    | Operand::Immediate(Immediate::Scalar(ScalarMaybeUninit::Uninit))
+                    | Operand::Immediate(Immediate::Scalar(ScalarMaybeUninit::Scalar(
+                        Scalar::Int(_),
+                    ))) => {}
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn remove_extraneous_tags(&mut self, tags: FxHashSet<SbTag>) {
+        let this = self.eval_context_mut();
+        this.memory.alloc_map().iter(|it| {
+            for (id, (_kind, alloc)) in it {
+                // Don't GC any allocation which has been exposed
+                if this.machine.intptrcast.borrow().is_exposed(*id) {
+                    continue;
+                }
+
+                alloc.extra.stacked_borrows.as_ref().unwrap().borrow_mut().collect(&tags);
+            }
+        });
+    }
+}