Auto merge of rust-lang#3837 - JoJoDeveloping:tb-compacting-provenanc…

…e-gc, r=RalfJung

Make Tree Borrows Provenance GC compact the tree

Follow-up on rust-lang#3833 and rust-lang#3835. In these PRs, the TB GC was fixed to no longer cause a stack overflow. One test that motivated it was the test `fill::horizontal_line` in [`tiny-skia`](https://github.com/RazrFalcon/tiny-skia). But not causing stack overflows was not a large improvents, since it did not fix the fundamental issue: The tree was too large. The test now ran, but it required gigabytes of memory and hours of time (only for it to be OOM-killed 🤬), whereas it finishes within 24 seconds in Stacked Borrows. With this merged, it finishes in about 40 seconds under TB.

The problem in that test was that it used [`slice::chunked`](https://doc.rust-lang.org/std/primitive.slice.html#method.chunks) to iterate a slice in chunks. That iterator is written to reborrow at each call to `next`, which creates a linear tree with a bunch of intermediary nodes, which also fragments the `RangeMap` for that allocation.

The solution is to now compact the tree, so that these interior nodes are removed. Care is taken to not remove nodes that are protected, or that otherwise restrict their children.

I am currently only 99% sure that this is sound, and I do also think that this could compact even more. So `@Vanille-N` please also have a look at whether I got the compacting logic right.

For a more visual comparison, [here is a gist](https://gist.github.com/JoJoDeveloping/ae4a7f7c29335a4c233ef42d2f267b01) of what the tree looks like at one point during that test, with and without compacting.

This new GC requires a different iteration order during accesses (since the current one can make the error messages non-deterministic), so it is rebased on top of rust-lang#3843 and requires that PR to be merged first.

src/tools/miri/bench-cargo-miri/slice-chunked/Cargo.lock

@@ -0,0 +1,7 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "slice-chunked"
+version = "0.1.0"

src/tools/miri/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]

src/tools/miri/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/tools/miri/src/borrow_tracker/mod.rs

@@ -71,6 +71,12 @@ pub struct FrameState {
 
 impl VisitProvenance for FrameState {
     fn visit_provenance(&self, visit: &mut VisitWith<'_>) {
+        // Visit all protected tags. At least in Tree Borrows,
+        // protected tags can not be GC'd because they still have
+        // an access coming when the protector ends. Additionally,
+        // the tree compacting mechanism of TB's GC relies on the fact
+        // that all protected tags are marked as live for correctness,
+        // so we _have_ to visit them here.
         for (id, tag) in &self.protected_tags {
             visit(Some(*id), Some(*tag));
         }

src/tools/miri/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,44 @@ 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,
+            // Replacing a state by itself is always okay, even if the child state is protected.
+            (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 (and itself).
+            (Frozen, Frozen) => true,
+            (Frozen, Disabled) => true,
+            (Frozen, _) => false,
+            // Disabled can not be replaced by anything else.
+            (Disabled, Disabled) => true,
+            (Disabled, _) => false,
+        }
+    }
 }
 
 impl PermTransition {

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

@@ -128,6 +128,22 @@ impl LocationState {
         Ok(transition)
     }
 
+    /// Like `perform_access`, but ignores the concrete error cause and also uses state-passing
+    /// rather than a mutable reference. As such, it returns `Some(x)` if the transition succeeded,
+    /// or `None` if there was an error.
+    #[cfg(test)]
+    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,60 @@ 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();
+
+        // We never want to replace the root node, as it is also kept in `root_ptr_tags`.
+        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::FrameExtra::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();
+        // Check that for that one child, `can_be_replaced_by_child` holds for the permission
+        // on all locations.
+        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.
+    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 +953,21 @@ 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 `idx` node 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) {
+                            // `nextchild` is our grandchild, and will become our direct child.
+                            // Delete the in-between node, `idx`.
+                            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/tools/miri/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

src/tools/miri/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();
 }