Properly fix rust-lang#3846 by resetting parents on lazy node creation

This commit supplies a real fix, which makes retags more complicated, at the benefit of
making accesses more performant.

Co-authored-by: Ralf Jung <[email protected]>

src/tools/miri/src/borrow_tracker/tree_borrows/foreign_access_skipping.rs

@@ -0,0 +1,108 @@
+use super::AccessKind;
+use super::tree::AccessRelatedness;
+
+/// To speed up tree traversals, we want to skip traversing subtrees when we know the traversal will have no effect.
+/// This is often the case for foreign accesses, since usually foreign accesses happen several times in a row, but also
+/// foreign accesses are idempotent. In particular, see tests `foreign_read_is_noop_after_foreign_write` and `all_transitions_idempotent`.
+/// Thus, for each node we keep track of the "strongest idempotent foreign access" (SIFA), i.e. which foreign access can be skipped.
+/// Note that for correctness, it is not required that this is the strongest access, just any access it is idempotent under. In particular, setting
+/// it to `None` is always correct, but the point of this optimization is to have it be as strong as possible so that more accesses can be skipped.
+/// This enum represents the kinds of values we store:
+/// - `None` means that the node (and its subtrees) are not (guaranteed to be) idempotent under any foreign access.
+/// - `Read` means that the node (and its subtrees) are idempotent under foreign reads, but not (yet / necessarily) under foreign writes.
+/// - `Write` means that the node (and its subtrees) are idempotent under foreign writes. This also implies that it is idempotent under foreign
+///   reads, since reads are stronger than writes (see test `foreign_read_is_noop_after_foreign_write`). In other words, this node can be skipped
+///   for all foreign accesses.
+///
+/// Since a traversal does not just visit a node, but instead the entire subtree, the SIFA field for a given node indicates that the access to
+/// *the entire subtree* rooted at that node can be skipped. In order for this to work, we maintain the global invariant that at
+/// each location, the SIFA at each child must be stronger than that at the parent. For normal reads and writes, this is easily accomplished by
+/// tracking each foreign access as it occurs, so that then the next access can be skipped. This also obviously maintains the invariant, because
+/// if a node undergoes a foreign access, then all its children also see this as a foreign access. However, the invariant is broken during retags,
+/// because retags act across the entire allocation, but only emit a read event across a specific range. This means that for all nodes outside that
+/// range, the invariant is potentially broken, since a new child with a weaker SIFA is inserted. Thus, during retags, special care is taken to
+/// "manually" reset the parent's SIFA to be at least as strong as the new child's. This is accomplished with the `ensure_no_stronger_than` method.
+///
+/// Note that we derive Ord and PartialOrd, so the order in which variants are listed below matters:
+/// None < Read < Write. Do not change that order. See the `test_order` test.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Default)]
+pub enum IdempotentForeignAccess {
+    #[default]
+    None,
+    Read,
+    Write,
+}
+
+impl IdempotentForeignAccess {
+    /// Returns true if a node where the strongest idempotent foreign access is `self`
+    /// can skip the access `happening_next`. Note that if this returns
+    /// `true`, then the entire subtree will be skipped.
+    pub fn can_skip_foreign_access(self, happening_next: IdempotentForeignAccess) -> bool {
+        debug_assert!(happening_next.is_foreign());
+        // This ordering is correct. Intuitively, if the last access here was
+        // a foreign write, everything can be skipped, since after a foreign write,
+        // all further foreign accesses are idempotent
+        happening_next <= self
+    }
+
+    /// Updates `self` to account for a foreign access.
+    pub fn record_new(&mut self, just_happened: IdempotentForeignAccess) {
+        if just_happened.is_local() {
+            // If the access is local, reset it.
+            *self = IdempotentForeignAccess::None;
+        } else {
+            // Otherwise, keep it or stengthen it.
+            *self = just_happened.max(*self);
+        }
+    }
+
+    /// Returns true if this access is local.
+    pub fn is_local(self) -> bool {
+        matches!(self, IdempotentForeignAccess::None)
+    }
+
+    /// Returns true if this access is foreign, i.e. not local.
+    pub fn is_foreign(self) -> bool {
+        !self.is_local()
+    }
+
+    /// Constructs a foreign access from an `AccessKind`
+    pub fn from_foreign(acc: AccessKind) -> IdempotentForeignAccess {
+        match acc {
+            AccessKind::Read => Self::Read,
+            AccessKind::Write => Self::Write,
+        }
+    }
+
+    /// Usually, tree traversals have an `AccessKind` and an `AccessRelatedness`.
+    /// This methods converts these into the corresponding `IdempotentForeignAccess`, to be used
+    /// to e.g. invoke `can_skip_foreign_access`.
+    pub fn from_acc_and_rel(acc: AccessKind, rel: AccessRelatedness) -> IdempotentForeignAccess {
+        if rel.is_foreign() { Self::from_foreign(acc) } else { Self::None }
+    }
+
+    /// During retags, the SIFA needs to be weakened to account for children with weaker SIFAs being inserted.
+    /// Thus, this method is called from the bottom up on each parent, until it returns false, which means the
+    /// "children have stronger SIFAs" invariant is restored.
+    pub fn ensure_no_stronger_than(&mut self, strongest_allowed: IdempotentForeignAccess) -> bool {
+        if *self > strongest_allowed {
+            *self = strongest_allowed;
+            true
+        } else {
+            false
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::IdempotentForeignAccess;
+
+    #[test]
+    fn test_order() {
+        // The internal logic relies on this order.
+        // Do not change.
+        assert!(IdempotentForeignAccess::None < IdempotentForeignAccess::Read);
+        assert!(IdempotentForeignAccess::Read < IdempotentForeignAccess::Write);
+    }
+}

src/tools/miri/src/borrow_tracker/tree_borrows/mod.rs

@@ -9,6 +9,7 @@ use crate::concurrency::data_race::NaReadType;
 use crate::*;
 
 pub mod diagnostics;
+mod foreign_access_skipping;
 mod perms;
 mod tree;
 mod unimap;
@@ -296,7 +297,14 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             this.machine.current_span(),
         )?;
         // Record the parent-child pair in the tree.
-        tree_borrows.new_child(orig_tag, new_tag, new_perm.initial_state, range, span)?;
+        tree_borrows.new_child(
+            orig_tag,
+            new_tag,
+            new_perm.initial_state,
+            range,
+            span,
+            new_perm.protector.is_some(),
+        )?;
         drop(tree_borrows);
 
         // Also inform the data race model (but only if any bytes are actually affected).

src/tools/miri/src/borrow_tracker/tree_borrows/perms.rs

@@ -48,6 +48,7 @@ enum PermissionPriv {
     Disabled,
 }
 use self::PermissionPriv::*;
+use super::foreign_access_skipping::IdempotentForeignAccess;
 
 impl PartialOrd for PermissionPriv {
     /// PermissionPriv is ordered by the reflexive transitive closure of
@@ -87,6 +88,28 @@ impl PermissionPriv {
     fn compatible_with_protector(&self) -> bool {
         !matches!(self, ReservedIM)
     }
+
+    /// See `foreign_access_skipping.rs`. Computes the SIFA of a permission.
+    fn strongest_idempotent_foreign_access(&self, prot: bool) -> IdempotentForeignAccess {
+        match self {
+            // A protected non-conflicted Reserved will become conflicted under a foreign read,
+            // and is hence not idempotent under it.
+            ReservedFrz { conflicted } if prot && !conflicted => IdempotentForeignAccess::None,
+            // Otherwise, foreign reads do not affect Reserved
+            ReservedFrz { .. } => IdempotentForeignAccess::Read,
+            // Famously, ReservedIM survives foreign writes. It is never protected.
+            ReservedIM if prot => unreachable!("Protected ReservedIM should not exist!"),
+            ReservedIM => IdempotentForeignAccess::Write,
+            // Active changes on any foreign access (becomes Frozen/Disabled).
+            Active => IdempotentForeignAccess::None,
+            // Frozen survives foreign reads, but not writes.
+            Frozen => IdempotentForeignAccess::Read,
+            // Disabled survives foreign reads and writes. It survives them
+            // even if protected, because a protected `Disabled` is not initialized
+            // and does therefore not trigger UB.
+            Disabled => IdempotentForeignAccess::Write,
+        }
+    }
 }
 
 /// This module controls how each permission individually reacts to an access.
@@ -305,6 +328,13 @@ impl Permission {
             (Disabled, _) => false,
         }
     }
+
+    /// Returns the strongest foreign action this node survives (without change),
+    /// where `prot` indicates if it is protected.
+    /// See `foreign_access_skipping`
+    pub fn strongest_idempotent_foreign_access(&self, prot: bool) -> IdempotentForeignAccess {
+        self.inner.strongest_idempotent_foreign_access(prot)
+    }
 }
 
 impl PermTransition {
@@ -575,7 +605,7 @@ mod propagation_optimization_checks {
     impl Exhaustive for AccessRelatedness {
         fn exhaustive() -> Box<dyn Iterator<Item = Self>> {
             use AccessRelatedness::*;
-            Box::new(vec![This, StrictChildAccess, AncestorAccess, DistantAccess].into_iter())
+            Box::new(vec![This, StrictChildAccess, AncestorAccess, CousinAccess].into_iter())
         }
     }
 
@@ -634,6 +664,35 @@ mod propagation_optimization_checks {
         }
     }
 
+    #[test]
+    #[rustfmt::skip]
+    fn permission_sifa_is_correct() {
+        // Tests that `strongest_idempotent_foreign_access` is correct. See `foreign_access_skipping.rs`.
+        for perm in PermissionPriv::exhaustive() {
+            // Assert that adding a protector makes it less idempotent.
+            if perm.compatible_with_protector() {
+                assert!(perm.strongest_idempotent_foreign_access(true) <= perm.strongest_idempotent_foreign_access(false));
+            }
+            for prot in bool::exhaustive() {
+                if prot {
+                    precondition!(perm.compatible_with_protector());
+                }
+                let access = perm.strongest_idempotent_foreign_access(prot);
+                // We now assert it is idempotent, and never causes UB.
+                // First, if the SIFA includes foreign reads, assert it is idempotent under foreign reads.
+                if access >= IdempotentForeignAccess::Read {
+                    // We use `CousinAccess` here. We could also use `AncestorAccess`, since `transition::perform_access` treats these the same.
+                    // The only place they are treated differently is in protector end accesses, but these are not handled here.
+                    assert_eq!(perm, transition::perform_access(AccessKind::Read, AccessRelatedness::CousinAccess, perm, prot).unwrap());
+                }
+                // Then, if the SIFA includes foreign writes, assert it is idempotent under foreign writes.
+                if access >= IdempotentForeignAccess::Write {
+                    assert_eq!(perm, transition::perform_access(AccessKind::Write, AccessRelatedness::CousinAccess, perm, prot).unwrap());
+                }
+            }
+        }
+    }
+
     #[test]
     // Check that all transitions are consistent with the order on PermissionPriv,
     // i.e. Reserved -> Active -> Frozen -> Disabled