Make Tree Borrows Provenance GC no longer produce stack overflows

Author: JoJoDeveloping

Diff for: 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"

Diff for: 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
         }
     }
 }

Diff for: 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
     }
 }