TB: tree traversal

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

@@ -115,6 +115,151 @@ pub(super) struct Node {
     pub debug_info: NodeDebugInfo,
 }
 
+/// Data given to the transition function
+struct NodeAppArgs<'node> {
+    /// Node on which the transition is currently being applied
+    node: &'node Node,
+    /// Mutable access to its permissions
+    perm: UniEntry<'node, LocationState>,
+    /// Relative position of the access
+    rel_pos: AccessRelatedness,
+}
+/// Data given to the error handler
+struct ErrHandlerArgs<'node, InErr> {
+    /// Kind of error that occurred
+    error_kind: InErr,
+    /// Tag that triggered the error (not the tag that was accessed,
+    /// rather the parent tag that had insufficient permissions or the
+    /// non-parent tag that had a protector).
+    faulty_tag: &'node NodeDebugInfo,
+}
+/// Internal contents of `Tree` with the minimum of mutable access for
+/// the purposes of the tree traversal functions: the permissions (`perms`) can be
+/// updated but not the tree structure (`tag_mapping` and `nodes`)
+struct TreeVisitor<'tree> {
+    tag_mapping: &'tree UniKeyMap<BorTag>,
+    nodes: &'tree UniValMap<Node>,
+    perms: &'tree mut UniValMap<LocationState>,
+}
+
+/// Whether to continue exploring the children recursively or not.
+enum ContinueTraversal {
+    Recurse,
+    SkipChildren,
+}
+
+impl<'tree> TreeVisitor<'tree> {
+    // Applies `f_propagate` to every vertex of the tree top-down in the following order: first
+    // all ancestors of `start`, then `start` itself, then children of `start`, then the rest.
+    // This ensures that errors are triggered in the following order
+    // - first invalid accesses with insufficient permissions, closest to the root first,
+    // - then protector violations, closest to `start` first.
+    //
+    // `f_propagate` should follow the following format: for a given `Node` it updates its
+    // `Permission` depending on the position relative to `start` (given by an
+    // `AccessRelatedness`).
+    // It outputs whether the tree traversal for this subree should continue or not.
+    fn traverse_parents_this_children_others<InnErr, OutErr>(
+        mut self,
+        start: BorTag,
+        f_propagate: impl Fn(NodeAppArgs<'_>) -> Result<ContinueTraversal, InnErr>,
+        err_builder: impl Fn(ErrHandlerArgs<'_, InnErr>) -> OutErr,
+    ) -> Result<(), OutErr>
+where {
+        struct TreeVisitAux<NodeApp, ErrHandler> {
+            f_propagate: NodeApp,
+            err_builder: ErrHandler,
+            stack: Vec<(UniIndex, AccessRelatedness)>,
+        }
+        impl<NodeApp, InnErr, OutErr, ErrHandler> TreeVisitAux<NodeApp, ErrHandler>
+        where
+            NodeApp: Fn(NodeAppArgs<'_>) -> Result<ContinueTraversal, InnErr>,
+            ErrHandler: Fn(ErrHandlerArgs<'_, InnErr>) -> OutErr,
+        {
+            fn pop(&mut self) -> Option<(UniIndex, AccessRelatedness)> {
+                self.stack.pop()
+            }
+
+            /// Apply the function to the current `tag`, and push its children
+            /// to the stack of future tags to visit.
+            fn exec_and_visit(
+                &mut self,
+                this: &mut TreeVisitor<'_>,
+                tag: UniIndex,
+                exclude: Option<UniIndex>,
+                rel_pos: AccessRelatedness,
+            ) -> Result<(), OutErr> {
+                // 1. apply the propagation function
+                let node = this.nodes.get(tag).unwrap();
+                let recurse =
+                    (self.f_propagate)(NodeAppArgs { node, perm: this.perms.entry(tag), rel_pos })
+                        .map_err(|error_kind| {
+                            (self.err_builder)(ErrHandlerArgs {
+                                error_kind,
+                                faulty_tag: &node.debug_info,
+                            })
+                        })?;
+                // 2. add the children to the stack for future traversal
+                if matches!(recurse, ContinueTraversal::Recurse) {
+                    let child_rel = rel_pos.for_child();
+                    for &child in node.children.iter() {
+                        // some child might be excluded from here and handled separately
+                        if Some(child) != exclude {
+                            self.stack.push((child, child_rel));
+                        }
+                    }
+                }
+                Ok(())
+            }
+        }
+
+        let start_idx = self.tag_mapping.get(&start).unwrap();
+        let mut stack = TreeVisitAux { f_propagate, err_builder, stack: Vec::new() };
+        {
+            let mut path_ascend = Vec::new();
+            // First climb to the root while recording the path
+            let mut curr = start_idx;
+            while let Some(ancestor) = self.nodes.get(curr).unwrap().parent {
+                path_ascend.push((ancestor, curr));
+                curr = ancestor;
+            }
+            // Then descend:
+            // - execute f_propagate on each node
+            // - record children in visit
+            while let Some((ancestor, next_in_path)) = path_ascend.pop() {
+                // Explore ancestors in descending order.
+                // `next_in_path` is excluded from the recursion because it
+                // will be the `ancestor` of the next iteration.
+                // It also needs a different `AccessRelatedness` than the other
+                // children of `ancestor`.
+                stack.exec_and_visit(
+                    &mut self,
+                    ancestor,
+                    Some(next_in_path),
+                    AccessRelatedness::StrictChildAccess,
+                )?;
+            }
+        };
+        // All (potentially zero) ancestors have been explored, call f_propagate on start
+        stack.exec_and_visit(&mut self, start_idx, None, AccessRelatedness::This)?;
+        // up to this point we have never popped from `stack`, hence if the
+        // path to the root is `root = p(n) <- p(n-1)... <- p(1) <- p(0) = start`
+        // then now `stack` contains
+        // `[<children(p(n)) except p(n-1)> ... <children(p(1)) except p(0)> <children(p(0))>]`,
+        // all of which are for now unexplored.
+        // This is the starting point of a standard DFS which will thus
+        // explore all non-ancestors of `start` in the following order:
+        // - all descendants of `start`;
+        // - then the unexplored descendants of `parent(start)`;
+        // ...
+        // - until finally the unexplored descendants of `root`.
+        while let Some((tag, rel_pos)) = stack.pop() {
+            stack.exec_and_visit(&mut self, tag, None, rel_pos)?;
+        }
+        Ok(())
+    }
+}
+
 impl Tree {
     /// Create a new tree, with only a root pointer.
     pub fn new(root_tag: BorTag, size: Size) -> Self {
@@ -177,6 +322,200 @@ impl<'tcx> Tree {
         Ok(())
     }
 
+    /// Deallocation requires
+    /// - a pointer that permits write accesses
+    /// - the absence of Strong Protectors anywhere in the allocation
+    pub fn dealloc(
+        &mut self,
+        tag: BorTag,
+        range: AllocRange,
+        global: &GlobalState,
+    ) -> InterpResult<'tcx> {
+        self.perform_access(AccessKind::Write, tag, range, global)?;
+        let access_info = &self.nodes.get(self.tag_mapping.get(&tag).unwrap()).unwrap().debug_info;
+        for (_range, perms) in self.rperms.iter_mut(range.start, range.size) {
+            TreeVisitor { nodes: &self.nodes, tag_mapping: &self.tag_mapping, perms }
+                .traverse_parents_this_children_others(
+                    tag,
+                    |args: NodeAppArgs<'_>| -> Result<ContinueTraversal, TransitionError> {
+                        let NodeAppArgs { node, .. } = args;
+                        if global.borrow().protected_tags.get(&node.tag)
+                            == Some(&ProtectorKind::StrongProtector)
+                        {
+                            Err(TransitionError::ProtectedDealloc)
+                        } else {
+                            Ok(ContinueTraversal::Recurse)
+                        }
+                    },
+                    |args: ErrHandlerArgs<'_, TransitionError>| -> InterpErrorInfo<'tcx> {
+                        let ErrHandlerArgs { error_kind, faulty_tag } = args;
+                        TbError {
+                            faulty_tag,
+                            access_kind: AccessKind::Write,
+                            error_kind,
+                            tag_of_access: access_info,
+                        }
+                        .build()
+                    },
+                )?;
+        }
+        Ok(())
+    }
+
+    /// Maps the following propagation procedure to each range:
+    /// - initialize if needed;
+    /// - compute new state after transition;
+    /// - check that there is no protector that would forbid this;
+    /// - record this specific location as accessed.
+    pub fn perform_access(
+        &mut self,
+        access_kind: AccessKind,
+        tag: BorTag,
+        range: AllocRange,
+        global: &GlobalState,
+    ) -> InterpResult<'tcx> {
+        let access_info = &self.nodes.get(self.tag_mapping.get(&tag).unwrap()).unwrap().debug_info;
+        for (_range, perms) in self.rperms.iter_mut(range.start, range.size) {
+            TreeVisitor { nodes: &self.nodes, tag_mapping: &self.tag_mapping, perms }
+                .traverse_parents_this_children_others(
+                    tag,
+                    |args: NodeAppArgs<'_>| -> Result<ContinueTraversal, TransitionError> {
+                        let NodeAppArgs { node, mut perm, rel_pos } = args;
+
+                        let old_state =
+                            perm.or_insert_with(|| LocationState::new(node.default_initial_perm));
+
+                        // Optimize the tree traversal.
+                        // The optimization here consists of observing thanks to the tests
+                        // `foreign_read_is_noop_after_write` and `all_transitions_idempotent`
+                        // that if we apply twice in a row the effects of a foreign access
+                        // we can skip some branches.
+                        // "two foreign accesses in a row" occurs when `perm.latest_foreign_access` is `Some(_)`
+                        // AND the `rel_pos` of the current access corresponds to a foreign access.
+                        if rel_pos.is_foreign() {
+                            let new_access_noop =
+                                match (old_state.latest_foreign_access, access_kind) {
+                                    // Previously applied transition makes the new one a guaranteed
+                                    // noop in the two following cases:
+                                    // (1) justified by `foreign_read_is_noop_after_write`
+                                    (Some(AccessKind::Write), AccessKind::Read) => true,
+                                    // (2) justified by `all_transitions_idempotent`
+                                    (Some(old), new) if old == new => true,
+                                    // In all other cases there has been a recent enough
+                                    // child access that the effects of the new foreign access
+                                    // need to be applied to this subtree.
+                                    _ => false,
+                                };
+                            if new_access_noop {
+                                // Abort traversal if the new transition is indeed guaranteed
+                                // to be noop.
+                                return Ok(ContinueTraversal::SkipChildren);
+                            } else {
+                                // Otherwise propagate this time, and also record the
+                                // access that just occurred so that we can skip the propagation
+                                // next time.
+                                old_state.latest_foreign_access = Some(access_kind);
+                            }
+                        } else {
+                            // A child access occurred, this breaks the streak of "two foreign
+                            // accesses in a row" and we reset this field.
+                            old_state.latest_foreign_access = None;
+                        }
+
+                        let old_perm = old_state.permission;
+                        let protected = global.borrow().protected_tags.contains_key(&node.tag);
+                        let new_perm =
+                            Permission::perform_access(access_kind, rel_pos, old_perm, protected)
+                                .ok_or(TransitionError::ChildAccessForbidden(old_perm))?;
+                        if protected
+                            // Can't trigger Protector on uninitialized locations
+                            && old_state.initialized
+                            && !old_perm.protector_allows_transition(new_perm)
+                        {
+                            return Err(TransitionError::ProtectedTransition(old_perm, new_perm));
+                        }
+                        old_state.permission = new_perm;
+                        old_state.initialized |= !rel_pos.is_foreign();
+                        Ok(ContinueTraversal::Recurse)
+                    },
+                    |args: ErrHandlerArgs<'_, TransitionError>| -> InterpErrorInfo<'tcx> {
+                        let ErrHandlerArgs { error_kind, faulty_tag } = args;
+                        TbError { faulty_tag, access_kind, error_kind, tag_of_access: access_info }
+                            .build()
+                    },
+                )?;
+        }
+        Ok(())
+    }
+}
+
+/// Integration with the BorTag garbage collector
+impl Tree {
+    pub fn remove_unreachable_tags(&mut self, live_tags: &FxHashSet<BorTag>) {
+        assert!(self.keep_only_needed(self.root, live_tags)); // root can't be removed
+    }
+
+    /// 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.
+    ///
+    /// 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
+    /// children of dead tags to the nearest live parent, but it must be done with care
+    /// not to remove UB.
+    ///
+    /// Example: Consider the tree `root - parent - child`, with `parent: Frozen` and
+    /// `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 in self.rperms.iter_mut_all() {
+                perms.remove(idx);
+            }
+            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
+        }
+    }
+}
+
+impl VisitTags for Tree {
+    fn visit_tags(&self, visit: &mut dyn FnMut(BorTag)) {
+        // To ensure that the root never gets removed, we visit it
+        // (the `root` node of `Tree` is not an `Option<_>`)
+        visit(self.nodes.get(self.root).unwrap().tag)
+    }
+}
+
 /// Relative position of the access
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum AccessRelatedness {