TB: Track permissions on the byte-level

src/borrow_tracker/tree_borrows/mod.rs

@@ -3,6 +3,7 @@ use rustc_middle::mir::{Mutability, RetagKind};
 use rustc_middle::ty::layout::HasTypingEnv;
 use rustc_middle::ty::{self, Ty};
 
+use self::tree::LocationState;
 use crate::borrow_tracker::{GlobalState, GlobalStateInner, ProtectorKind};
 use crate::concurrency::data_race::NaReadType;
 use crate::*;
@@ -113,16 +114,19 @@ impl<'tcx> Tree {
 
 /// Policy for a new borrow.
 #[derive(Debug, Clone, Copy)]
-struct NewPermission {
-    /// Which permission should the pointer start with.
-    initial_state: Permission,
+pub struct NewPermission {
+    /// Permission for the frozen part of the range.
+    freeze_perm: Permission,
+    /// Whether a read access should be performed on the frozen part on a retag.
+    freeze_access: bool,
+    /// Permission for the non-frozen part of the range.
+    nonfreeze_perm: Permission,
+    /// Whether a read access should be performed on the non-frozen
+    /// part on a retag.
+    nonfreeze_access: bool,
     /// Whether this pointer is part of the arguments of a function call.
     /// `protector` is `Some(_)` for all pointers marked `noalias`.
     protector: Option<ProtectorKind>,
-    /// Whether a read should be performed on a retag.  This should be `false`
-    /// for `Cell` because this could cause data races when using thread-safe
-    /// data types like `Mutex<T>`.
-    initial_read: bool,
 }
 
 impl<'tcx> NewPermission {
@@ -133,27 +137,42 @@ impl<'tcx> NewPermission {
         kind: RetagKind,
         cx: &crate::MiriInterpCx<'tcx>,
     ) -> Option<Self> {
-        let ty_is_freeze = pointee.is_freeze(*cx.tcx, cx.typing_env());
         let ty_is_unpin = pointee.is_unpin(*cx.tcx, cx.typing_env());
         let is_protected = kind == RetagKind::FnEntry;
-        // As demonstrated by `tests/fail/tree_borrows/reservedim_spurious_write.rs`,
-        // interior mutability and protectors interact poorly.
-        // To eliminate the case of Protected Reserved IM we override interior mutability
-        // in the case of a protected reference: protected references are always considered
-        // "freeze" in their reservation phase.
-        let (initial_state, initial_read) = match mutability {
+        let protector = is_protected.then_some(ProtectorKind::StrongProtector);
+
+        Some(match mutability {
             Mutability::Mut if ty_is_unpin =>
-                (Permission::new_reserved(ty_is_freeze, is_protected), true),
-            Mutability::Not if ty_is_freeze => (Permission::new_frozen(), true),
-            Mutability::Not if !ty_is_freeze => (Permission::new_cell(), false),
-            // Raw pointers never enter this function so they are not handled.
-            // However raw pointers are not the only pointers that take the parent
-            // tag, this also happens for `!Unpin` `&mut`s, which are excluded above.
+                NewPermission {
+                    freeze_perm: Permission::new_reserved(
+                        /* ty_is_freeze */ true,
+                        is_protected,
+                    ),
+                    freeze_access: true,
+                    nonfreeze_perm: Permission::new_reserved(
+                        /* ty_is_freeze */ false,
+                        is_protected,
+                    ),
+                    // If we have a mutable reference, then the non-frozen part will
+                    // have state `ReservedIM` or `Reserved`, which can have an initial read access
+                    // performed on it because you cannot have multiple mutable borrows.
+                    nonfreeze_access: true,
+                    protector,
+                },
+            Mutability::Not =>
+                NewPermission {
+                    freeze_perm: Permission::new_frozen(),
+                    freeze_access: true,
+                    nonfreeze_perm: Permission::new_cell(),
+                    // If it is a shared reference, then the non-frozen
+                    // part will have state `Cell`, which should not have an initial access,
+                    // as this can cause data races when using thread-safe data types like
+                    // `Mutex<T>`.
+                    nonfreeze_access: false,
+                    protector,
+                },
             _ => return None,
-        };
-
-        let protector = is_protected.then_some(ProtectorKind::StrongProtector);
-        Some(Self { initial_state, protector, initial_read })
+        })
     }
 
     /// Compute permission for `Box`-like type (`Box` always, and also `Unique` if enabled).
@@ -168,13 +187,17 @@ impl<'tcx> NewPermission {
         pointee.is_unpin(*cx.tcx, cx.typing_env()).then_some(()).map(|()| {
             // Regular `Unpin` box, give it `noalias` but only a weak protector
             // because it is valid to deallocate it within the function.
-            let ty_is_freeze = pointee.is_freeze(*cx.tcx, cx.typing_env());
-            let protected = kind == RetagKind::FnEntry;
-            let initial_state = Permission::new_reserved(ty_is_freeze, protected);
-            Self {
-                initial_state,
-                protector: protected.then_some(ProtectorKind::WeakProtector),
-                initial_read: true,
+            let is_protected = kind == RetagKind::FnEntry;
+            let protector = is_protected.then_some(ProtectorKind::WeakProtector);
+            NewPermission {
+                freeze_perm: Permission::new_reserved(/* ty_is_freeze */ true, is_protected),
+                freeze_access: true,
+                nonfreeze_perm: Permission::new_reserved(
+                    /* ty_is_freeze */ false,
+                    is_protected,
+                ),
+                nonfreeze_access: true,
+                protector,
             }
         })
     }
