Miri: Skip over GlobalAllocs when collecting

Author: saethlin

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -57,10 +57,20 @@ impl<T: fmt::Display> fmt::Display for MemoryKind<T> {
     }
 }
 
+#[derive(Clone, Copy)]
+pub enum Liveness {
+    Live,
+    Dead,
+    Global,
+}
+
 /// The return value of `get_alloc_info` indicates the "kind" of the allocation.
+#[derive(Clone, Copy)]
 pub enum AllocKind {
     /// A regular live data allocation.
     LiveData,
+    /// A regular live data allocation that will never be deallocated.
+    GlobalLiveData,
     /// A function allocation (that fn ptrs point to).
     Function,
     /// A (symbolic) vtable allocation.
@@ -69,6 +79,24 @@ pub enum AllocKind {
     Dead,
 }
 
+impl AllocKind {
+    pub fn is_global(self) -> bool {
+        use AllocKind::*;
+        match self {
+            GlobalLiveData | Function | VTable => true,
+            LiveData | Dead => false,
+        }
+    }
+
+    pub fn is_live_data(self) -> bool {
+        use AllocKind::*;
+        match self {
+            GlobalLiveData | LiveData => true,
+            Function | VTable | Dead => false,
+        }
+    }
+}
+
 /// The value of a function pointer.
 #[derive(Debug, Copy, Clone)]
 pub enum FnVal<'tcx, Other> {
@@ -706,10 +734,17 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     /// Check whether an allocation is live. This is faster than calling
     /// [`InterpCx::get_alloc_info`] if all you need to check is whether the kind is
     /// [`AllocKind::Dead`] because it doesn't have to look up the type and layout of statics.
-    pub fn is_alloc_live(&self, id: AllocId) -> bool {
-        self.tcx.try_get_global_alloc(id).is_some()
-            || self.memory.alloc_map.contains_key_ref(&id)
-            || self.memory.extra_fn_ptr_map.contains_key(&id)
+    pub fn is_alloc_live(&self, id: AllocId) -> Liveness {
+        if self.memory.alloc_map.contains_key_ref(&id) {
+            return Liveness::Live;
+        }
+        if self.tcx.try_get_global_alloc(id).is_some() {
+            return Liveness::Global;
+        }
+        if self.memory.extra_fn_ptr_map.contains_key(&id) {
+            return Liveness::Live;
+        }
+        Liveness::Dead
     }
 
     /// Obtain the size and alignment of an allocation, even if that allocation has
@@ -744,13 +779,13 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                     .expect("statics should not have generic parameters");
                 let layout = self.tcx.layout_of(ParamEnv::empty().and(ty)).unwrap();
                 assert!(layout.is_sized());
-                (layout.size, layout.align.abi, AllocKind::LiveData)
+                (layout.size, layout.align.abi, AllocKind::GlobalLiveData)
             }
             Some(GlobalAlloc::Memory(alloc)) => {
                 // Need to duplicate the logic here, because the global allocations have
                 // different associated types than the interpreter-local ones.
                 let alloc = alloc.inner();
-                (alloc.size(), alloc.align, AllocKind::LiveData)
+                (alloc.size(), alloc.align, AllocKind::GlobalLiveData)
             }
             Some(GlobalAlloc::Function(_)) => bug!("We already checked function pointers above"),
             Some(GlobalAlloc::VTable(..)) => {

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -25,7 +25,7 @@ pub use self::intern::{
     intern_const_alloc_for_constprop, intern_const_alloc_recursive, InternKind,
 };
 pub use self::machine::{compile_time_machine, AllocMap, Machine, MayLeak, StackPopJump};
-pub use self::memory::{AllocKind, AllocRef, AllocRefMut, FnVal, Memory, MemoryKind};
+pub use self::memory::{AllocKind, AllocRef, AllocRefMut, FnVal, Liveness, Memory, MemoryKind};
 pub use self::operand::{ImmTy, Immediate, OpTy, Readable};
 pub use self::place::{MPlaceTy, MemPlaceMeta, PlaceTy, Writeable};
 pub use self::projection::{OffsetMode, Projectable};

src/tools/miri/src/borrow_tracker/mod.rs

@@ -91,7 +91,7 @@ pub struct GlobalStateInner {
     /// Table storing the "base" tag for each allocation.
     /// The base tag is the one used for the initial pointer.
     /// We need this in a separate table to handle cyclic statics.
-    base_ptr_tags: FxHashMap<AllocId, BorTag>,
+    base_ptr_tags: FxHashMap<AllocId, (BorTag, bool)>,
     /// Next unused call ID (for protectors).
     next_call_id: CallId,
     /// All currently protected tags.
@@ -222,7 +222,7 @@ impl GlobalStateInner {
     }
 
     pub fn base_ptr_tag(&mut self, id: AllocId, machine: &MiriMachine<'_, '_>) -> BorTag {
-        self.base_ptr_tags.get(&id).copied().unwrap_or_else(|| {
+        self.base_ptr_tags.get(&id).map(|(id, _is_global)| *id).unwrap_or_else(|| {
             let tag = self.new_ptr();
             if self.tracked_pointer_tags.contains(&tag) {
                 machine.emit_diagnostic(NonHaltingDiagnostic::CreatedPointerTag(
@@ -232,13 +232,13 @@ impl GlobalStateInner {
                 ));
             }
             trace!("New allocation {:?} has base tag {:?}", id, tag);
-            self.base_ptr_tags.try_insert(id, tag).unwrap();
+            self.base_ptr_tags.try_insert(id, (tag, false)).unwrap();
             tag
         })
     }
 
     pub fn remove_unreachable_allocs(&mut self, allocs: &LiveAllocs<'_, '_, '_>) {
-        self.base_ptr_tags.retain(|id, _| allocs.is_live(*id));
+        self.base_ptr_tags.retain(|id, (_, is_global)| allocs.is_live(*id, is_global));
     }
 }
 
@@ -382,7 +382,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             // pointer is readable, and the implicit read access inserted
             // will never cause UB on the pointer itself.
             let (_, _, kind) = this.get_alloc_info(*alloc_id);
-            if matches!(kind, AllocKind::LiveData) {
+            if kind.is_live_data() {
                 let alloc_extra = this.get_alloc_extra(*alloc_id).unwrap();
                 let alloc_borrow_tracker = &alloc_extra.borrow_tracker.as_ref().unwrap();
                 alloc_borrow_tracker.release_protector(&this.machine, borrow_tracker, *tag)?;

src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs

@@ -648,32 +648,27 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
             };
 
             let (_size, _align, alloc_kind) = this.get_alloc_info(alloc_id);
-            match alloc_kind {
-                AllocKind::LiveData => {
-                    // This should have alloc_extra data, but `get_alloc_extra` can still fail
-                    // if converting this alloc_id from a global to a local one
-                    // uncovers a non-supported `extern static`.
-                    let extra = this.get_alloc_extra(alloc_id)?;
-                    let mut stacked_borrows = extra
-                        .borrow_tracker_sb()
-                        .borrow_mut();
-                    // Note that we create a *second* `DiagnosticCxBuilder` below for the actual retag.
-                    // FIXME: can this be done cleaner?
-                    let dcx = DiagnosticCxBuilder::retag(
-                        &this.machine,
-                        retag_info,
-                        new_tag,
-                        orig_tag,
-                        alloc_range(base_offset, size),
-                    );
-                    let mut dcx = dcx.build(&mut stacked_borrows.history, base_offset);
-                    dcx.log_creation();
-                    if new_perm.protector().is_some() {
-                        dcx.log_protector();
-                    }
-                },
-                AllocKind::Function | AllocKind::VTable | AllocKind::Dead => {
-                    // No stacked borrows on these allocations.
+            if alloc_kind.is_live_data() {
+                // This should have alloc_extra data, but `get_alloc_extra` can still fail
+                // if converting this alloc_id from a global to a local one
+                // uncovers a non-supported `extern static`.
+                let extra = this.get_alloc_extra(alloc_id)?;
+                let mut stacked_borrows = extra
+                    .borrow_tracker_sb()
+                    .borrow_mut();
+                // Note that we create a *second* `DiagnosticCxBuilder` below for the actual retag.
+                // FIXME: can this be done cleaner?
+                let dcx = DiagnosticCxBuilder::retag(
+                    &this.machine,
+                    retag_info,
+                    new_tag,
+                    orig_tag,
+                    alloc_range(base_offset, size),
+                );
+                let mut dcx = dcx.build(&mut stacked_borrows.history, base_offset);
+                dcx.log_creation();
+                if new_perm.protector().is_some() {
+                    dcx.log_protector();
                 }
             }
             Ok(())
@@ -1012,18 +1007,13 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // This is okay because accessing them is UB anyway, no need for any Stacked Borrows checks.
         // NOT using `get_alloc_extra_mut` since this might be a read-only allocation!
         let (_size, _align, kind) = this.get_alloc_info(alloc_id);
-        match kind {
-            AllocKind::LiveData => {
-                // This should have alloc_extra data, but `get_alloc_extra` can still fail
-                // if converting this alloc_id from a global to a local one
-                // uncovers a non-supported `extern static`.
-                let alloc_extra = this.get_alloc_extra(alloc_id)?;
-                trace!("Stacked Borrows tag {tag:?} exposed in {alloc_id:?}");
-                alloc_extra.borrow_tracker_sb().borrow_mut().exposed_tags.insert(tag);
-            }
-            AllocKind::Function | AllocKind::VTable | AllocKind::Dead => {
-                // No stacked borrows on these allocations.
-            }
+        if kind.is_live_data() {
+            // This should have alloc_extra data, but `get_alloc_extra` can still fail
+            // if converting this alloc_id from a global to a local one
+            // uncovers a non-supported `extern static`.
+            let alloc_extra = this.get_alloc_extra(alloc_id)?;
+            trace!("Stacked Borrows tag {tag:?} exposed in {alloc_id:?}");
+            alloc_extra.borrow_tracker_sb().borrow_mut().exposed_tags.insert(tag);
         }
         Ok(())
     }

src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs

@@ -285,7 +285,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
         }
 
         let alloc_kind = this.get_alloc_info(alloc_id).2;
-        if !matches!(alloc_kind, AllocKind::LiveData) {
+        if !alloc_kind.is_live_data() {
             assert_eq!(ptr_size, Size::ZERO); // we did the deref check above, size has to be 0 here
             // There's not actually any bytes here where accesses could even be tracked.
             // Just produce the new provenance, nothing else to do.
@@ -538,18 +538,13 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // This is okay because accessing them is UB anyway, no need for any Tree Borrows checks.
         // NOT using `get_alloc_extra_mut` since this might be a read-only allocation!
         let (_size, _align, kind) = this.get_alloc_info(alloc_id);
-        match kind {
-            AllocKind::LiveData => {
-                // This should have alloc_extra data, but `get_alloc_extra` can still fail
-                // if converting this alloc_id from a global to a local one
-                // uncovers a non-supported `extern static`.
-                let alloc_extra = this.get_alloc_extra(alloc_id)?;
-                trace!("Tree Borrows tag {tag:?} exposed in {alloc_id:?}");
-                alloc_extra.borrow_tracker_tb().borrow_mut().expose_tag(tag);
-            }
-            AllocKind::Function | AllocKind::VTable | AllocKind::Dead => {
-                // No tree borrows on these allocations.
-            }
+        if kind.is_live_data() {
+            // This should have alloc_extra data, but `get_alloc_extra` can still fail
+            // if converting this alloc_id from a global to a local one
+            // uncovers a non-supported `extern static`.
+            let alloc_extra = this.get_alloc_extra(alloc_id)?;
+            trace!("Tree Borrows tag {tag:?} exposed in {alloc_id:?}");
+            alloc_extra.borrow_tracker_tb().borrow_mut().expose_tag(tag);
         }
         Ok(())
     }

src/tools/miri/src/intptrcast.rs

@@ -35,7 +35,7 @@ pub struct GlobalStateInner {
     /// `AllocExtra` because function pointers also have a base address, and
     /// they do not have an `AllocExtra`.
     /// This is the inverse of `int_to_ptr_map`.
-    base_addr: FxHashMap<AllocId, u64>,
+    base_addr: FxHashMap<AllocId, (u64, bool)>,
     /// Whether an allocation has been exposed or not. This cannot be put
     /// into `AllocExtra` for the same reason as `base_addr`.
     exposed: FxHashSet<AllocId>,
@@ -78,7 +78,7 @@ impl GlobalStateInner {
     pub fn remove_unreachable_allocs(&mut self, allocs: &LiveAllocs<'_, '_, '_>) {
         // `exposed` and `int_to_ptr_map` are cleared immediately when an allocation
         // is freed, so `base_addr` is the only one we have to clean up based on the GC.
-        self.base_addr.retain(|id, _| allocs.is_live(*id));
+        self.base_addr.retain(|id, (_, is_global)| allocs.is_live(*id, is_global));
     }
 }
 
@@ -125,7 +125,7 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         // We only use this provenance if it has been exposed.
         if global_state.exposed.contains(&alloc_id) {
             // This must still be live, since we remove allocations from `int_to_ptr_map` when they get freed.
-            debug_assert!(ecx.is_alloc_live(alloc_id));
+            debug_assert!(!matches!(ecx.is_alloc_live(alloc_id), Liveness::Dead));
             Some(alloc_id)
         } else {
             None
@@ -138,7 +138,7 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         let global_state = &mut *global_state;
 
         Ok(match global_state.base_addr.entry(alloc_id) {
-            Entry::Occupied(entry) => *entry.get(),
+            Entry::Occupied(entry) => entry.get().0,
             Entry::Vacant(entry) => {
                 let (size, align, kind) = ecx.get_alloc_info(alloc_id);
                 // This is either called immediately after allocation (and then cached), or when
@@ -161,7 +161,7 @@ trait EvalContextExtPriv<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     .checked_add(slack)
                     .ok_or_else(|| err_exhaust!(AddressSpaceFull))?;
                 let base_addr = align_addr(base_addr, align.bytes());
-                entry.insert(base_addr);
+                entry.insert((base_addr, kind.is_global()));
                 trace!(
                     "Assigning base address {:#x} to allocation {:?} (size: {}, align: {}, slack: {})",
                     base_addr,
@@ -204,7 +204,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
         // Exposing a dead alloc is a no-op, because it's not possible to get a dead allocation
         // via int2ptr.
-        if !ecx.is_alloc_live(alloc_id) {
+        if matches!(ecx.is_alloc_live(alloc_id), Liveness::Dead) {
             return Ok(());
         }
         trace!("Exposing allocation id {alloc_id:?}");
@@ -312,7 +312,7 @@ impl GlobalStateInner {
         // returns a dead allocation.
         // To avoid a linear scan we first look up the address in `base_addr`, and then find it in
         // `int_to_ptr_map`.
-        let addr = *self.base_addr.get(&dead_id).unwrap();
+        let addr = self.base_addr.get(&dead_id).unwrap().0;
         let pos = self.int_to_ptr_map.binary_search_by_key(&addr, |(addr, _)| *addr).unwrap();
         let removed = self.int_to_ptr_map.remove(pos);
         assert_eq!(removed, (addr, dead_id)); // double-check that we removed the right thing

src/tools/miri/src/machine.rs

@@ -530,7 +530,7 @@ pub struct MiriMachine<'mir, 'tcx> {
 
     /// The spans we will use to report where an allocation was created and deallocated in
     /// diagnostics.
-    pub(crate) allocation_spans: RefCell<FxHashMap<AllocId, (Span, Option<Span>)>>,
+    pub(crate) allocation_spans: RefCell<FxHashMap<AllocId, (Span, Option<Span>, bool)>>,
 }
 
 impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
@@ -781,14 +781,14 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
         self.allocation_spans
             .borrow()
             .get(&alloc_id)
-            .map(|(allocated, _deallocated)| allocated.data())
+            .map(|(allocated, _deallocated, _)| allocated.data())
     }
 
     pub(crate) fn deallocated_span(&self, alloc_id: AllocId) -> Option<SpanData> {
         self.allocation_spans
             .borrow()
             .get(&alloc_id)
-            .and_then(|(_allocated, deallocated)| *deallocated)
+            .and_then(|(_allocated, deallocated, _)| *deallocated)
             .map(Span::data)
     }
 }
@@ -1120,7 +1120,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
             ecx.machine
                 .allocation_spans
                 .borrow_mut()
-                .insert(id, (ecx.machine.current_span(), None));
+                .insert(id, (ecx.machine.current_span(), None, false));
         }
 
         Ok(Cow::Owned(alloc))
@@ -1258,7 +1258,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         if let Some(borrow_tracker) = &mut alloc_extra.borrow_tracker {
             borrow_tracker.before_memory_deallocation(alloc_id, prove_extra, range, machine)?;
         }
-        if let Some((_, deallocated_at)) = machine.allocation_spans.borrow_mut().get_mut(&alloc_id)
+        if let Some((_, deallocated_at, _)) = machine.allocation_spans.borrow_mut().get_mut(&alloc_id)
         {
             *deallocated_at = Some(machine.current_span());
         }
@@ -1447,7 +1447,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         };
         let local_decl = &ecx.active_thread_stack()[frame].body.local_decls[local];
         let span = local_decl.source_info.span;
-        ecx.machine.allocation_spans.borrow_mut().insert(alloc_id, (span, None));
+        ecx.machine.allocation_spans.borrow_mut().insert(alloc_id, (span, None, false));
         Ok(())
     }
 }

src/tools/miri/src/provenance_gc.rs

@@ -156,9 +156,21 @@ pub struct LiveAllocs<'a, 'mir, 'tcx> {
 }
 
 impl LiveAllocs<'_, '_, '_> {
-    pub fn is_live(&self, id: AllocId) -> bool {
-        self.collected.contains(&id) ||
-        self.ecx.is_alloc_live(id)
+    pub fn is_live(&self, id: AllocId, is_global: &mut bool) -> bool {
+        if *is_global {
+            return true;
+        }
+        if self.collected.contains(&id) {
+            return true;
+        }
+        match self.ecx.is_alloc_live(id) {
+            Liveness::Dead => false,
+            Liveness::Live => true,
+            Liveness::Global => {
+                *is_global = true;
+                true
+            }
+        }
     }
 }
 
@@ -200,7 +212,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: MiriInterpCxExt<'mir, 'tcx> {
             ecx: this,
             collected: allocs,
         };
-        this.machine.allocation_spans.borrow_mut().retain(|id, _| allocs.is_live(*id));
+        this.machine.allocation_spans.borrow_mut().retain(|id, (.., is_live)| allocs.is_live(*id, is_live));
         this.machine.intptrcast.borrow_mut().remove_unreachable_allocs(&allocs);
         if let Some(borrow_tracker) = &this.machine.borrow_tracker {
             borrow_tracker.borrow_mut().remove_unreachable_allocs(&allocs);