make tree_visitor generic

royAmmerschuber · royAmmerschuber · commit 9622840f7819 · 2025-12-04T12:43:15.000+01:00
diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/tree.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/tree.rs
@@ -20,8 +20,7 @@ use smallvec::SmallVec;
 
 use super::Permission;
 use super::diagnostics::{
-    self, AccessCause, NodeDebugInfo, TbError, TransitionError,
-    no_valid_exposed_references_error,
+    self, AccessCause, NodeDebugInfo, TbError, TransitionError, no_valid_exposed_references_error,
 };
 use super::foreign_access_skipping::IdempotentForeignAccess;
 use super::perms::PermTransition;
@@ -562,42 +561,43 @@ impl<'tcx> Tree {
             // Checks the tree containing `idx` for strong protector violations.
             // It does this in traversal order.
             let mut check_tree = |idx| {
-                TreeVisitor { nodes: &mut self.nodes, loc }.traverse_this_parents_children_other(
-                    idx,
-                    // Visit all children, skipping none.
-                    |_| ContinueTraversal::Recurse,
-                    |args: NodeAppArgs<'_>| {
-                        let node = args.nodes.get(args.idx).unwrap();
+                TreeVisitor { nodes: &mut self.nodes, data: loc }
+                    .traverse_this_parents_children_other(
+                        idx,
+                        // Visit all children, skipping none.
+                        |_| ContinueTraversal::Recurse,
+                        |args: NodeAppArgs<'_, _>| {
+                            let node = args.nodes.get(args.idx).unwrap();
 
-                        let perm = args
-                            .loc
-                            .perms
-                            .get(args.idx)
-                            .copied()
-                            .unwrap_or_else(|| node.default_location_state());
-                        if global.borrow().protected_tags.get(&node.tag)
-                        == Some(&ProtectorKind::StrongProtector)
-                        // Don't check for protector if it is a Cell (see `unsafe_cell_deallocate` in `interior_mutability.rs`).
-                        // Related to https://github.com/rust-lang/rust/issues/55005.
-                        && !perm.permission.is_cell()
-                        // Only trigger UB if the accessed bit is set, i.e. if the protector is actually protecting this offset. See #4579.
-                        && perm.accessed
-                        {
-                            Err(TbError {
-                                conflicting_info: &node.debug_info,
-                                access_cause: diagnostics::AccessCause::Dealloc,
-                                alloc_id,
-                                error_offset: loc_range.start,
-                                error_kind: TransitionError::ProtectedDealloc,
-                                accessed_info: start_idx
-                                    .map(|idx| &args.nodes.get(idx).unwrap().debug_info),
+                            let perm = args
+                                .data
+                                .perms
+                                .get(args.idx)
+                                .copied()
+                                .unwrap_or_else(|| node.default_location_state());
+                            if global.borrow().protected_tags.get(&node.tag)
+                                == Some(&ProtectorKind::StrongProtector)
+                                // Don't check for protector if it is a Cell (see `unsafe_cell_deallocate` in `interior_mutability.rs`).
+                                // Related to https://github.com/rust-lang/rust/issues/55005.
+                                && !perm.permission.is_cell()
+                                // Only trigger UB if the accessed bit is set, i.e. if the protector is actually protecting this offset. See #4579.
+                                && perm.accessed
+                            {
+                                Err(TbError {
+                                    conflicting_info: &node.debug_info,
+                                    access_cause: diagnostics::AccessCause::Dealloc,
+                                    alloc_id,
+                                    error_offset: loc_range.start,
+                                    error_kind: TransitionError::ProtectedDealloc,
+                                    accessed_info: start_idx
+                                        .map(|idx| &args.nodes.get(idx).unwrap().debug_info),
+                                }
+                                .build())
+                            } else {
+                                Ok(())
                             }
-                            .build())
-                        } else {
-                            Ok(())
-                        }
-                    },
-                )
+                        },
+                    )
             };
             // If we have a start index we first check its subtree in traversal order.
             // This results in us showing the error of the closest node instead of an
@@ -967,16 +967,16 @@ impl<'tcx> LocationTree {
         //
         // `loc_range` is only for diagnostics (it is the range of
         // the `RangeMap` on which we are currently working).
-        let node_skipper = |args: &NodeAppArgs<'_>| -> ContinueTraversal {
+        let node_skipper = |args: &NodeAppArgs<'_, LocationTree>| -> ContinueTraversal {
             let node = args.nodes.get(args.idx).unwrap();
-            let perm = args.loc.perms.get(args.idx);
+            let perm = args.data.perms.get(args.idx);
 
             let old_state = perm.copied().unwrap_or_else(|| node.default_location_state());
             old_state.skip_if_known_noop(access_kind, args.rel_pos)
         };
-        let node_app = |args: NodeAppArgs<'_>| {
+        let node_app = |args: NodeAppArgs<'_, LocationTree>| {
             let node = args.nodes.get_mut(args.idx).unwrap();
-            let mut perm = args.loc.perms.entry(args.idx);
+            let mut perm = args.data.perms.entry(args.idx);
 
             let state = perm.or_insert(node.default_location_state());
 
