Tree Borrows: improve protector end access child skipping

src/borrow_tracker/tree_borrows/tree.rs

@@ -661,6 +661,7 @@ impl<'tcx> Tree {
                 global,
                 ChildrenVisitMode::VisitChildrenOfAccessed,
                 &diagnostics,
+                None, // min_exposed_child only matters for protector end access,
             )?;
         }
         interp_ok(())
@@ -686,6 +687,12 @@ impl<'tcx> Tree {
 
         let source_idx = self.tag_mapping.get(&tag).unwrap();
 
+        let min_exposed_child = if self.roots.len() > 1 {
+            LocationTree::get_min_exposed_child(source_idx, &self.nodes)
+        } else {
+            None
+        };
+
         // This is a special access through the entire allocation.
         // It actually only affects `accessed` locations, so we need
         // to filter on those before initiating the traversal.
@@ -716,6 +723,7 @@ impl<'tcx> Tree {
                     global,
                     ChildrenVisitMode::SkipChildrenOfAccessed,
                     &diagnostics,
+                    min_exposed_child,
                 )?;
             }
         }
@@ -876,9 +884,35 @@ impl Tree {
 }
 
 impl<'tcx> LocationTree {
+    /// Returns the smallest exposed tag, if any, that is a transitive child of `root`.
+    fn get_min_exposed_child(root: UniIndex, nodes: &UniValMap<Node>) -> Option<BorTag> {
+        // We cannot use the wildcard datastructure to improve this lookup. This is because
+        // the datastructure only tracks enabled nodes and we need to also consider disabled ones.
+        let mut stack = vec![root];
+        let mut min_tag = None;
+        while let Some(idx) = stack.pop() {
+            let node = nodes.get(idx).unwrap();
+            if min_tag.is_some_and(|min| min < node.tag) {
+                continue;
+            }
+            stack.extend_from_slice(node.children.as_slice());
+            if node.is_exposed {
+                min_tag = match min_tag {
+                    Some(prev) if prev < node.tag => Some(prev),
+                    _ => Some(node.tag),
+                };
+            }
+        }
+        min_tag
+    }
+
     /// Performs an access on this location.
     /// * `access_source`: The index, if any, where the access came from.
     /// * `visit_children`: Whether to skip updating the children of `access_source`.
