Make Tree Borrows Provenance GC compact the tree

bench-cargo-miri/slice-chunked/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "slice-chunked"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]

bench-cargo-miri/slice-chunked/src/main.rs

@@ -0,0 +1,26 @@
+//! This is a small example using slice::chunks, which creates a very large Tree Borrows tree.
+//! Thanks to ##3837, the GC now compacts the tree, so this test can be run in a reasonable time again.
+//! The actual code is adapted from tiny_skia, see https://github.com/RazrFalcon/tiny-skia/blob/master/src/pixmap.rs#L121
+//! To make this benchmark demonstrate the effectiveness, run with MIRIFLAGS="-Zmiri-tree-borrows -Zmiri-provenance-gc=100"
+
+const N: usize = 1000;
+
+fn input_vec() -> Vec<u8> {
+    vec![0; N]
+}
+
+fn main() {
+    let data_len = 2 * N;
+    let mut rgba_data = Vec::with_capacity(data_len);
+    let img_data = input_vec();
+    for slice in img_data.chunks(2) {
+        let gray = slice[0];
+        let alpha = slice[1];
+        rgba_data.push(gray);
+        rgba_data.push(gray);
+        rgba_data.push(gray);
+        rgba_data.push(alpha);
+    }
+
+    assert_eq!(rgba_data.len(), data_len);
+}

src/borrow_tracker/tree_borrows/perms.rs

