Store protectors outside Item, pack Tag and Perm

Previously, Item was a struct of a NonZeroU64, an Option which was
usually unset or irrelevant, and a 4-variant enum. So collectively, the
size of an Item was 24 bytes, but only 8 bytes were used for the most
part.

So this takes advantage of the fact that it is probably impossible to
exhaust the total space of SbTags, and steals 3 bits from it to pack the
whole struct into a single u64. This bit-packing means that we reduce
peak memory usage when Miri goes memory-bound by ~3x. We also get CPU
performance improvements of varying size, because not only are we simply
accessing less memory, we can now compare a Vec<Item> using a memcmp
because it does not have any padding.

Author: saethlin

src/machine.rs

@@ -54,15 +54,18 @@ pub struct FrameData<'tcx> {
     /// for the start of this frame. When we finish executing this frame,
     /// we use this to register a completed event with `measureme`.
     pub timing: Option<measureme::DetachedTiming>,
+
+    pub protected_tags: Vec<SbTag>,
 }
 
 impl<'tcx> std::fmt::Debug for FrameData<'tcx> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         // Omitting `timing`, it does not support `Debug`.
-        let FrameData { call_id, catch_unwind, timing: _ } = self;
+        let FrameData { call_id, catch_unwind, timing: _, protected_tags } = self;
         f.debug_struct("FrameData")
             .field("call_id", call_id)
             .field("catch_unwind", catch_unwind)
+            .field("protected_tags", protected_tags)
             .finish()
     }
 }
@@ -892,7 +895,7 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'mir, 'tcx> {
             stacked_borrows.borrow_mut().new_call()
         });
 
-        let extra = FrameData { call_id, catch_unwind: None, timing };
+        let extra = FrameData { call_id, catch_unwind: None, timing, protected_tags: Vec::new() };
         Ok(frame.with_extra(extra))
     }
 
