Auto merge of rust-lang#3833 - JoJoDeveloping:tb-fix-stack-overflow, …

…r=RalfJung

Make Tree Borrows Provenance GC no longer produce stack overflows

Most functions operating on Tree Borrows' trees are carefully written to not cause stack overflows due to too much recursion. The one exception is [`Tree::keep_only_needed`](https://github.com/rust-lang/miri/blob/94f5588fafcc7d59fce60ca8f7af0208e6f618d4/src/borrow_tracker/tree_borrows/tree.rs#L724), which just uses regular recursion.
This function is part of the provenance GC, so it is called regularly for every allocation in the program.

Tests show that this is a problem in practice. For example, the test `fill::horizontal_line` in crate `tiny-skia` (version 0.11.4) is such a test.

This PR changes this, this test no now longer crashes. Instead, it succeeds (after a _long_ time).

src/tools/miri/Cargo.toml

@@ -20,7 +20,7 @@ doctest = false # and no doc tests
 [dependencies]
 getrandom = { version = "0.2", features = ["std"] }
 rand = "0.8"
-smallvec = "1.7"
+smallvec = { version = "1.7", features = ["drain_filter"] }
 aes = { version = "0.8.3", features = ["hazmat"] }
 measureme = "11"
 ctrlc = "3.2.5"

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

@@ -10,7 +10,7 @@
 //!   and the relative position of the access;
 //! - idempotency properties asserted in `perms.rs` (for optimizations)
 
-use std::fmt;
+use std::{fmt, mem};
 
 use smallvec::SmallVec;
 
