Auto merge of #3013 - RalfJung:tree-borrow-data-races, r=RalfJung

tree borrows: consider some retags as writes for the purpose of data races

Turns out not all retags can be freely reordered. Those that cannot should be considered writes for the data race model, to aid with optimizations.

Only the last commit is the actual change; the ones before that are some refactoring I couldn't stop myself from doing.

Author: bors

Diff for: src/tools/miri/src/borrow_tracker/stacked_borrows/mod.rs

@@ -604,16 +604,15 @@ impl<'mir: 'ecx, 'tcx: 'mir, 'ecx> EvalContextPrivExt<'mir, 'tcx, 'ecx>
 {
 }
 trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'mir, 'tcx> {
-    /// Returns the `AllocId` the reborrow was done in, if some actual borrow stack manipulation
-    /// happened.
+    /// Returns the provenance that should be used henceforth.
     fn sb_reborrow(
         &mut self,
         place: &MPlaceTy<'tcx, Provenance>,
         size: Size,
         new_perm: NewPermission,
         new_tag: BorTag,
         retag_info: RetagInfo, // diagnostics info about this retag
-    ) -> InterpResult<'tcx, Option<AllocId>> {
+    ) -> InterpResult<'tcx, Option<Provenance>> {
         let this = self.eval_context_mut();
         // Ensure we bail out if the pointer goes out-of-bounds (see miri#1050).
         this.check_ptr_access_align(place.ptr, size, Align::ONE, CheckInAllocMsg::InboundsTest)?;
@@ -695,11 +694,14 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
             // pointer tagging for example all calls to get_unchecked on them are invalid.
             if let Ok((alloc_id, base_offset, orig_tag)) = this.ptr_try_get_alloc_id(place.ptr) {
                 log_creation(this, Some((alloc_id, base_offset, orig_tag)))?;
-                return Ok(Some(alloc_id));
+                // Still give it the new provenance, it got retagged after all.
+                return Ok(Some(Provenance::Concrete { alloc_id, tag: new_tag }));
+            } else {
+                // This pointer doesn't come with an AllocId. :shrug:
+                log_creation(this, None)?;
+                // Provenance unchanged.
+                return Ok(place.ptr.provenance);
             }
-            // This pointer doesn't come with an AllocId. :shrug:
-            log_creation(this, None)?;
-            return Ok(None);
         }
 
         let (alloc_id, base_offset, orig_tag) = this.ptr_get_alloc_id(place.ptr)?;
@@ -804,7 +806,7 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
             }
         }
 
-        Ok(Some(alloc_id))
+        Ok(Some(Provenance::Concrete { alloc_id, tag: new_tag }))
     }
 
     /// Retags an individual pointer, returning the retagged version.
@@ -831,25 +833,10 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
         let new_tag = this.machine.borrow_tracker.as_mut().unwrap().get_mut().new_ptr();
 
         // Reborrow.
-        let alloc_id = this.sb_reborrow(&place, size, new_perm, new_tag, info)?;
+        let new_prov = this.sb_reborrow(&place, size, new_perm, new_tag, info)?;
 
         // Adjust pointer.
-        let new_place = place.map_provenance(|p| {
-            p.map(|prov| {
-                match alloc_id {
-                    Some(alloc_id) => {
-                        // If `reborrow` could figure out the AllocId of this ptr, hard-code it into the new one.
-                        // Even if we started out with a wildcard, this newly retagged pointer is tied to that allocation.
-                        Provenance::Concrete { alloc_id, tag: new_tag }
-                    }
-                    None => {
-                        // Looks like this has to stay a wildcard pointer.
-                        assert!(matches!(prov, Provenance::Wildcard));
-                        Provenance::Wildcard
-                    }
-                }
-            })
-        });
+        let new_place = place.map_provenance(|_| new_prov);
 
         // Return new pointer.
         Ok(ImmTy::from_immediate(new_place.to_ref(this), val.layout))

Diff for: src/tools/miri/src/borrow_tracker/tree_borrows/diagnostics.rs

