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`. 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, whereas it finishes within seconds in Stacked
Borrows.

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.

Author: JoJoDeveloping

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