@@ -699,18 +699,24 @@ impl<'tcx> Tree {
 /// Integration with the BorTag garbage collector
 impl Tree {
     pub fn remove_unreachable_tags(&mut self, live_tags: &FxHashSet<BorTag>) {
-        let root_is_needed = self.keep_only_needed(self.root, live_tags); // root can't be removed
-        assert!(root_is_needed);
+        self.remove_useless_children(self.root, live_tags);
         // Right after the GC runs is a good moment to check if we can
         // merge some adjacent ranges that were made equal by the removal of some
         // tags (this does not necessarily mean that they have identical internal representations,
         // see the `PartialEq` impl for `UniValMap`)
         self.rperms.merge_adjacent_thorough();
     }
 
+    /// Checks if a node is useless and should be GC'ed.
+    /// A node is useless if it has no children and also the tag is no longer live.
+    fn is_useless(&self, idx: UniIndex, live: &FxHashSet<BorTag>) -> bool {
+        let node = self.nodes.get(idx).unwrap();
+        node.children.is_empty() && !live.contains(&node.tag)
+    }
+
     /// Traverses the entire tree looking for useless tags.
-    /// Returns true iff the tag it was called on is still live or has live children,
-    /// and removes from the tree all tags that have no live children.
+    /// Removes from the tree all useless child nodes of root.
+    /// It will not delete the root itself.
     ///
     /// NOTE: This leaves in the middle of the tree tags that are unreachable but have
     /// reachable children. There is a potential for compacting the tree by reassigning
@@ -721,42 +727,60 @@ impl Tree {
     /// `child: Reserved`. This tree can exist. If we blindly delete `parent` and reassign
     /// `child` to be a direct child of `root` then Writes to `child` are now permitted
     /// whereas they were not when `parent` was still there.
-    fn keep_only_needed(&mut self, idx: UniIndex, live: &FxHashSet<BorTag>) -> bool {
-        let node = self.nodes.get(idx).unwrap();
-        // FIXME: this function does a lot of cloning, a 2-pass approach is possibly
-        // more efficient. It could consist of
-        // 1. traverse the Tree, collect all useless tags in a Vec
-        // 2. traverse the Vec, remove all tags previously selected
-        // Bench it.
-        let children: SmallVec<_> = node
-            .children
-            .clone()
-            .into_iter()
-            .filter(|child| self.keep_only_needed(*child, live))
-            .collect();
-        let no_children = children.is_empty();
-        let node = self.nodes.get_mut(idx).unwrap();
-        node.children = children;
-        if !live.contains(&node.tag) && no_children {
-            // All of the children and this node are unreachable, delete this tag
-            // from the tree (the children have already been deleted by recursive
-            // calls).
-            // Due to the API of UniMap we must absolutely call
-            // `UniValMap::remove` for the key of this tag on *all* maps that used it
-            // (which are `self.nodes` and every range of `self.rperms`)
-            // before we can safely apply `UniValMap::forget` to truly remove
-            // the tag from the mapping.
-            let tag = node.tag;
-            self.nodes.remove(idx);
-            for (_perms_range, perms) in self.rperms.iter_mut_all() {
-                perms.remove(idx);
+    fn remove_useless_children(&mut self, root: UniIndex, live: &FxHashSet<BorTag>) {
+        // To avoid stack overflows, we roll our own stack.
+        // Each element in the stack consists of the current tag, and the number of the
+        // next child to be processed.
+
+        // The other functions are written using the `TreeVisitorStack`, but that does not work here
+        // since we need to 1) do a post-traversal and 2) remove nodes from the tree.
+        // Since we do a post-traversal (by deleting nodes only after handling all children),
+        // we also need to be a bit smarter than "pop node, push all children."
+        let mut stack = vec![(root, 0)];
+        while let Some((tag, nth_child)) = stack.last_mut() {
+            let node = self.nodes.get(*tag).unwrap();
+            if *nth_child < node.children.len() {
+                // Visit the child by pushing it to the stack.
+                // Also increase `nth_child` so that when we come back to the `tag` node, we
+                // look at the next child.
+                let next_child = node.children[*nth_child];
+                *nth_child += 1;
+                stack.push((next_child, 0));
+                continue;
+            } else {
+                // We have processed all children of `node`, so now it is time to process `node` itself.
+                // First, get the current children of `node`. To appease the borrow checker,
+                // we have to temporarily move the list out of the node, and then put the
+                // list of remaining children back in.
+                let mut children_of_node =
+                    mem::take(&mut self.nodes.get_mut(*tag).unwrap().children);
+                // Remove all useless children, and save them for later.
+                // The closure needs `&self` and the loop below needs `&mut self`, so we need to `collect`
+                // in to a temporary list.
+                let to_remove: Vec<_> =
+                    children_of_node.drain_filter(|x| self.is_useless(*x, live)).collect();
+                // Put back the now-filtered vector.
+                self.nodes.get_mut(*tag).unwrap().children = children_of_node;
+                // Now, all that is left is unregistering the children saved in `to_remove`.
+                for idx in to_remove {
+                    // 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);
+                }
+                // We are done, the parent can continue.
+                stack.pop();
+                continue;
             }
-            self.tag_mapping.remove(&tag);
-            // The tag has been deleted, inform the caller
-            false
-        } else {
-            // The tag is still live or has live children, it must be kept
-            true
         }
     }
 }

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

@@ -12,7 +12,7 @@
 
 #![allow(dead_code)]
 
-use std::hash::Hash;
+use std::{hash::Hash, mem};
 
 use rustc_data_structures::fx::FxHashMap;
 
@@ -187,13 +187,16 @@ impl<V> UniValMap<V> {
         self.data.get_mut(idx.idx as usize).and_then(Option::as_mut)
     }
 
-    /// Delete any value associated with this index. Ok even if the index
-    /// has no associated value.
-    pub fn remove(&mut self, idx: UniIndex) {
+    /// Delete any value associated with this index.
+    /// Returns None if the value was not present, otherwise
+    /// returns the previously stored value.
+    pub fn remove(&mut self, idx: UniIndex) -> Option<V> {
         if idx.idx as usize >= self.data.len() {
-            return;
+            return None;
         }
-        self.data[idx.idx as usize] = None;
+        let mut res = None;
+        mem::swap(&mut res, &mut self.data[idx.idx as usize]);
+        res
     }
 }