@@ -218,7 +218,7 @@ impl<'tcx> Tree {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy)]
 pub(super) enum TransitionError {
     /// This access is not allowed because some parent tag has insufficient permissions.
     /// For example, if a tag is `Frozen` and encounters a child write this will

Diff for: src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs

@@ -117,7 +117,7 @@ impl<'tcx> NewPermission {
         let ty_is_freeze = pointee.is_freeze(*cx.tcx, cx.param_env());
         let ty_is_unpin = pointee.is_unpin(*cx.tcx, cx.param_env());
         let initial_state = match mutability {
-            Mutability::Mut if ty_is_unpin => Permission::new_unique_2phase(ty_is_freeze),
+            Mutability::Mut if ty_is_unpin => Permission::new_reserved(ty_is_freeze),
             Mutability::Not if ty_is_freeze => Permission::new_frozen(),
             // Raw pointers never enter this function so they are not handled.
             // However raw pointers are not the only pointers that take the parent
@@ -146,7 +146,7 @@ impl<'tcx> NewPermission {
             let ty_is_freeze = ty.is_freeze(*cx.tcx, cx.param_env());
             Self {
                 zero_size,
-                initial_state: Permission::new_unique_2phase(ty_is_freeze),
+                initial_state: Permission::new_reserved(ty_is_freeze),
                 protector: (kind == RetagKind::FnEntry).then_some(ProtectorKind::WeakProtector),
             }
         })
@@ -161,22 +161,14 @@ impl<'mir: 'ecx, 'tcx: 'mir, 'ecx> EvalContextPrivExt<'mir, 'tcx, 'ecx>
 {
 }
 trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'mir, 'tcx> {
-    /// Returns the `AllocId` the reborrow was done in, if there is some actual
-    /// memory associated with this pointer. Returns `None` if there is no actual
-    /// memory allocated. Also checks that the reborrow of size `ptr_size` is
-    /// within bounds of the allocation.
-    ///
-    /// Also returns the tag that the pointer should get, which is essentially
-    /// `if new_perm.is_some() { new_tag } else { parent_tag }` along with
-    /// some logging (always) and fake reads (if `new_perm` is
-    /// `Some(NewPermission { perform_read_access: true }`).
+    /// Returns the provenance that should be used henceforth.
     fn tb_reborrow(
         &mut self,
         place: &MPlaceTy<'tcx, Provenance>, // parent tag extracted from here
         ptr_size: Size,
         new_perm: NewPermission,
         new_tag: BorTag,
-    ) -> InterpResult<'tcx, Option<(AllocId, BorTag)>> {
+    ) -> InterpResult<'tcx, Option<Provenance>> {
         let this = self.eval_context_mut();
         // Ensure we bail out if the pointer goes out-of-bounds (see miri#1050).
         this.check_ptr_access_align(
@@ -222,13 +214,14 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
                     place.layout.ty,
                 );
                 log_creation(this, None)?;
-                return Ok(None);
+                // Keep original provenance.
+                return Ok(place.ptr.provenance);
             }
         };
         log_creation(this, Some((alloc_id, base_offset, parent_prov)))?;
 
         let orig_tag = match parent_prov {
-            ProvenanceExtra::Wildcard => return Ok(None), // TODO: handle wildcard pointers
+            ProvenanceExtra::Wildcard => return Ok(place.ptr.provenance), // TODO: handle wildcard pointers
             ProvenanceExtra::Concrete(tag) => tag,
         };
 
@@ -255,31 +248,54 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
                 .insert(new_tag, protect);
         }
 
+        let alloc_kind = this.get_alloc_info(alloc_id).2;
+        if !matches!(alloc_kind, AllocKind::LiveData) {
+            assert_eq!(ptr_size, Size::ZERO); // we did the deref check above, size has to be 0 here
+            // There's not actually any bytes here where accesses could even be tracked.
+            // Just produce the new provenance, nothing else to do.
+            return Ok(Some(Provenance::Concrete { alloc_id, tag: new_tag }));
+        }
+
         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();
 
         // All reborrows incur a (possibly zero-sized) read access to the parent