@@ -130,7 +130,7 @@ mod transition {
             Active =>
                 if protected {
                     // We wrote, someone else reads -- that's bad.
-                    // (If this is initialized, this move-to-protected will mean insta-UB.)
+                    // (Since Active is always initialized, this move-to-protected will mean insta-UB.)
                     Disabled
                 } else {
                     // We don't want to disable here to allow read-read reordering: it is crucial
@@ -267,6 +267,43 @@ impl Permission {
         transition::perform_access(kind, rel_pos, old_state, protected)
             .map(|new_state| PermTransition { from: old_state, to: new_state })
     }
+
+    /// During a provenance GC, we want to compact the tree.
+    /// For this, we want to merge nodes upwards if they have a singleton parent.
+    /// But we need to be careful: If the parent is Frozen, and the child is Reserved,
+    /// we can not do such a merge. In general, such a merge is possible if the parent
+    /// allows similar accesses, and in particular if the parent never causes UB on its
+    /// own. This is enforced by a test, namely `tree_compacting_is_sound`. See that
+    /// test for more information.
+    /// This method is only sound if the parent is not protected. We never attempt to
+    /// remove protected parents.
+    pub fn can_be_replaced_by_child(self, child: Self) -> bool {
+        match (self.inner, child.inner) {
+            // ReservedIM can be replaced by anything, as it allows all
+            // transitions.
+            (ReservedIM, _) => true,
+            // Reserved (as parent, where conflictedness does not matter)
+            // can be replaced by all but ReservedIM,
+            // since ReservedIM alone would survive foreign writes
+            (ReservedFrz { .. }, ReservedIM) => false,
+            (ReservedFrz { .. }, _) => true,
+            // Active can not be replaced by something surviving
+            // foreign reads and then remaining writable
+            (Active, ReservedIM) => false,
+            (Active, ReservedFrz { .. }) => false,
+            (Active, Active) => true,
+            // Active can be replaced by Frozen, since it is not protected
+            (Active, Frozen) => true,
+            (Active, Disabled) => true,
+            // Frozen can only be replaced by Disabled
+            (Frozen, Frozen) => true,
+            (Frozen, Disabled) => true,
+            (Frozen, _) => false,
+            // Disabled can not be replaced by anything else
+            (Disabled, Disabled) => true,
+            (Disabled, _) => false,
+        }
+    }
 }
 
 impl PermTransition {

src/borrow_tracker/tree_borrows/tree.rs

@@ -128,6 +128,22 @@ impl LocationState {
         Ok(transition)
     }
 
+    /// Like `perform_access`, but ignores the diagnostics, and also is pure.
+    /// As such, it returns `Some(x)` if the transition succeeded, or `None`
+    /// if there was an error.
+    #[allow(unused)]
+    fn perform_access_no_fluff(
+        mut self,
+        access_kind: AccessKind,
+        rel_pos: AccessRelatedness,
+        protected: bool,
+    ) -> Option<Self> {
+        match self.perform_access(access_kind, rel_pos, protected) {
+            Ok(_) => Some(self),
+            Err(_) => None,
+        }
+    }
+
     // Helper to optimize the tree traversal.
     // The optimization here consists of observing thanks to the tests
     // `foreign_read_is_noop_after_foreign_write` and `all_transitions_idempotent`,
@@ -840,6 +856,57 @@ impl Tree {
         node.children.is_empty() && !live.contains(&node.tag)
     }
 
+    /// Checks whether a node can be replaced by its only child.
+    /// If so, returns the index of said only child.
+    /// If not, returns none.
+    fn can_be_replaced_by_single_child(
+        &self,
+        idx: UniIndex,
+        live: &FxHashSet<BorTag>,
+    ) -> Option<UniIndex> {
+        let node = self.nodes.get(idx).unwrap();
+        if node.children.len() != 1 || live.contains(&node.tag) || node.parent.is_none() {
+            return None;
+        }
+        // Since protected nodes are never GC'd (see `borrow_tracker::GlobalStateInner::visit_provenance`),
+        // we know that `node` is not protected because otherwise `live` would
+        // have contained `node.tag`.
+        let child_idx = node.children[0];
+        let child = self.nodes.get(child_idx).unwrap();
+        for (_, data) in self.rperms.iter_all() {
+            let parent_perm =
+                data.get(idx).map(|x| x.permission).unwrap_or_else(|| node.default_initial_perm);
+            let child_perm = data
+                .get(child_idx)
+                .map(|x| x.permission)
+                .unwrap_or_else(|| child.default_initial_perm);
+            if !parent_perm.can_be_replaced_by_child(child_perm) {
+                return None;
+            }
+        }
+
+        Some(child_idx)
+    }
+
+    /// Properly removes a node.
+    /// The node to be removed should not otherwise be usable. It also
+    /// should have no children, but this is not checked, so that nodes
+    /// whose children were rotated somewhere else can be deleted without
+    /// having to first modify them to clear that array.
+    /// otherwise (i.e. the GC should have marked it as removable).
+    fn remove_useless_node(&mut self, this: UniIndex) {
+        // Due to the API of UniMap we must make sure to call
+        // `UniValMap::remove` for the key of this node on *all* maps that used it
+        // (which are `self.nodes` and every range of `self.rperms`)
+        // before we can safely apply `UniKeyMap::remove` to truly remove
+        // this tag from the `tag_mapping`.
+        let node = self.nodes.remove(this).unwrap();
+        for (_perms_range, perms) in self.rperms.iter_mut_all() {
+            perms.remove(this);
+        }
+        self.tag_mapping.remove(&node.tag);
+    }
+
     /// Traverses the entire tree looking for useless tags.
     /// Removes from the tree all useless child nodes of root.
     /// It will not delete the root itself.
@@ -883,23 +950,20 @@ impl Tree {
                 // Remove all useless children.
                 children_of_node.retain_mut(|idx| {
                     if self.is_useless(*idx, live) {
-                        // Note: In the rest of this comment, "this node" refers to `idx`.
-                        // This node has no more children (if there were any, they have already been removed).
-                        // It is also unreachable as determined by the GC, so we can remove it everywhere.
-                        // Due to the API of UniMap we must make sure to call
-                        // `UniValMap::remove` for the key of this node on *all* maps that used it
-                        // (which are `self.nodes` and every range of `self.rperms`)
-                        // before we can safely apply `UniKeyMap::remove` to truly remove
-                        // this tag from the `tag_mapping`.
-                        let node = self.nodes.remove(*idx).unwrap();
-                        for (_perms_range, perms) in self.rperms.iter_mut_all() {
-                            perms.remove(*idx);
-                        }
-                        self.tag_mapping.remove(&node.tag);
-                        // now delete it
+                        // delete it everywhere else
+                        self.remove_useless_node(*idx);
+                        // and delete it from children_of_node
                         false
                     } else {
-                        // do nothing, but retain
+                        if let Some(nextchild) = self.can_be_replaced_by_single_child(*idx, live) {
+                            // delete the in-between child
+                            self.remove_useless_node(*idx);
+                            // set the new child's parent
+                            self.nodes.get_mut(nextchild).unwrap().parent = Some(*tag);
+                            // save the new child in children_of_node
+                            *idx = nextchild;
+                        }
+                        // retain it
                         true
                     }
                 });

src/borrow_tracker/tree_borrows/tree/tests.rs

@@ -64,6 +64,71 @@ fn all_read_accesses_commute() {
     }
 }
 
+fn as_foreign_or_child(related: AccessRelatedness) -> &'static str {
+    if related.is_foreign() { "foreign" } else { "child" }
+}
+
+fn as_protected(b: bool) -> &'static str {
+    if b { " (protected)" } else { "" }
+}
+
+fn as_lazy_or_init(b: bool) -> &'static str {
+    if b { "initialized" } else { "lazy" }
+}
+
+/// Test that tree compacting (as performed by the GC) is sound.
+/// Specifically, the GC will replace a parent by a child if the parent is not
+/// protected, and if `can_be_replaced_by_child(parent, child)` is true.
+/// To check that this is sound, the function must be a simulation, i.e.
+/// if both are accessed, the results must still be in simulation, and also
+/// if an access is UB, it must also be UB if done only at the child.
+#[test]
+fn tree_compacting_is_sound() {
+    // The parent is unprotected
+    let parent_protected = false;
+    for ([parent, child], child_protected) in <([LocationState; 2], bool)>::exhaustive() {
+        if child_protected {
+            precondition!(child.compatible_with_protector())
+        }
+        precondition!(parent.permission().can_be_replaced_by_child(child.permission()));
+        for (kind, rel) in <(AccessKind, AccessRelatedness)>::exhaustive() {
+            let new_parent = parent.perform_access_no_fluff(kind, rel, parent_protected);
+            let new_child = child.perform_access_no_fluff(kind, rel, child_protected);
+            match (new_parent, new_child) {
+                (Some(np), Some(nc)) => {
+                    assert!(
+                        np.permission().can_be_replaced_by_child(nc.permission()),
+                        "`can_be_replaced_by_child` is not a simulation: on a {} {} to a {} parent and a {} {}{} child, the parent becomes {}, the child becomes {}, and these are not in simulation!",
+                        as_foreign_or_child(rel),
+                        kind,
+                        parent.permission(),
+                        as_lazy_or_init(child.is_initialized()),
+                        child.permission(),
+                        as_protected(child_protected),
+                        np.permission(),
+                        nc.permission()
+                    )
+                }
+                (_, None) => {
+                    // the child produced UB, this is fine no matter what the parent does
+                }
+                (None, Some(nc)) => {
+                    panic!(
+                        "`can_be_replaced_by_child` does not have the UB property: on a {} {} to a(n) {} parent and a(n) {} {}{} child, only the parent causes UB, while the child becomes {}, and it is not allowed for only the parent to cause UB!",
+                        as_foreign_or_child(rel),
+                        kind,
+                        parent.permission(),
+                        as_lazy_or_init(child.is_initialized()),
+                        child.permission(),
+                        as_protected(child_protected),
+                        nc.permission()
+                    )
+                }
+            }
+        }
+    }
+}
+
 #[test]
 #[rustfmt::skip]
 // Ensure that of 2 accesses happen, one foreign and one a child, and we are protected, that we

tests/pass-dep/concurrency/linux-futex.rs

@@ -158,7 +158,9 @@ fn wait_wake() {
         );
     }
 
-    assert!((200..1000).contains(&start.elapsed().as_millis()));
+    // When running this in stress-gc mode, things can take quite long.
+    // So the timeout is 3000 ms.
+    assert!((200..3000).contains(&start.elapsed().as_millis()));
     t.join().unwrap();
 }