slight simplifications for borrow tracking

3 files changed, 36 insertions, 58 deletions

diff --git a/src/tools/miri/src/borrow_tracker/mod.rs b/src/tools/miri/src/borrow_tracker/mod.rs
index 69f9bd09f26..78332879e5e 100644
--- a/src/tools/miri/src/borrow_tracker/mod.rs
+++ b/src/tools/miri/src/borrow_tracker/mod.rs
@@ -55,9 +55,9 @@ impl fmt::Debug for BorTag {
     }
 }
 
-/// Per-frame data for borrow tracking
+/// Per-call-stack-frame data for borrow tracking
 #[derive(Debug)]
-pub struct FrameExtra {
+pub struct FrameState {
     /// The ID of the call this frame corresponds to.
     pub call_id: CallId,
 
@@ -72,7 +72,7 @@ pub struct FrameExtra {
     pub protected_tags: SmallVec<[BorTag; 2]>,
 }
 
-impl VisitTags for FrameExtra {
+impl VisitTags for FrameState {
     fn visit_tags(&self, _visit: &mut dyn FnMut(BorTag)) {
         // `protected_tags` are fine to GC.
     }
@@ -190,14 +190,14 @@ impl GlobalStateInner {
         id
     }
 
-    pub fn new_frame(&mut self, machine: &MiriMachine<'_, '_>) -> FrameExtra {
+    pub fn new_frame(&mut self, machine: &MiriMachine<'_, '_>) -> FrameState {
         let call_id = self.next_call_id;
         trace!("new_frame: Assigning call ID {}", call_id);
         if self.tracked_call_ids.contains(&call_id) {
             machine.emit_diagnostic(NonHaltingDiagnostic::CreatedCallId(call_id));
         }
         self.next_call_id = NonZeroU64::new(call_id.get() + 1).unwrap();
-        FrameExtra { call_id, protected_tags: SmallVec::new() }
+        FrameState { call_id, protected_tags: SmallVec::new() }
     }
 
     pub fn end_call(&mut self, frame: &machine::FrameData<'_>) {
@@ -253,10 +253,10 @@ impl GlobalStateInner {
         alloc_size: Size,
         kind: MemoryKind<machine::MiriMemoryKind>,
         machine: &MiriMachine<'_, '_>,
-    ) -> AllocExtra {
+    ) -> AllocState {
         match self.borrow_tracker_method {
             BorrowTrackerMethod::StackedBorrows =>
-                AllocExtra::StackedBorrows(Box::new(RefCell::new(Stacks::new_allocation(
+                AllocState::StackedBorrows(Box::new(RefCell::new(Stacks::new_allocation(
                     id, alloc_size, self, kind, machine,
                 )))),
         }
@@ -292,24 +292,30 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
 
 /// Extra per-allocation data for borrow tracking
 #[derive(Debug, Clone)]
-pub enum AllocExtra {
+pub enum AllocState {
     /// Data corresponding to Stacked Borrows
-    StackedBorrows(Box<RefCell<stacked_borrows::AllocExtra>>),
+    StackedBorrows(Box<RefCell<stacked_borrows::AllocState>>),
 }
 
-impl AllocExtra {
-    pub fn assert_sb(&self) -> &RefCell<stacked_borrows::AllocExtra> {
-        match self {
-            AllocExtra::StackedBorrows(ref sb) => sb,
+impl machine::AllocExtra {
+    #[track_caller]
+    pub fn borrow_tracker_sb(&self) -> &RefCell<stacked_borrows::AllocState> {
+        match self.borrow_tracker {
+            Some(AllocState::StackedBorrows(ref sb)) => sb,
+            _ => panic!("expected Stacked Borrows borrow tracking, got something else"),
         }
     }
 
-    pub fn assert_sb_mut(&mut self) -> &mut RefCell<stacked_borrows::AllocExtra> {
-        match self {
-            AllocExtra::StackedBorrows(ref mut sb) => sb,
+    #[track_caller]
+    pub fn borrow_tracker_sb_mut(&mut self) -> &mut RefCell<stacked_borrows::AllocState> {
+        match self.borrow_tracker {
+            Some(AllocState::StackedBorrows(ref mut sb)) => sb,
+            _ => panic!("expected Stacked Borrows borrow tracking, got something else"),
         }
     }
+}
 
+impl AllocState {
     pub fn before_memory_read<'tcx>(
         &self,
         alloc_id: AllocId,
@@ -318,7 +324,7 @@ impl AllocExtra {
         machine: &MiriMachine<'_, 'tcx>,
     ) -> InterpResult<'tcx> {
         match self {
-            AllocExtra::StackedBorrows(sb) =>
+            AllocState::StackedBorrows(sb) =>
                 sb.borrow_mut().before_memory_read(alloc_id, prov_extra, range, machine),
         }
     }
@@ -331,7 +337,7 @@ impl AllocExtra {
         machine: &mut MiriMachine<'_, 'tcx>,
     ) -> InterpResult<'tcx> {
         match self {
-            AllocExtra::StackedBorrows(sb) =>
+            AllocState::StackedBorrows(sb) =>
                 sb.get_mut().before_memory_write(alloc_id, prov_extra, range, machine),
         }
     }
@@ -344,22 +350,22 @@ impl AllocExtra {
         machine: &mut MiriMachine<'_, 'tcx>,
     ) -> InterpResult<'tcx> {
         match self {
-            AllocExtra::StackedBorrows(sb) =>
+            AllocState::StackedBorrows(sb) =>
                 sb.get_mut().before_memory_deallocation(alloc_id, prov_extra, range, machine),
         }
     }
 
     pub fn remove_unreachable_tags(&self, tags: &FxHashSet<BorTag>) {
         match self {
-            AllocExtra::StackedBorrows(sb) => sb.borrow_mut().remove_unreachable_tags(tags),
+            AllocState::StackedBorrows(sb) => sb.borrow_mut().remove_unreachable_tags(tags),
         }
     }
 }
 
-impl VisitTags for AllocExtra {
+impl VisitTags for AllocState {
     fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
         match self {
-            AllocExtra::StackedBorrows(sb) => sb.visit_tags(visit),
+            AllocState::StackedBorrows(sb) => sb.visit_tags(visit),
         }
     }
 }
diff --git a/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs b/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs
index ec3be398a2c..50c2ad75ca7 100644
--- a/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs
+++ b/src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs
@@ -25,7 +25,7 @@ mod stack;
 pub use stack::Stack;
 pub mod diagnostics;
 
-pub type AllocExtra = Stacks;
+pub type AllocState = Stacks;
 
 /// Extra per-allocation state.
 #[derive(Clone, Debug)]
@@ -500,10 +500,6 @@ impl Stacks {
         })?;
         Ok(())
     }
-
-    fn expose_tag(&mut self, tag: BorTag) {
-        self.exposed_tags.insert(tag);
-    }
 }
 
 /// Retagging/reborrowing.  There is some policy in here, such as which permissions
@@ -567,10 +563,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
                     // uncovers a non-supported `extern static`.
                     let extra = this.get_alloc_extra(alloc_id)?;
                     let mut stacked_borrows = extra
-                        .borrow_tracker
-                        .as_ref()
-                        .expect("We should have borrow tracking data")
-                        .assert_sb()
+                        .borrow_tracker_sb()
                         .borrow_mut();
                     // Note that we create a *second* `DiagnosticCxBuilder` below for the actual retag.
                     // FIXME: can this be done cleaner?
@@ -681,12 +674,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
                 // We have to use shared references to alloc/memory_extra here since
                 // `visit_freeze_sensitive` needs to access the global state.
                 let alloc_extra = this.get_alloc_extra(alloc_id)?;
-                let mut stacked_borrows = alloc_extra
-                    .borrow_tracker
-                    .as_ref()
-                    .expect("We should have borrow tracking data")
-                    .assert_sb()
-                    .borrow_mut();
+                let mut stacked_borrows = alloc_extra.borrow_tracker_sb().borrow_mut();
                 this.visit_freeze_sensitive(place, size, |mut range, frozen| {
                     // Adjust range.
                     range.start += base_offset;
@@ -736,12 +724,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
         // Note that this asserts that the allocation is mutable -- but since we are creating a
         // mutable pointer, that seems reasonable.
         let (alloc_extra, machine) = this.get_alloc_extra_mut(alloc_id)?;
-        let stacked_borrows = alloc_extra
-            .borrow_tracker
-            .as_mut()
-            .expect("We should have borrow tracking data")
-            .assert_sb_mut()
-            .get_mut();
+        let stacked_borrows = alloc_extra.borrow_tracker_sb_mut().get_mut();
         let item = Item::new(new_tag, perm, protect.is_some());
         let range = alloc_range(base_offset, size);
         let global = machine.borrow_tracker.as_ref().unwrap().borrow();
@@ -993,13 +976,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 // 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
-                    .as_ref()
-                    .expect("We should have borrow tracking data")
-                    .assert_sb()
-                    .borrow_mut()
-                    .expose_tag(tag);
+                alloc_extra.borrow_tracker_sb().borrow_mut().exposed_tags.insert(tag);
             }
             AllocKind::Function | AllocKind::VTable | AllocKind::Dead => {
                 // No stacked borrows on these allocations.
@@ -1011,12 +988,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
     fn print_stacks(&mut self, alloc_id: AllocId) -> InterpResult<'tcx> {
         let this = self.eval_context_mut();
         let alloc_extra = this.get_alloc_extra(alloc_id)?;
-        let stacks = alloc_extra
-            .borrow_tracker
-            .as_ref()
-            .expect("We should have borrow tracking data")
-            .assert_sb()
-            .borrow();
+        let stacks = alloc_extra.borrow_tracker_sb().borrow();
         for (range, stack) in stacks.stacks.iter_all() {
             print!("{range:?}: [");
             if let Some(bottom) = stack.unknown_bottom() {
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
index 920ee3e1ef1..79106d538b7 100644
--- a/src/tools/miri/src/machine.rs
+++ b/src/tools/miri/src/machine.rs
@@ -39,7 +39,7 @@ pub const STACK_SIZE: u64 = 16 * PAGE_SIZE; // whatever
 /// Extra data stored with each stack frame
 pub struct FrameData<'tcx> {
     /// Extra data for Stacked Borrows.
-    pub borrow_tracker: Option<borrow_tracker::FrameExtra>,
+    pub borrow_tracker: Option<borrow_tracker::FrameState>,
 
     /// If this is Some(), then this is a special "catch unwind" frame (the frame of `try_fn`
     /// called by `try`). When this frame is popped during unwinding a panic,
@@ -255,7 +255,7 @@ impl ProvenanceExtra {
 #[derive(Debug, Clone)]
 pub struct AllocExtra {
     /// Global state of the borrow tracker, if enabled.
-    pub borrow_tracker: Option<borrow_tracker::AllocExtra>,
+    pub borrow_tracker: Option<borrow_tracker::AllocState>,
     /// Data race detection via the use of a vector-clock,
     ///  this is only added if it is enabled.
     pub data_race: Option<data_race::AllocExtra>,