@@ -985,7 +985,7 @@ impl<'tcx> LocationTree {
                 .perform_transition(
                     args.idx,
                     args.nodes,
-                    &mut args.loc.wildcard_accesses,
+                    &mut args.data.wildcard_accesses,
                     access_kind,
                     access_cause,
                     access_range,
@@ -1007,7 +1007,7 @@ impl<'tcx> LocationTree {
                 })
         };
 
-        let visitor = TreeVisitor { nodes, loc: self };
+        let visitor = TreeVisitor { nodes, data: self };
         match visit_children {
             ChildrenVisitMode::VisitChildrenOfAccessed =>
                 visitor.traverse_this_parents_children_other(access_source, node_skipper, node_app),
@@ -1061,16 +1061,16 @@ impl<'tcx> LocationTree {
         // marked as having an exposed foreign node, but actually that foreign node cannot be
         // the source of the access due to `max_local_tag`. The wildcard tracking cannot know
         // about `max_local_tag` so we will incorrectly assume that this might be a foreign access.
-        TreeVisitor { loc: self, nodes }.traverse_children_this(
+        TreeVisitor { data: self, nodes }.traverse_children_this(
             root,
             |args| -> ContinueTraversal {
                 let node = args.nodes.get(args.idx).unwrap();
-                let perm = args.loc.perms.get(args.idx);
+                let perm = args.data.perms.get(args.idx);
 
                 let old_state = perm.copied().unwrap_or_else(|| node.default_location_state());
                 // If we know where, relative to this node, the wildcard access occurs,
                 // then check if we can skip the entire subtree.
-                if let Some(relatedness) = get_relatedness(args.idx, node, args.loc)
+                if let Some(relatedness) = get_relatedness(args.idx, node, args.data)
                     && let Some(relatedness) = relatedness.to_relatedness()
                 {
                     // We can use the usual SIFA machinery to skip nodes.
@@ -1084,7 +1084,7 @@ impl<'tcx> LocationTree {
 
                 let protected = global.borrow().protected_tags.contains_key(&node.tag);
 
-                let Some(wildcard_relatedness) = get_relatedness(args.idx, node, args.loc) else {
+                let Some(wildcard_relatedness) = get_relatedness(args.idx, node, args.data) else {
                     // There doesn't exist a valid exposed reference for this access to
                     // happen through.
                     // This can only happen if `root` is the main root: We set
@@ -1105,13 +1105,13 @@ impl<'tcx> LocationTree {
                     return Ok(());
                 };
 
-                let mut entry = args.loc.perms.entry(args.idx);
+                let mut entry = args.data.perms.entry(args.idx);
                 let perm = entry.or_insert(node.default_location_state());
                 // We know the exact relatedness, so we can actually do precise checks.
                 perm.perform_transition(
                     args.idx,
                     args.nodes,
-                    &mut args.loc.wildcard_accesses,
+                    &mut args.data.wildcard_accesses,
                     access_kind,
                     access_cause,
                     access_range,
diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/tree_visitor.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/tree_visitor.rs
@@ -1,23 +1,25 @@
-use super::tree::{AccessRelatedness, LocationTree, Node};
+use std::marker::PhantomData;
+
+use super::tree::{AccessRelatedness, Node};
 use super::unimap::{UniIndex, UniValMap};
 
 /// Data given to the transition function
-pub struct NodeAppArgs<'visit> {
+pub struct NodeAppArgs<'visit, T> {
     /// The index of the current node.
     pub idx: UniIndex,
     /// Relative position of the access.
     pub rel_pos: AccessRelatedness,
     /// The node map of this tree.
     pub nodes: &'visit mut UniValMap<Node>,
-    /// The permissions map of this tree.
-    pub loc: &'visit mut LocationTree,
+    /// Additional data we want to be able to modify in f_propagate and read in f_continue.
+    pub data: &'visit mut T,
 }
 /// Internal contents of `Tree` with the minimum of mutable access for
 /// For soundness do not modify the children or parent indexes of nodes
 /// during traversal.
-pub struct TreeVisitor<'tree> {
+pub struct TreeVisitor<'tree, T> {
     pub nodes: &'tree mut UniValMap<Node>,
-    pub loc: &'tree mut LocationTree,
+    pub data: &'tree mut T,
 }
 
 /// Whether to continue exploring the children recursively or not.
@@ -41,7 +43,7 @@ enum RecursionState {
 /// Stack of nodes left to explore in a tree traversal.
 /// See the docs of `traverse_this_parents_children_other` for details on the
 /// traversal order.
-struct TreeVisitorStack<NodeContinue, NodeApp> {
+struct TreeVisitorStack<NodeContinue, NodeApp, T> {
     /// Function describing whether to continue at a tag.
     /// This is only invoked for foreign accesses.
     f_continue: NodeContinue,
@@ -56,36 +58,37 @@ struct TreeVisitorStack<NodeContinue, NodeApp> {
     /// This is just an artifact of how you hand-roll recursion,
     /// it does not have a deeper meaning otherwise.
     stack: Vec<(UniIndex, AccessRelatedness, RecursionState)>,
+    phantom: PhantomData<T>,
 }
 
-impl<NodeContinue, NodeApp, Err> TreeVisitorStack<NodeContinue, NodeApp>
+impl<NodeContinue, NodeApp, T, Err> TreeVisitorStack<NodeContinue, NodeApp, T>
 where
-    NodeContinue: Fn(&NodeAppArgs<'_>) -> ContinueTraversal,
-    NodeApp: FnMut(NodeAppArgs<'_>) -> Result<(), Err>,
+    NodeContinue: Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
+    NodeApp: FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
 {
     fn should_continue_at(
         &self,
-        this: &mut TreeVisitor<'_>,
+        this: &mut TreeVisitor<'_, T>,
         idx: UniIndex,
         rel_pos: AccessRelatedness,
     ) -> ContinueTraversal {
-        let args = NodeAppArgs { idx, rel_pos, nodes: this.nodes, loc: this.loc };
+        let args = NodeAppArgs { idx, rel_pos, nodes: this.nodes, data: this.data };
         (self.f_continue)(&args)
     }
 
     fn propagate_at(
         &mut self,
-        this: &mut TreeVisitor<'_>,
+        this: &mut TreeVisitor<'_, T>,
         idx: UniIndex,
         rel_pos: AccessRelatedness,
     ) -> Result<(), Err> {
-        (self.f_propagate)(NodeAppArgs { idx, rel_pos, nodes: this.nodes, loc: this.loc })
+        (self.f_propagate)(NodeAppArgs { idx, rel_pos, nodes: this.nodes, data: this.data })
     }
 
     /// Returns the root of this tree.
     fn go_upwards_from_accessed(
         &mut self,
-        this: &mut TreeVisitor<'_>,
+        this: &mut TreeVisitor<'_, T>,
         accessed_node: UniIndex,
         visit_children: ChildrenVisitMode,
     ) -> Result<UniIndex, Err> {
@@ -136,7 +139,7 @@ where
         Ok(last_node)
     }
 
-    fn finish_foreign_accesses(&mut self, this: &mut TreeVisitor<'_>) -> Result<(), Err> {
+    fn finish_foreign_accesses(&mut self, this: &mut TreeVisitor<'_, T>) -> Result<(), Err> {
         while let Some((idx, rel_pos, step)) = self.stack.last_mut() {
             let idx = *idx;
             let rel_pos = *rel_pos;
@@ -173,11 +176,11 @@ where
     }
 
     fn new(f_continue: NodeContinue, f_propagate: NodeApp) -> Self {
-        Self { f_continue, f_propagate, stack: Vec::new() }
+        Self { f_continue, f_propagate, stack: Vec::new(), phantom: PhantomData }
     }
 }
 
-impl<'tree> TreeVisitor<'tree> {
+impl<'tree, T> TreeVisitor<'tree, T> {
     /// Applies `f_propagate` to every vertex of the tree in a piecewise bottom-up way: First, visit
     /// all ancestors of `start_idx` (starting with `start_idx` itself), then children of `start_idx`, then the rest,
     /// going bottom-up in each of these two "pieces" / sections.
@@ -219,8 +222,8 @@ impl<'tree> TreeVisitor<'tree> {
     pub fn traverse_this_parents_children_other<Err>(
         mut self,
         start_idx: UniIndex,
-        f_continue: impl Fn(&NodeAppArgs<'_>) -> ContinueTraversal,
-        f_propagate: impl FnMut(NodeAppArgs<'_>) -> Result<(), Err>,
+        f_continue: impl Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
+        f_propagate: impl FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
     ) -> Result<UniIndex, Err> {
         let mut stack = TreeVisitorStack::new(f_continue, f_propagate);
         // Visits the accessed node itself, and all its parents, i.e. all nodes
@@ -245,8 +248,8 @@ impl<'tree> TreeVisitor<'tree> {
     pub fn traverse_nonchildren<Err>(
         mut self,
         start_idx: UniIndex,
-        f_continue: impl Fn(&NodeAppArgs<'_>) -> ContinueTraversal,
-        f_propagate: impl FnMut(NodeAppArgs<'_>) -> Result<(), Err>,
+        f_continue: impl Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
+        f_propagate: impl FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
     ) -> Result<UniIndex, Err> {
         let mut stack = TreeVisitorStack::new(f_continue, f_propagate);
         // Visits the accessed node itself, and all its parents, i.e. all nodes
@@ -271,8 +274,8 @@ impl<'tree> TreeVisitor<'tree> {
     pub fn traverse_children_this<Err>(
         mut self,
         start_idx: UniIndex,
-        f_continue: impl Fn(&NodeAppArgs<'_>) -> ContinueTraversal,
-        f_propagate: impl FnMut(NodeAppArgs<'_>) -> Result<(), Err>,
+        f_continue: impl Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
+        f_propagate: impl FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
     ) -> Result<(), Err> {
         let mut stack = TreeVisitorStack::new(f_continue, f_propagate);
 
