Store protectors outside Item, pack Tag and Perm

saethlin · saethlin · commit 08e0754f2efb · 2022-07-03T01:05:28.000-04:00
diff --git a/src/machine.rs b/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()
     }
 }
@@ -859,7 +862,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))
     }
 
diff --git a/src/shims/panic.rs b/src/shims/panic.rs
@@ -124,7 +124,11 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
 
         trace!("handle_stack_pop(extra = {:?}, unwinding = {})", extra, unwinding);
         if let Some(stacked_borrows) = &this.machine.stacked_borrows {
-            stacked_borrows.borrow_mut().end_call(extra.call_id);
+            let mut sb = stacked_borrows.borrow_mut();
+            sb.end_call(extra.call_id);
+            for tag in extra.protected_tags.drain(..) {
+                sb.protected_tags.remove(&tag);
+            }
         }
 
         // We only care about `catch_panic` if we're unwinding - if we're doing a normal
diff --git a/src/stacked_borrows.rs b/src/stacked_borrows.rs
@@ -99,18 +99,11 @@ pub struct Item {
     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 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)?;
-        }
-        write!(f, "]")?;
-        Ok(())
+        write!(f, "[{:?} for {:?}]", self.perm, self.tag)
     }
 }
 
@@ -138,6 +131,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: FxHashSet<SbTag>,
     /// The pointer ids to trace
     tracked_pointer_tags: HashSet<SbTag>,
     /// The call ids to trace
@@ -202,6 +197,7 @@ impl GlobalStateInner {
             base_ptr_tags: FxHashMap::default(),
             next_call_id: NonZeroU64::new(1).unwrap(),
             active_calls: FxHashSet::default(),
+            protected_tags: FxHashSet::default(),
             tracked_pointer_tags,
             tracked_call_ids,
             retag_fields,
@@ -232,10 +228,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();
@@ -332,24 +324,22 @@ impl<'tcx> Stack {
             ));
         }
 
-        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 global.protected_tags.contains(&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: {:?}",
+                        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,
+                ))?
             }
         }
         Ok(())
@@ -577,7 +567,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 { perm, tag };
         let stack = Stack::new(item);
 
         Stacks {
@@ -792,7 +782,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,
@@ -803,6 +792,15 @@ 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);
+            this.machine.stacked_borrows.as_mut().unwrap().get_mut().protected_tags.insert(new_tag);
+        }
+        // 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!
@@ -839,14 +837,7 @@ trait EvalContextPrivExt<'mir, 'tcx: 'mir>: crate::MiriEvalContextExt<'mir, 'tcx
                     } else {
                         Permission::SharedReadWrite
                     };
-                    let protector = if frozen {
-                        protector
-                    } else {
-                        // We do not protect inside UnsafeCell.
-                        // This fixes https://github.com/rust-lang/rust/issues/55005.
-                        None
-                    };
-                    let item = Item { perm, tag: new_tag, protector };
+                    let item = Item { perm, tag: new_tag };
                     let mut global = this.machine.stacked_borrows.as_ref().unwrap().borrow_mut();
                     stacked_borrows.for_each(range, |offset, stack, history, exposed_tags| {
                         stack.grant(
@@ -872,7 +863,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 { perm, tag: new_tag };
         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`
diff --git a/src/stacked_borrows/stack.rs b/src/stacked_borrows/stack.rs