-        {
-            let global = &this.machine.borrow_tracker.as_ref().unwrap();
-            let span = this.machine.current_span();
-            tree_borrows.perform_access(
-                AccessKind::Read,
-                orig_tag,
-                range,
-                global,
-                span,
-                diagnostics::AccessCause::Reborrow,
-            )?;
+        tree_borrows.perform_access(
+            AccessKind::Read,
+            orig_tag,
+            range,
+            this.machine.borrow_tracker.as_ref().unwrap(),
+            this.machine.current_span(),
+            diagnostics::AccessCause::Reborrow,
+        )?;
+        // Record the parent-child pair in the tree.
+        tree_borrows.new_child(orig_tag, new_tag, new_perm.initial_state, range, span)?;
+        drop(tree_borrows);
+
+        // Also inform the data race model (but only if any bytes are actually affected).
+        if range.size.bytes() > 0 {
             if let Some(data_race) = alloc_extra.data_race.as_ref() {
-                data_race.read(alloc_id, range, &this.machine)?;
+                // We sometimes need to make it a write, since not all retags commute with reads!
+                // FIXME: Is that truly the semantics we want? Some optimizations are likely to be
+                // very unhappy without this. We'd tsill ge some UB just by picking a suitable
+                // interleaving, but wether UB happens can depend on whether a write occurs in the
+                // future...
+                let is_write = new_perm.initial_state.is_active()
+                    || (new_perm.initial_state.is_resrved() && new_perm.protector.is_some());
+                if is_write {
+                    // Need to get mutable access to alloc_extra.
+                    // (Cannot always do this as we can do read-only reborrowing on read-only allocations.)
+                    let (alloc_extra, machine) = this.get_alloc_extra_mut(alloc_id)?;
+                    alloc_extra.data_race.as_mut().unwrap().write(alloc_id, range, machine)?;
+                } else {
+                    data_race.read(alloc_id, range, &this.machine)?;
+                }
             }
         }
 
-        // Record the parent-child pair in the tree.
-        tree_borrows.new_child(orig_tag, new_tag, new_perm.initial_state, range, span)?;
-        Ok(Some((alloc_id, new_tag)))
+        Ok(Some(Provenance::Concrete { alloc_id, tag: new_tag }))
     }
 
     /// Retags an individual pointer, returning the retagged version.
@@ -315,25 +331,10 @@ trait EvalContextPrivExt<'mir: 'ecx, 'tcx: 'mir, 'ecx>: crate::MiriInterpCxExt<'
         let new_tag = this.machine.borrow_tracker.as_mut().unwrap().get_mut().new_ptr();
 
         // Compute the actual reborrow.
-        let reborrowed = this.tb_reborrow(&place, reborrow_size, new_perm, new_tag)?;
+        let new_prov = this.tb_reborrow(&place, reborrow_size, new_perm, new_tag)?;
 
         // Adjust pointer.
-        let new_place = place.map_provenance(|p| {
-            p.map(|prov| {
-                match reborrowed {
-                    Some((alloc_id, actual_tag)) => {
-                        // If `reborrow` could figure out the AllocId of this ptr, hard-code it into the new one.
-                        // Even if we started out with a wildcard, this newly retagged pointer is tied to that allocation.
-                        Provenance::Concrete { alloc_id, tag: actual_tag }
-                    }
-                    None => {
-                        // Looks like this has to stay a wildcard pointer.
-                        assert!(matches!(prov, Provenance::Wildcard));
-                        Provenance::Wildcard
-                    }
-                }
-            })
-        });
+        let new_place = place.map_provenance(|_| new_prov);
 
         // Return new pointer.
         Ok(ImmTy::from_immediate(new_place.to_ref(this), val.layout))

Diff for: src/tools/miri/src/borrow_tracker/tree_borrows/perms.rs

