Merge pull request #4540 from RalfJung/tb-refactors

tree borrows: refactor new-permission logic

6 files changed, 116 insertions, 144 deletions

diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs
index ad2a67160f4..04b951cd77c 100644
--- a/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs
+++ b/src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs
@@ -125,81 +125,64 @@ pub struct NewPermission {
     /// Whether a read access should be performed on the non-frozen
     /// part on a retag.
     nonfreeze_access: bool,
+    /// Permission for memory outside the range.
+    outside_perm: Permission,
     /// Whether this pointer is part of the arguments of a function call.
     /// `protector` is `Some(_)` for all pointers marked `noalias`.
     protector: Option<ProtectorKind>,
 }
 
 impl<'tcx> NewPermission {
-    /// Determine NewPermission of the reference from the type of the pointee.
-    fn from_ref_ty(
+    /// Determine NewPermission of the reference/Box from the type of the pointee.
+    ///
+    /// A `ref_mutability` of `None` indicates a `Box` type.
+    fn new(
         pointee: Ty<'tcx>,
-        mutability: Mutability,
-        kind: RetagKind,
+        ref_mutability: Option<Mutability>,
+        retag_kind: RetagKind,
         cx: &crate::MiriInterpCx<'tcx>,
     ) -> Option<Self> {
         let ty_is_unpin = pointee.is_unpin(*cx.tcx, cx.typing_env());
-        let is_protected = kind == RetagKind::FnEntry;
-        let protector = is_protected.then_some(ProtectorKind::StrongProtector);
-
-        Some(match mutability {
-            Mutability::Mut if ty_is_unpin =>
-                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 ty_is_freeze = pointee.is_freeze(*cx.tcx, cx.typing_env());
+        let is_protected = retag_kind == RetagKind::FnEntry;
 
-    /// Compute permission for `Box`-like type (`Box` always, and also `Unique` if enabled).
-    /// These pointers allow deallocation so need a different kind of protector not handled
-    /// by `from_ref_ty`.
-    fn from_unique_ty(
-        ty: Ty<'tcx>,
-        kind: RetagKind,
-        cx: &crate::MiriInterpCx<'tcx>,
-    ) -> Option<Self> {
-        let pointee = ty.builtin_deref(true).unwrap();
-        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 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,
-            }
+        if matches!(ref_mutability, Some(Mutability::Mut) | None if !ty_is_unpin) {
+            // Mutable reference / Box to pinning type: retagging is a NOP.
+            // FIXME: with `UnsafePinned`, this should do proper per-byte tracking.
+            return None;
+        }
+
+        let freeze_perm = match ref_mutability {
+            // Shared references are frozen.
+            Some(Mutability::Not) => Permission::new_frozen(),
+            // Mutable references and Boxes are reserved.
+            _ => Permission::new_reserved_frz(),
+        };
+        let nonfreeze_perm = match ref_mutability {
+            // Shared references are "transparent".
+            Some(Mutability::Not) => Permission::new_cell(),
+            // *Protected* mutable references and boxes are reserved without regarding for interior mutability.
+            _ if is_protected => Permission::new_reserved_frz(),
+            // Unprotected mutable references and boxes start in `ReservedIm`.
+            _ => Permission::new_reserved_im(),
+        };
+
+        // Everything except for `Cell` gets an initial access.
+        let initial_access = |perm: &Permission| !perm.is_cell();
+
+        Some(NewPermission {
+            freeze_perm,
+            freeze_access: initial_access(&freeze_perm),
+            nonfreeze_perm,
+            nonfreeze_access: initial_access(&nonfreeze_perm),
+            outside_perm: if ty_is_freeze { freeze_perm } else { nonfreeze_perm },
+            protector: is_protected.then_some(if ref_mutability.is_some() {
+                // Strong protector for references
+                ProtectorKind::StrongProtector
+            } else {
+                // Weak protector for boxes
+                ProtectorKind::WeakProtector
+            }),
         })
     }
 }
@@ -313,16 +296,6 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
         let span = this.machine.current_span();
 
-        // Store initial permissions and their corresponding range.
-        let mut perms_map: DedupRangeMap<LocationState> = DedupRangeMap::new(
-            ptr_size,
-            LocationState::new_accessed(Permission::new_disabled(), IdempotentForeignAccess::None), // this will be overwritten
-        );
-        // Keep track of whether the node has any part that allows for interior mutability.
-        // FIXME: This misses `PhantomData<UnsafeCell<T>>` which could be considered a marker
-        // for requesting interior mutability.
-        let mut has_unsafe_cell = false;
-
         // When adding a new node, the SIFA of its parents needs to be updated, potentially across
         // the entire memory range. For the parts that are being accessed below, the access itself
         // trivially takes care of that. However, we have to do some more work to also deal with
@@ -350,66 +323,48 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             .get_tree_borrows_params()
             .precise_interior_mut;
 
-        let default_perm = if !precise_interior_mut {
-            // NOTE: Using `ty_is_freeze` doesn't give the same result as going through the range
-            // and computing `has_unsafe_cell`.  This is because of zero-sized `UnsafeCell`, for which
-            // `has_unsafe_cell` is false, but `!ty_is_freeze` is true.
-            let ty_is_freeze = place.layout.ty.is_freeze(*this.tcx, this.typing_env());
-            let (perm, access) = if ty_is_freeze {
+        // Compute initial "inside" permissions.
+        let loc_state = |frozen: bool| -> LocationState {
+            let (perm, access) = if frozen {
                 (new_perm.freeze_perm, new_perm.freeze_access)
             } else {
                 (new_perm.nonfreeze_perm, new_perm.nonfreeze_access)
             };
             let sifa = perm.strongest_idempotent_foreign_access(protected);
-            let new_loc = if access {
+            if access {
                 LocationState::new_accessed(perm, sifa)
             } else {
                 LocationState::new_non_accessed(perm, sifa)
-            };
-
-            for (_loc_range, loc) in perms_map.iter_mut_all() {
-                *loc = new_loc;
             }
-
-            perm
+        };
+        let perms_map = if !precise_interior_mut {
+            // For `!Freeze` types, just pretend the entire thing is an `UnsafeCell`.
+            let ty_is_freeze = place.layout.ty.is_freeze(*this.tcx, this.typing_env());
+            let state = loc_state(ty_is_freeze);
+            DedupRangeMap::new(ptr_size, state)
         } else {
+            // The initial state will be overwritten by the visitor below.
+            let mut perms_map: DedupRangeMap<LocationState> = DedupRangeMap::new(
+                ptr_size,
+                LocationState::new_accessed(
+                    Permission::new_disabled(),
+                    IdempotentForeignAccess::None,
+                ),
+            );
             this.visit_freeze_sensitive(place, ptr_size, |range, frozen| {
-                has_unsafe_cell = has_unsafe_cell || !frozen;
-
-                // 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 sifa = perm.strongest_idempotent_foreign_access(protected);
-                // NOTE: Currently, `access` is false if and only if `perm` is Cell, so this `if`
-                // doesn't not change whether any code is UB or not. We could just always use
-                // `new_accessed` and everything would stay the same. But that seems conceptually
-                // odd, so we keep the initial "accessed" bit of the `LocationState` in sync with whether
-                // a read access is performed below.
-                let new_loc = if access {
-                    LocationState::new_accessed(perm, sifa)
-                } else {
-                    LocationState::new_non_accessed(perm, sifa)
-                };
-
-                // Store initial permissions.
+                let state = loc_state(frozen);
                 for (_loc_range, loc) in perms_map.iter_mut(range.start, range.size) {
-                    *loc = new_loc;
+                    *loc = state;
                 }
-
                 interp_ok(())
             })?;
-
-            // Allow lazily writing to surrounding data if we found an `UnsafeCell`.
-            if has_unsafe_cell { new_perm.nonfreeze_perm } else { new_perm.freeze_perm }
+            perms_map
         };
 
         let alloc_extra = this.get_alloc_extra(alloc_id)?;
         let mut tree_borrows = alloc_extra.borrow_tracker_tb().borrow_mut();
 
-        for (perm_range, perm) in perms_map.iter_mut_all() {
+        for (perm_range, perm) in perms_map.iter_all() {
             if perm.is_accessed() {
                 // Some reborrows incur a read access to the parent.
                 // Adjust range to be relative to allocation start (rather than to `place`).
@@ -447,7 +402,7 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             orig_tag,
             new_tag,
             perms_map,
-            default_perm,
+            new_perm.outside_perm,
             protected,
             span,
         )?;
@@ -514,7 +469,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let this = self.eval_context_mut();
         let new_perm = match val.layout.ty.kind() {
             &ty::Ref(_, pointee, mutability) =>
-                NewPermission::from_ref_ty(pointee, mutability, kind, this),
+                NewPermission::new(pointee, Some(mutability), kind, this),
             _ => None,
         };
         if let Some(new_perm) = new_perm {
@@ -571,8 +526,9 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             fn visit_box(&mut self, box_ty: Ty<'tcx>, place: &PlaceTy<'tcx>) -> InterpResult<'tcx> {
                 // Only boxes for the global allocator get any special treatment.
                 if box_ty.is_box_global(*self.ecx.tcx) {
+                    let pointee = place.layout.ty.builtin_deref(true).unwrap();
                     let new_perm =
-                        NewPermission::from_unique_ty(place.layout.ty, self.kind, self.ecx);
+                        NewPermission::new(pointee, /* not a ref */ None, self.kind, self.ecx);
                     self.retag_ptr_inplace(place, new_perm)?;
                 }
                 interp_ok(())
@@ -591,7 +547,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 match place.layout.ty.kind() {
                     &ty::Ref(_, pointee, mutability) => {
                         let new_perm =
-                            NewPermission::from_ref_ty(pointee, mutability, self.kind, self.ecx);
+                            NewPermission::new(pointee, Some(mutability), self.kind, self.ecx);
                         self.retag_ptr_inplace(place, new_perm)?;
                     }
                     ty::RawPtr(_, _) => {
@@ -643,14 +599,11 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             // 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_perm: Permission::new_reserved_frz(),
             freeze_access: true,
-            nonfreeze_perm: Permission::new_reserved(
-                /* ty_is_freeze */ false, /* protected */ true,
-            ),
+            nonfreeze_perm: Permission::new_reserved_frz(),
             nonfreeze_access: true,
+            outside_perm: Permission::new_reserved_frz(),
             protector: Some(ProtectorKind::StrongProtector),
         };
         this.tb_retag_place(place, new_perm)
diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/perms.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/perms.rs
index 38863ca0734..390435e58d1 100644
--- a/src/tools/miri/src/borrow_tracker/tree_borrows/perms.rs
+++ b/src/tools/miri/src/borrow_tracker/tree_borrows/perms.rs
@@ -272,28 +272,15 @@ impl Permission {
 
     /// Default initial permission of a reborrowed mutable reference that is either
     /// protected or not interior mutable.
-    fn new_reserved_frz() -> Self {
+    pub fn new_reserved_frz() -> Self {
         Self { inner: ReservedFrz { conflicted: false } }
     }
 
     /// Default initial permission of an unprotected interior mutable reference.
-    fn new_reserved_im() -> Self {
+    pub fn new_reserved_im() -> Self {
         Self { inner: ReservedIM }
     }
 
-    /// Wrapper around `new_reserved_frz` and `new_reserved_im` that decides
-    /// which to call based on the interior mutability and the retag kind (whether there
-    /// 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() }
-    }
-
     /// Default initial permission of a reborrowed shared reference.
     pub fn new_frozen() -> Self {
         Self { inner: Frozen }
diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/tree/tests.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/tree/tests.rs
index bb3fc2d80b3..d9b3696e4f8 100644
--- a/src/tools/miri/src/borrow_tracker/tree_borrows/tree/tests.rs
+++ b/src/tools/miri/src/borrow_tracker/tree_borrows/tree/tests.rs
@@ -610,7 +610,7 @@ mod spurious_read {
             },
             y: LocStateProt {
                 state: LocationState::new_non_accessed(
-                    Permission::new_reserved(/* freeze */ true, /* protected */ true),
+                    Permission::new_reserved_frz(),
                     IdempotentForeignAccess::default(),
                 ),
                 prot: true,
diff --git a/src/tools/miri/tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.rs b/src/tools/miri/tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.rs
new file mode 100644
index 00000000000..7d51050f32b
--- /dev/null
+++ b/src/tools/miri/tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.rs
@@ -0,0 +1,9 @@
+//@compile-flags: -Zmiri-tree-borrows
+
+fn main() {
+    // Since the "inside" part is `!Freeze`, the permission to mutate is gone.
+    let pair = ((), 1);
+    let x = &pair.0;
+    let ptr = (&raw const *x).cast::<i32>().cast_mut();
+    unsafe { ptr.write(0) }; //~ERROR: /write access .* forbidden/
+}
diff --git a/src/tools/miri/tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.stderr b/src/tools/miri/tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.stderr
new file mode 100644
index 00000000000..e9800468c57
--- /dev/null
+++ b/src/tools/miri/tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.stderr
@@ -0,0 +1,21 @@
+error: Undefined Behavior: write access through <TAG> at ALLOC[0x0] is forbidden
+  --> tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.rs:LL:CC
+   |
+LL |     unsafe { ptr.write(0) };
+   |              ^^^^^^^^^^^^ Undefined Behavior occurred here
+   |
+   = help: this indicates a potential bug in the program: it performed an invalid operation, but the Tree Borrows rules it violated are still experimental
+   = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/tree-borrows.md for further information
+   = help: the accessed tag <TAG> has state Frozen which forbids this child write access
+help: the accessed tag <TAG> was created here, in the initial state Frozen
+  --> tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.rs:LL:CC
+   |
+LL |     let x = &pair.0;
+   |             ^^^^^^^
+   = note: BACKTRACE (of the first span):
+   = note: inside `main` at tests/fail/tree_borrows/frozen-lazy-write-to-surrounding.rs:LL:CC
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+
diff --git a/src/tools/miri/tests/pass/tree_borrows/cell-lazy-write-to-surrounding.rs b/src/tools/miri/tests/pass/tree_borrows/cell-lazy-write-to-surrounding.rs
index abe08f2cd22..7352784ac7a 100644
--- a/src/tools/miri/tests/pass/tree_borrows/cell-lazy-write-to-surrounding.rs
+++ b/src/tools/miri/tests/pass/tree_borrows/cell-lazy-write-to-surrounding.rs
@@ -14,9 +14,11 @@ fn main() {
     foo(&arr[0]);
 
     let pair = (Cell::new(1), 1);
-    // TODO: Ideally, this would result in UB since the second element
-    // in `pair` is Frozen.  We would need some way to express a
-    // "shared reference with permission to access surrounding
-    // interior mutable data".
     foo(&pair.0);
+
+    // As long as the "inside" part is `!Freeze`, the permission to mutate the "outside" is preserved.
+    let pair = (Cell::new(()), 1);
+    let x = &pair.0;
+    let ptr = (&raw const *x).cast::<i32>().cast_mut();
+    unsafe { ptr.write(0) };
 }