+    /// * `min_exposed_child`: any subtree that could be a child of `access_source` must
+    ///   have a root with a larger tag than this. If None then no other subtree can be
+    ///   a child of `access_source`.
+    ///   This only gets read if `visit_children==SkipChildrenOfAccessed`.
     fn perform_access(
         &mut self,
         roots: impl Iterator<Item = UniIndex>,
@@ -888,6 +922,7 @@ impl<'tcx> LocationTree {
         global: &GlobalState,
         visit_children: ChildrenVisitMode,
         diagnostics: &DiagnosticInfo,
+        min_exposed_child: Option<BorTag>,
     ) -> InterpResult<'tcx> {
         let accessed_root = if let Some(idx) = access_source {
             Some(self.perform_normal_access(
@@ -906,11 +941,21 @@ impl<'tcx> LocationTree {
         };
 
         let accessed_root_tag = accessed_root.map(|idx| nodes.get(idx).unwrap().tag);
-        if matches!(visit_children, ChildrenVisitMode::SkipChildrenOfAccessed) {
-            // FIXME: approximate which roots could be children of the accessed node and only skip them instead of all other trees.
-            return interp_ok(());
-        }
         for root in roots {
+            let tag = nodes.get(root).unwrap().tag;
+            // On a protector release access we skip the children of the accessed tag so that
+            // we can correctly return references from functions. However, if the tag has
+            // exposed children then some of the wildcard subtrees could also be children of
+            // the accessed node and would also need to be skipped. We can narrow down which
+            // child trees are children by comparing their root tag to the minimum exposed
+            // child of the accessed node. As the parent tag is always smaller than the child
+            // tag this means we only need to skip subtrees with a root tag larger than
+            // `min_exposed_child`.
+            if let Some(min_exposed_child) = min_exposed_child
+                && tag > min_exposed_child
+            {
+                break;
+            }
             // We don't perform a wildcard access on the tree we already performed a
             // normal access on.
             if Some(root) == accessed_root {

tests/fail/tree_borrows/wildcard/protector_release.rs

@@ -0,0 +1,77 @@
+//@compile-flags: -Zmiri-tree-borrows -Zmiri-permissive-provenance
+use std::cell::UnsafeCell;
+
+/// This is a variant of the test in `../protector-write-lazy.rs`, but with
+/// wildcard references.
+/// Checks that a protector release transitions a foreign reference on a reference,
+/// if we can determine that it's not a child of the released tag.
+/// For this version we know the tag is not a child, because its wildcard root has
+/// a smaller tag then the released reference.
+pub fn main() {
+    // We need two locations so that we can create a new reference
+    // that is foreign to an already active tag.
+    let mut x: UnsafeCell<[u32; 2]> = UnsafeCell::new([32, 33]);
+    let ref1 = &mut x;
+    let cell_ptr = ref1.get() as *mut u32;
+
+    let int = ref1 as *mut UnsafeCell<[u32; 2]> as usize;
+    let wild = int as *mut UnsafeCell<u32>;
+
+    let ref2 = unsafe { &mut *cell_ptr };
+
+    // `ref3` gets created before the protected ref `arg4`.
+    let ref3 = unsafe { &mut *wild.wrapping_add(1) };
+
+    let protect = |arg4: &mut u32| {
+        // Activates arg4. This would disable ref3  at [0] if it wasn't a cell.
+        *arg4 = 41;
+
+        // Creates an exposed child of arg4.
+        let ref5 = &mut *arg4;
+        let _int = ref5 as *mut u32 as usize;
+
+        // This creates ref6 from ref3 at [1], so that it doesn't disable arg4 at [0].
+        let ref6 = unsafe { &mut *ref3.get() };
+
+        //    ┌───────────┐
+        //    │    ref1*  │
+        //    │ Cel │ Cel │           *
+        //    └─────┬─────┘           │
+        //          │                 │
+        //          │                 │
+        //          ▼                 ▼
+        //    ┌───────────┐     ┌───────────┐
+        //    │ ref2      │     │    ref3   │
+        //    │ Act │ Res │     │ Cel │ Cel │
+        //    └─────┬─────┘     └─────┬─────┘
+        //          │                 │
+        //          │                 │
+        //          ▼                 ▼
+        //    ┌───────────┐     ┌───────────┐
+        //    │ arg4      │     │       ref6│
+        //    │ Act │ Res │     │ Res │ Res │
+        //    └─────┬─────┘     └───────────┘
+        //          │
+        //          │
+        //          ▼
+        //    ┌───────────┐
+        //    │ ref5*     │
+        //    │ Res │ Res │
+        //    └───────────┘
+
+        // Creates a pointer to [0] with the provenance of ref6.
+        return (ref6 as *mut u32).wrapping_sub(1);
+
+        // Protector release on arg4 happens here.
+        // This should cause a foreign write on all foreign nodes,
+        // unless they could be a child of arg4.
+        // arg4 has an exposed child, so some tags with a wildcard
+        // ancestor could be its children.
+        // However, the root of ref6 was created before arg4, so it
+        // cannot be a child of it. So it should get disabled
+        // (at location [0]).
+    };
+
+    let ptr = protect(ref2);
+    let _fail = unsafe { *ptr }; //~ ERROR: /read access through .* is forbidden/
+}

tests/fail/tree_borrows/wildcard/protector_release.stderr

@@ -0,0 +1,25 @@
+error: Undefined Behavior: read access through <TAG> at ALLOC[0x0] is forbidden
+  --> tests/fail/tree_borrows/wildcard/protector_release.rs:LL:CC
+   |
+LL |     let _fail = unsafe { *ptr };
+   |                          ^^^^ 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 Disabled which forbids this child read access
+help: the accessed tag <TAG> was created here, in the initial state Reserved
+  --> tests/fail/tree_borrows/wildcard/protector_release.rs:LL:CC
+   |
+LL |         let ref6 = unsafe { &mut *ref3.get() };
+   |                             ^^^^^^^^^^^^^^^^
+help: the accessed tag <TAG> later transitioned to Disabled due to a protector release (acting as a foreign write access) on every location previously accessed by this tag
+  --> tests/fail/tree_borrows/wildcard/protector_release.rs:LL:CC
+   |
+LL |     };
+   |      ^
+   = help: this transition corresponds to a loss of read and write permissions
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+

tests/fail/tree_borrows/wildcard/protector_release2.rs

@@ -0,0 +1,77 @@
+//@compile-flags: -Zmiri-tree-borrows -Zmiri-permissive-provenance
+use std::cell::UnsafeCell;
+
+/// This is a variant of the test in `../protector-write-lazy.rs`, but with
+/// wildcard references.
+/// Checks that a protector release transitions a foreign reference on a reference,
+/// if we can determine that it's not a child of the released tag.
+/// For this version we know the tag is not a child, because its wildcard root has
+/// a smaller tag then the exposed child of the protected tag.
+pub fn main() {
+    // We need two locations so that we can create a new reference
+    // that is foreign to an already active tag.
+    let mut x: UnsafeCell<[u32; 2]> = UnsafeCell::new([32, 33]);
+    let ref1 = &mut x;
+    let cell_ptr = ref1.get() as *mut u32;
+
+    let int = ref1 as *mut UnsafeCell<[u32; 2]> as usize;
+    let wild = int as *mut UnsafeCell<u32>;
+
+    let ref2 = unsafe { &mut *cell_ptr };
+
+    let protect = |arg3: &mut u32| {
+        // `ref4` gets created after the protected ref `arg3` but before the exposed `ref5`.
+        let ref4 = unsafe { &mut *wild.wrapping_add(1) };
+
+        // Activates arg4. This would disable ref3  at [0] if it wasn't a cell.
+        *arg3 = 41;
+
+        // Creates an exposed child of arg3.
+        let ref5 = &mut *arg3;
+        let _int = ref5 as *mut u32 as usize;
+
+        // This creates ref6 from ref4 at [1], so that it doesn't disable arg3 at [0].
+        let ref6 = unsafe { &mut *ref4.get() };
+
+        //    ┌───────────┐
+        //    │    ref1*  │
+        //    │ Cel │ Cel │           *
+        //    └─────┬─────┘           │
+        //          │                 │
+        //          │                 │
+        //          ▼                 ▼
+        //    ┌───────────┐     ┌───────────┐
+        //    │ ref2      │     │       ref4│
+        //    │ Act │ Res │     │ Cel │ Cel │
+        //    └─────┬─────┘     └─────┬─────┘
+        //          │                 │
+        //          │                 │
+        //          ▼                 ▼
+        //    ┌───────────┐     ┌───────────┐
+        //    │ arg3      │     │       ref6│
+        //    │ Act │ Res │     │ Res │ Res │
+        //    └─────┬─────┘     └───────────┘
+        //          │
+        //          │
+        //          ▼
+        //    ┌───────────┐
+        //    │ ref5*     │
+        //    │ Res │ Res │
+        //    └───────────┘
+
+        // Creates a pointer to [0] with the provenance of ref6.
+        return (ref6 as *mut u32).wrapping_sub(1);
+
+        // Protector release on arg3 happens here.
+        // This should cause a foreign write on all foreign nodes,
+        // unless they could be a child of arg3.
+        // arg3 has an exposed child, so some tags with a wildcard
+        // ancestor could be its children.
+        // However, the root of ref6 was created before the exposed
+        // child ref5, so it cannot be a child of it. So it should
+        // get disabled (at location [0]).
+    };
+
+    let ptr = protect(ref2);
+    let _fail = unsafe { *ptr }; //~ ERROR: /read access through .* is forbidden/
+}

tests/fail/tree_borrows/wildcard/protector_release2.stderr

@@ -0,0 +1,25 @@
+error: Undefined Behavior: read access through <TAG> at ALLOC[0x0] is forbidden
+  --> tests/fail/tree_borrows/wildcard/protector_release2.rs:LL:CC
+   |
+LL |     let _fail = unsafe { *ptr };
+   |                          ^^^^ 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 Disabled which forbids this child read access
+help: the accessed tag <TAG> was created here, in the initial state Reserved
+  --> tests/fail/tree_borrows/wildcard/protector_release2.rs:LL:CC
+   |
+LL |         let ref6 = unsafe { &mut *ref4.get() };
+   |                             ^^^^^^^^^^^^^^^^
+help: the accessed tag <TAG> later transitioned to Disabled due to a protector release (acting as a foreign write access) on every location previously accessed by this tag
+  --> tests/fail/tree_borrows/wildcard/protector_release2.rs:LL:CC
+   |
+LL |     };
+   |      ^
+   = help: this transition corresponds to a loss of read and write permissions
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+