@@ -134,25 +134,32 @@ pub struct PermTransition {
 
 impl Permission {
     /// Default initial permission of the root of a new tree.
-    pub fn new_root() -> Self {
+    pub fn new_active() -> Self {
         Self { inner: Active }
     }
 
     /// Default initial permission of a reborrowed mutable reference.
-    pub fn new_unique_2phase(ty_is_freeze: bool) -> Self {
+    pub fn new_reserved(ty_is_freeze: bool) -> Self {
         Self { inner: Reserved { ty_is_freeze } }
     }
 
-    /// Default initial permission for return place.
-    pub fn new_active() -> Self {
-        Self { inner: Active }
-    }
-
     /// Default initial permission of a reborrowed shared reference
     pub fn new_frozen() -> Self {
         Self { inner: Frozen }
     }
 
+    pub fn is_active(self) -> bool {
+        matches!(self.inner, Active)
+    }
+
+    pub fn is_resrved(self) -> bool {
+        matches!(self.inner, Reserved { .. })
+    }
+
+    pub fn is_frozen(self) -> bool {
+        matches!(self.inner, Frozen)
+    }
+
     /// Apply the transition to the inner PermissionPriv.
     pub fn perform_access(
         kind: AccessKind,
@@ -438,7 +445,7 @@ mod propagation_optimization_checks {
     }
 
     #[test]
-    fn foreign_read_is_noop_after_write() {
+    fn foreign_read_is_noop_after_foreign_write() {
         use transition::*;
         let old_access = AccessKind::Write;
         let new_access = AccessKind::Read;

Diff for: src/tools/miri/src/borrow_tracker/tree_borrows/tree.rs

@@ -110,7 +110,7 @@ impl LocationState {
 
     // Helper to optimize the tree traversal.
     // The optimization here consists of observing thanks to the tests
-    // `foreign_read_is_noop_after_write` and `all_transitions_idempotent`,
+    // `foreign_read_is_noop_after_foreign_write` and `all_transitions_idempotent`,
     // that there are actually just three possible sequences of events that can occur
     // in between two child accesses that produce different results.
     //
@@ -139,7 +139,7 @@ impl LocationState {
             let new_access_noop = match (self.latest_foreign_access, access_kind) {
                 // Previously applied transition makes the new one a guaranteed
                 // noop in the two following cases:
-                // (1) justified by `foreign_read_is_noop_after_write`
+                // (1) justified by `foreign_read_is_noop_after_foreign_write`
                 (Some(AccessKind::Write), AccessKind::Read) => true,
                 // (2) justified by `all_transitions_idempotent`
                 (Some(old), new) if old == new => true,
@@ -376,7 +376,7 @@ where {
 impl Tree {
     /// Create a new tree, with only a root pointer.
     pub fn new(root_tag: BorTag, size: Size, span: Span) -> Self {
-        let root_perm = Permission::new_root();
+        let root_perm = Permission::new_active();
         let mut tag_mapping = UniKeyMap::default();
         let root_idx = tag_mapping.insert(root_tag);
         let nodes = {
@@ -670,7 +670,8 @@ impl AccessRelatedness {
 mod commutation_tests {
     use super::*;
     impl LocationState {
-        pub fn all_without_access() -> impl Iterator<Item = Self> {
+        pub fn all() -> impl Iterator<Item = Self> {
+            // We keep `latest_foreign_access` at `None` as that's just a cache.
             Permission::all().flat_map(|permission| {
                 [false, true].into_iter().map(move |initialized| {
                     Self { permission, initialized, latest_foreign_access: None }
@@ -695,12 +696,12 @@ mod commutation_tests {
                 // Any protector state works, but we can't move reads across function boundaries
                 // so the two read accesses occur under the same protector.
                 for &protected in &[true, false] {
-                    for loc in LocationState::all_without_access() {
+                    for loc in LocationState::all() {
                         // Apply 1 then 2. Failure here means that there is UB in the source
                         // and we skip the check in the target.
                         let mut loc12 = loc;
-                        let Ok(_) = loc12.perform_access(kind, rel1, protected) else { continue; };
-                        let Ok(_) = loc12.perform_access(kind, rel2, protected) else { continue; };
+                        let Ok(_) = loc12.perform_access(kind, rel1, protected) else { continue };
+                        let Ok(_) = loc12.perform_access(kind, rel2, protected) else { continue };
 
                         // If 1 followed by 2 succeeded, then 2 followed by 1 must also succeed...
                         let mut loc21 = loc;
@@ -718,4 +719,33 @@ mod commutation_tests {
             }
         }
     }
+
+    #[test]
+    #[rustfmt::skip]
+    // Ensure that of 2 accesses happen, one foreign and one a child, and we are protected, that we
+    // get UB unless they are both reads.
+    fn protected_enforces_noalias() {
+        for rel1 in AccessRelatedness::all() {
+            for rel2 in AccessRelatedness::all() {
+                if rel1.is_foreign() == rel2.is_foreign() {
+                    // We want to check pairs of accesses where one is foreign and one is not.
+                    continue;
+                }
+                for kind1 in AccessKind::all() {
+                    for kind2 in AccessKind::all() {
+                        for mut state in LocationState::all() {
+                            let protected = true;
+                            let Ok(_) = state.perform_access(kind1, rel1, protected) else { continue };
+                            let Ok(_) = state.perform_access(kind2, rel2, protected) else { continue };
+                            // If these were both allowed, it must have been two reads.
+                            assert!(
+                                kind1 == AccessKind::Read && kind2 == AccessKind::Read,
+                                "failed to enforce noalias between two accesses that are not both reads"
+                            );
+                        }
+                    }
+                }
+            }
+        }
+    }
 }

Diff for: src/tools/miri/src/machine.rs

@@ -1219,7 +1219,8 @@ impl<'mir, 'tcx> Machine<'mir, 'tcx> for MiriMachine<'mir, 'tcx> {
         // If we have a borrow tracker, we also have it set up protection so that all reads *and
         // writes* during this call are insta-UB.
         if ecx.machine.borrow_tracker.is_some() {
-            if let Either::Left(place) = place.as_mplace_or_local() {
+            // Have to do `to_op` first because a `Place::Local` doesn't imply the local doesn't have an address.
+            if let Either::Left(place) = ecx.place_to_op(place)?.as_mplace_or_imm() {
                 ecx.protect_place(&place)?;
             } else {
                 // Locals that don't have their address taken are as protected as they can ever be.