@@ -193,9 +216,9 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         new_perm: NewPermission,
         new_tag: BorTag,
     ) -> InterpResult<'tcx, Option<Provenance>> {
+        // Make sure the new permissions make sense as the initial permissions of a fresh tag.
+        assert!(new_perm.freeze_perm.is_initial() && new_perm.nonfreeze_perm.is_initial());
         let this = self.eval_context_mut();
-        // Make sure the new permission makes sense as the initial permission of a fresh tag.
-        assert!(new_perm.initial_state.is_initial());
         // Ensure we bail out if the pointer goes out-of-bounds (see miri#1050).
         this.check_ptr_access(place.ptr(), ptr_size, CheckInAllocMsg::Dereferenceable)?;
 
@@ -206,7 +229,13 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             let global = this.machine.borrow_tracker.as_ref().unwrap().borrow();
             let ty = place.layout.ty;
             if global.tracked_pointer_tags.contains(&new_tag) {
-                let kind_str = format!("initial state {} (pointee type {ty})", new_perm.initial_state);
+                 let ty_is_freeze = ty.is_freeze(*this.tcx, this.typing_env());
+                 let kind_str =
+                     if ty_is_freeze {
+                         format!("initial state {} (pointee type {ty})", new_perm.freeze_perm)
+                     } else {
+                         format!("initial state {}/{} outside/inside UnsafeCell (pointee type {ty})", new_perm.freeze_perm, new_perm.nonfreeze_perm)
+                     };
                 this.emit_diagnostic(NonHaltingDiagnostic::CreatedPointerTag(
                     new_tag.inner(),
                     Some(kind_str),
@@ -285,43 +314,80 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
         let span = this.machine.current_span();
         let alloc_extra = this.get_alloc_extra(alloc_id)?;
-        let range = alloc_range(base_offset, ptr_size);
         let mut tree_borrows = alloc_extra.borrow_tracker_tb().borrow_mut();
+        let prot = new_perm.protector.is_some();
+
+        // Store initial permissions and their corresponding range.
+        let mut perms_map: RangeMap<LocationState> = RangeMap::new(
+            ptr_size,
+            LocationState::new_uninit(
+                Permission::new_disabled(),
+                foreign_access_skipping::IdempotentForeignAccess::None,
+            ),
+        );
+        // Keep track of whether the node has any part that allows for interior mutability.
+        let mut has_unsafe_cell = false;
+
+        this.visit_freeze_sensitive(place, ptr_size, |range, frozen| {
+            has_unsafe_cell = has_unsafe_cell || !frozen;
+
+            // Adjust range.
+            let mut range_with_offset = range;
+            range_with_offset.start += base_offset;
+
+            // We are only ever `Frozen` inside the frozen bits.
+            let (perm, access) = if frozen {
+                (new_perm.freeze_perm, new_perm.freeze_access)
+            } else {
+                (new_perm.nonfreeze_perm, new_perm.nonfreeze_access)
+            };
+            let strongest_idempotent = perm.strongest_idempotent_foreign_access(prot);
+
+            // Store initial permissions.
+            let perm_init = LocationState::new_init(perm, strongest_idempotent);
+            for (_perm_range, perm) in perms_map.iter_mut(range.start, range.size) {
+                *perm = perm_init;
+            }
+
+            // Some reborrows incur a read access to the parent.
+            if access {
+                tree_borrows.perform_access(
+                    orig_tag,
+                    Some((range_with_offset, AccessKind::Read, diagnostics::AccessCause::Reborrow)),
+                    this.machine.borrow_tracker.as_ref().unwrap(),
+                    alloc_id,
+                    this.machine.current_span(),
+                )?
+            }
+
+            // Also inform the data race model (but only if any bytes are actually affected).
+            if range.size.bytes() > 0 && access {
+                if let Some(data_race) = alloc_extra.data_race.as_vclocks_ref() {
+                    data_race.read(
+                        alloc_id,
+                        range_with_offset,
+                        NaReadType::Retag,
+                        Some(place.layout.ty),
+                        &this.machine,
+                    )?
+                }
+            }
+            interp_ok(())
+        })?;
 
-        // All reborrows incur a (possibly zero-sized) read access to the parent
-        if new_perm.initial_read {
-            tree_borrows.perform_access(
-                orig_tag,
-                Some((range, AccessKind::Read, diagnostics::AccessCause::Reborrow)),
-                this.machine.borrow_tracker.as_ref().unwrap(),
-                alloc_id,
-                this.machine.current_span(),
-            )?;
-        }
         // Record the parent-child pair in the tree.
         tree_borrows.new_child(
+            ptr_size,
+            base_offset,
             orig_tag,
             new_tag,
-            new_perm.initial_state,
-            range,
+            new_perm,
+            perms_map,
             span,
-            new_perm.protector.is_some(),
+            has_unsafe_cell,
         )?;
         drop(tree_borrows);
 
-        // Also inform the data race model (but only if any bytes are actually affected).
-        if range.size.bytes() > 0 && new_perm.initial_read {
-            if let Some(data_race) = alloc_extra.data_race.as_vclocks_ref() {
-                data_race.read(
-                    alloc_id,
-                    range,
-                    NaReadType::Retag,
-                    Some(place.layout.ty),
-                    &this.machine,
-                )?;
-            }
-        }
-
         interp_ok(Some(Provenance::Concrete { alloc_id, tag: new_tag }))
     }
 
@@ -508,15 +574,21 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn tb_protect_place(&mut self, place: &MPlaceTy<'tcx>) -> InterpResult<'tcx, MPlaceTy<'tcx>> {
         let this = self.eval_context_mut();
 
-        // Note: if we were to inline `new_reserved` below we would find out that
-        // `ty_is_freeze` is eventually unused because it appears in a `ty_is_freeze || true`.
-        // We are nevertheless including it here for clarity.
-        let ty_is_freeze = place.layout.ty.is_freeze(*this.tcx, this.typing_env());
         // Retag it. With protection! That is the entire point.
         let new_perm = NewPermission {
-            initial_state: Permission::new_reserved(ty_is_freeze, /* protected */ true),
+            // Note: If we are creating a protected Reserved, which can
+            // never be ReservedIM, the value of the `ty_is_freeze`
+            // argument doesn't matter
+            // (`ty_is_freeze || true` in `new_reserved` will always be `true`).
+            freeze_perm: Permission::new_reserved(
+                /* ty_is_freeze */ true, /* protected */ true,
+            ),
+            freeze_access: true,
+            nonfreeze_perm: Permission::new_reserved(
+                /* ty_is_freeze */ false, /* protected */ true,
+            ),
+            nonfreeze_access: true,
             protector: Some(ProtectorKind::StrongProtector),
-            initial_read: true,
         };
         this.tb_retag_place(place, new_perm)
     }

src/borrow_tracker/tree_borrows/perms.rs

@@ -94,6 +94,7 @@ impl PermissionPriv {
     }
 
     /// Reject `ReservedIM` that cannot exist in the presence of a protector.
+    #[cfg(test)]
     fn compatible_with_protector(&self) -> bool {
         // FIXME(TB-Cell): It is unclear what to do here.
         // `Cell` will occur with a protector but won't provide the guarantees
@@ -253,10 +254,6 @@ impl Permission {
     pub fn is_disabled(&self) -> bool {
         self.inner == Disabled
     }
-    /// Check if `self` is the post-child-write state of a pointer (is `Active`).
-    pub fn is_active(&self) -> bool {
-        self.inner == Active
-    }
     /// Check if `self` is the never-allow-writes-again state of a pointer (is `Frozen`).
     pub fn is_frozen(&self) -> bool {
         self.inner == Frozen
@@ -289,6 +286,11 @@ impl Permission {
     /// is a protector is relevant because being protected takes priority over being
     /// interior mutable)
     pub fn new_reserved(ty_is_freeze: bool, protected: bool) -> Self {
+        // As demonstrated by `tests/fail/tree_borrows/reservedim_spurious_write.rs`,
+        // interior mutability and protectors interact poorly.
+        // To eliminate the case of Protected Reserved IM we override interior mutability
+        // in the case of a protected reference: protected references are always considered
+        // "freeze" in their reservation phase.
         if ty_is_freeze || protected { Self::new_reserved_frz() } else { Self::new_reserved_im() }
     }
 
@@ -309,6 +311,7 @@ impl Permission {
     }
 
     /// Reject `ReservedIM` that cannot exist in the presence of a protector.
+    #[cfg(test)]
     pub fn compatible_with_protector(&self) -> bool {
         self.inner.compatible_with_protector()
     }
@@ -393,11 +396,6 @@ impl PermTransition {
         self.from <= self.to
     }
 
-    pub fn from(from: Permission, to: Permission) -> Option<Self> {
-        let t = Self { from: from.inner, to: to.inner };
-        t.is_possible().then_some(t)
-    }
-
     pub fn is_noop(self) -> bool {
         self.from == self.to
     }
@@ -407,11 +405,6 @@ impl PermTransition {
         (starting_point.inner == self.from).then_some(Permission { inner: self.to })
     }
 
-    /// Extract starting point of a transition
-    pub fn started(self) -> Permission {
-        Permission { inner: self.from }
-    }
-
     /// Determines if this transition would disable the permission.
     pub fn produces_disabled(self) -> bool {
         self.to == Disabled