@@ -936,7 +939,11 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for Evaluator<'mir, 'tcx> {
     ) -> InterpResult<'tcx, StackPopJump> {
         let timing = frame.extra.timing.take();
         if let Some(stacked_borrows) = &ecx.machine.stacked_borrows {
-            stacked_borrows.borrow_mut().end_call(frame.extra.call_id);
+            let mut sb = stacked_borrows.borrow_mut();
+            sb.end_call(frame.extra.call_id);
+            for tag in frame.extra.protected_tags.drain(..) {
+                sb.protected_tags.remove(&tag);
+            }
         }
         let res = ecx.handle_stack_pop_unwind(frame.extra, unwinding);
         if let Some(profiler) = ecx.machine.profiler.as_ref() {

src/stacked_borrows.rs

@@ -92,27 +92,90 @@ pub enum Permission {
     Disabled,
 }
 
-/// An item in the per-location borrow stack.
-#[derive(Copy, Clone, Hash, PartialEq, Eq)]
-pub struct Item {
-    /// The permission this item grants.
-    perm: Permission,
-    /// The pointers the permission is granted to.
-    tag: SbTag,
-    /// An optional protector, ensuring the item cannot get popped until `CallId` is over.
-    protector: Option<CallId>,
+impl Permission {
+    fn to_bits(self) -> u64 {
+        match self {
+            Permission::Unique => 0,
+            Permission::SharedReadWrite => 1,
+            Permission::SharedReadOnly => 2,
+            Permission::Disabled => 3,
+        }
+    }
+
+    fn from_bits(perm: u64) -> Self {
+        match perm {
+            0 => Permission::Unique,
+            1 => Permission::SharedReadWrite,
+            2 => Permission::SharedReadOnly,
+            3 => Permission::Disabled,
+            _ => unreachable!(),
+        }
+    }
 }
 
-impl fmt::Debug for Item {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "[{:?} for {:?}", self.perm, self.tag)?;
-        if let Some(call) = self.protector {
-            write!(f, " (call {})", call)?;
+mod item {
+    use super::{Permission, SbTag};
+    use std::fmt;
+
+    /// An item in the per-location borrow stack.
+    #[derive(Copy, Clone, Hash, PartialEq, Eq)]
+    pub struct Item(u64);
+
+    // An Item contains 3 bitfields:
+    // * Bits 0-61 store an SbTag
+    // * Bits 61-63 store a Permission
+    // * Bit 64 stores a flag which indicates if we have a protector
+    const TAG_MASK: u64 = u64::MAX >> 3;
+    const PERM_MASK: u64 = 0x3 << 61;
+    const PROTECTED_MASK: u64 = 0x1 << 63;
+
+    const PERM_SHIFT: u64 = 61;
+    const PROTECTED_SHIFT: u64 = 63;
+
+    impl Item {
+        pub fn new(tag: SbTag, perm: Permission, protected: bool) -> Self {
+            assert!(tag.0.get() <= TAG_MASK);
+            let packed_tag = tag.0.get();
+            let packed_perm = perm.to_bits() << PERM_SHIFT;
+            let packed_protected = (protected as u64) << PROTECTED_SHIFT;
+
+            let new = Self(packed_tag + packed_perm + packed_protected);
+
+            debug_assert!(new.tag() == tag);
+            debug_assert!(new.perm() == perm);
+            debug_assert!(new.protected() == protected);
+
+            new
+        }
+
+        /// The pointers the permission is granted to.
+        pub fn tag(self) -> SbTag {
+            unsafe { SbTag(std::num::NonZeroU64::new_unchecked(self.0 & TAG_MASK)) }
+        }
+
+        /// The permission this item grants.
+        pub fn perm(self) -> Permission {
+            Permission::from_bits((self.0 & PERM_MASK) >> PERM_SHIFT)
+        }
+
+        /// Whether or not there is a protector for this tag
+        pub fn protected(self) -> bool {
+            self.0 & PROTECTED_MASK > 0
+        }
+
+        /// Set the Permission stored in this Item to Permission::Disabled
+        pub fn set_disabled(&mut self) {
+            self.0 |= Permission::Disabled.to_bits() << PERM_SHIFT;
+        }
+    }
+
+    impl fmt::Debug for Item {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            write!(f, "[{:?} for {:?}]", self.perm(), self.tag())
         }
-        write!(f, "]")?;
-        Ok(())
     }
 }
+pub use item::Item;
 
 /// Extra per-allocation state.
 #[derive(Clone, Debug)]
@@ -138,6 +201,8 @@ pub struct GlobalStateInner {
     next_call_id: CallId,
     /// Those call IDs corresponding to functions that are still running.
     active_calls: FxHashSet<CallId>,
+    /// All tags currently protected
+    pub(crate) protected_tags: FxHashMap<SbTag, CallId>,
     /// The pointer ids to trace
     tracked_pointer_tags: HashSet<SbTag>,
     /// The call ids to trace
@@ -202,6 +267,7 @@ impl GlobalStateInner {
             base_ptr_tags: FxHashMap::default(),
             next_call_id: NonZeroU64::new(1).unwrap(),
             active_calls: FxHashSet::default(),
+            protected_tags: FxHashMap::default(),
             tracked_pointer_tags,
             tracked_call_ids,
             retag_fields,
@@ -230,10 +296,6 @@ impl GlobalStateInner {
         assert!(self.active_calls.remove(&id));
     }
 
-    fn is_active(&self, id: CallId) -> bool {
-        self.active_calls.contains(&id)
-    }
-
     pub fn base_ptr_tag(&mut self, id: AllocId) -> SbTag {
         self.base_ptr_tags.get(&id).copied().unwrap_or_else(|| {
             let tag = self.new_ptr();
@@ -287,7 +349,7 @@ impl<'tcx> Stack {
     /// Find the first write-incompatible item above the given one --
     /// i.e, find the height to which the stack will be truncated when writing to `granting`.
     fn find_first_write_incompatible(&self, granting: usize) -> usize {
-        let perm = self.get(granting).unwrap().perm;
+        let perm = self.get(granting).unwrap().perm();
         match perm {
             Permission::SharedReadOnly => bug!("Cannot use SharedReadOnly for writing"),
             Permission::Disabled => bug!("Cannot use Disabled for anything"),
@@ -299,7 +361,7 @@ impl<'tcx> Stack {
                 // The SharedReadWrite *just* above us are compatible, to skip those.
                 let mut idx = granting + 1;
                 while let Some(item) = self.get(idx) {
-                    if item.perm == Permission::SharedReadWrite {
+                    if item.perm() == Permission::SharedReadWrite {
                         // Go on.
                         idx += 1;
                     } else {
@@ -326,31 +388,36 @@ impl<'tcx> Stack {
         global: &GlobalStateInner,
         alloc_history: &mut AllocHistory,
     ) -> InterpResult<'tcx> {
-        if global.tracked_pointer_tags.contains(&item.tag) {
+        if global.tracked_pointer_tags.contains(&item.tag()) {
             register_diagnostic(NonHaltingDiagnostic::PoppedPointerTag(
                 *item,
                 provoking_access.map(|(tag, _alloc_range, _size, access)| (tag, access)),
             ));
         }
 
-        if let Some(call) = item.protector {
-            if global.is_active(call) {
-                if let Some((tag, _alloc_range, _offset, _access)) = provoking_access {
-                    Err(err_sb_ub(
-                        format!(
-                            "not granting access to tag {:?} because incompatible item is protected: {:?}",
-                            tag, item
-                        ),
-                        None,
-                        tag.and_then(|tag| alloc_history.get_logs_relevant_to(tag, Some(item.tag))),
-                    ))?
-                } else {
-                    Err(err_sb_ub(
-                        format!("deallocating while item is protected: {:?}", item),
-                        None,
-                        None,
-                    ))?
-                }
+        if !item.protected() {
+            return Ok(());
+        }
+
+        if let Some(call_id) = global.protected_tags.get(&item.tag()) {
+            if let Some((tag, _alloc_range, _offset, _access)) = provoking_access {
+                Err(err_sb_ub(
+                    format!(
+                        "not granting access to tag {:?} because incompatible item is protected: {:?} (call {:?})",
+                        tag, item, call_id
+                    ),
+                    None,
+                    tag.and_then(|tag| alloc_history.get_logs_relevant_to(tag, Some(item.tag()))),
+                ))?
+            } else {
+                Err(err_sb_ub(
+                    format!(
+                        "deallocating while item is protected: {:?} (call {:?})",
+                        item, call_id
+                    ),
+                    None,
+                    None,
+                ))?
             }
         }
         Ok(())
@@ -400,7 +467,7 @@ impl<'tcx> Stack {
                     global,
                     alloc_history,
                 )?;
-                alloc_history.log_invalidation(item.tag, alloc_range, current_span);
+                alloc_history.log_invalidation(item.tag(), alloc_range, current_span);
                 Ok(())
             })?;
         } else {
@@ -426,7 +493,7 @@ impl<'tcx> Stack {
                     global,
                     alloc_history,
                 )?;
-                alloc_history.log_invalidation(item.tag, alloc_range, current_span);
+                alloc_history.log_invalidation(item.tag(), alloc_range, current_span);
                 Ok(())
             })?;
         }
@@ -439,9 +506,9 @@ impl<'tcx> Stack {
             for i in 0..self.len() {
                 let item = self.get(i).unwrap();
                 // Skip disabled items, they cannot be matched anyway.
-                if !matches!(item.perm, Permission::Disabled) {
+                if !matches!(item.perm(), Permission::Disabled) {
                     // We are looking for a strict upper bound, so add 1 to this tag.
-                    max = cmp::max(item.tag.0.checked_add(1).unwrap(), max);
+                    max = cmp::max(item.tag().0.checked_add(1).unwrap(), max);
                 }
             }
             if let Some(unk) = self.unknown_bottom() {
@@ -505,7 +572,7 @@ impl<'tcx> Stack {
     ) -> InterpResult<'tcx> {
         // Figure out which access `perm` corresponds to.
         let access =
-            if new.perm.grants(AccessKind::Write) { AccessKind::Write } else { AccessKind::Read };
+            if new.perm().grants(AccessKind::Write) { AccessKind::Write } else { AccessKind::Read };
 
         // Now we figure out which item grants our parent (`derived_from`) this kind of access.
         // We use that to determine where to put the new item.
@@ -517,7 +584,7 @@ impl<'tcx> Stack {
         // Compute where to put the new item.
         // Either way, we ensure that we insert the new item in a way such that between
         // `derived_from` and the new one, there are only items *compatible with* `derived_from`.
-        let new_idx = if new.perm == Permission::SharedReadWrite {
+        let new_idx = if new.perm() == Permission::SharedReadWrite {
             assert!(
                 access == AccessKind::Write,
                 "this case only makes sense for stack-like accesses"
@@ -578,7 +645,7 @@ impl<'tcx> Stack {
 impl<'tcx> Stacks {
     /// Creates new stack with initial tag.
     fn new(size: Size, perm: Permission, tag: SbTag) -> Self {
-        let item = Item { perm, tag, protector: None };
+        let item = Item::new(tag, perm, false);
         let stack = Stack::new(item);
 
         Stacks {
@@ -808,7 +875,6 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
             });
         }
 
-        let protector = if protect { Some(this.frame().extra.call_id) } else { None };
         trace!(
             "reborrow: {} reference {:?} derived from {:?} (pointee {}): {:?}, size {}",
             kind,
@@ -819,6 +885,22 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
             size.bytes()
         );
 
+        // FIXME: This just protects everything, which is wrong. At the very least, we should not
+        // protect anything that contains an UnsafeCell.
+        if protect {
+            this.frame_mut().extra.protected_tags.push(new_tag);
+            let call_id = this.frame().extra.call_id;
+            this.machine
+                .stacked_borrows
+                .as_mut()
+                .unwrap()
+                .get_mut()
+                .protected_tags
+                .insert(new_tag, call_id);
+        }
+        // FIXME: can't hold the current span handle across the borrows of self above
+        let current_span = &mut this.machine.current_span();
+
         // Update the stacks.
         // Make sure that raw pointers and mutable shared references are reborrowed "weak":
         // There could be existing unique pointers reborrowed from them that should remain valid!
@@ -855,14 +937,14 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
                     } else {
                         Permission::SharedReadWrite
                     };
-                    let protector = if frozen {
-                        protector
+                    let protected = if frozen {
+                        protect
                     } else {
                         // We do not protect inside UnsafeCell.
                         // This fixes https://github.com/rust-lang/rust/issues/55005.
-                        None
+                        false
                     };
-                    let item = Item { perm, tag: new_tag, protector };
+                    let item = Item::new(new_tag, perm, protected);
                     let mut global = this.machine.stacked_borrows.as_ref().unwrap().borrow_mut();
                     stacked_borrows.for_each(range, |offset, stack, history, exposed_tags| {
                         stack.grant(
@@ -888,7 +970,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
             .as_mut()
             .expect("we should have Stacked Borrows data")
             .borrow_mut();
-        let item = Item { perm, tag: new_tag, protector };
+        let item = Item::new(new_tag, perm, protect);
         let range = alloc_range(base_offset, size);
         let mut global = machine.stacked_borrows.as_ref().unwrap().borrow_mut();
         let current_span = &mut machine.current_span(); // `get_alloc_extra_mut` invalidated our old `current_span`

src/stacked_borrows/diagnostics.rs

@@ -141,7 +141,7 @@ impl AllocHistory {
     ) -> InterpError<'tcx> {
         let action = format!(
             "trying to reborrow {derived_from:?} for {:?} permission at {alloc_id:?}[{:#x}]",
-            new.perm,
+            new.perm(),
             error_offset.bytes(),
         );
         err_sb_ub(
@@ -185,7 +185,7 @@ fn error_cause(stack: &Stack, tag: SbTagExtra) -> &'static str {
     if let SbTagExtra::Concrete(tag) = tag {
         if (0..stack.len())
             .map(|i| stack.get(i).unwrap())
-            .any(|item| item.tag == tag && item.perm != Permission::Disabled)
+            .any(|item| item.tag() == tag && item.perm() != Permission::Disabled)
         {
             ", but that tag only grants SharedReadOnly permission for this location"
         } else {