miri/borrow_tracker/tree_borrows/tree_visitor.rs
1use std::marker::PhantomData;
2
3use super::tree::{AccessRelatedness, Node};
4use super::unimap::{UniIndex, UniValMap};
5
6/// Data given to the transition function
7pub struct NodeAppArgs<'visit, T> {
8 /// The index of the current node.
9 pub idx: UniIndex,
10 /// Relative position of the access.
11 pub rel_pos: AccessRelatedness,
12 /// The node map of this tree.
13 pub nodes: &'visit mut UniValMap<Node>,
14 /// Additional data we want to be able to modify in f_propagate and read in f_continue.
15 pub data: &'visit mut T,
16}
17/// Internal contents of `Tree` with the minimum of mutable access for
18/// For soundness do not modify the children or parent indexes of nodes
19/// during traversal.
20pub struct TreeVisitor<'tree, T> {
21 pub nodes: &'tree mut UniValMap<Node>,
22 pub data: &'tree mut T,
23}
24
25/// Whether to continue exploring the children recursively or not.
26#[derive(Debug)]
27pub enum ContinueTraversal {
28 Recurse,
29 SkipSelfAndChildren,
30}
31
32#[derive(Clone, Copy, Debug)]
33pub enum ChildrenVisitMode {
34 VisitChildrenOfAccessed,
35 SkipChildrenOfAccessed,
36}
37
38enum RecursionState {
39 BeforeChildren,
40 AfterChildren,
41}
42
43/// Stack of nodes left to explore in a tree traversal.
44/// See the docs of `traverse_this_parents_children_other` for details on the
45/// traversal order.
46struct TreeVisitorStack<NodeContinue, NodeApp, T> {
47 /// Function describing whether to continue at a tag.
48 /// This is only invoked for foreign accesses.
49 f_continue: NodeContinue,
50 /// Function to apply to each tag.
51 f_propagate: NodeApp,
52 /// Mutable state of the visit: the tags left to handle.
53 /// Every tag pushed should eventually be handled,
54 /// and the precise order is relevant for diagnostics.
55 /// Since the traversal is piecewise bottom-up, we need to
56 /// remember whether we're here initially, or after visiting all children.
57 /// The last element indicates this.
58 /// This is just an artifact of how you hand-roll recursion,
59 /// it does not have a deeper meaning otherwise.
60 stack: Vec<(UniIndex, AccessRelatedness, RecursionState)>,
61 phantom: PhantomData<T>,
62}
63
64impl<NodeContinue, NodeApp, T, Err> TreeVisitorStack<NodeContinue, NodeApp, T>
65where
66 NodeContinue: Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
67 NodeApp: FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
68{
69 fn should_continue_at(
70 &self,
71 this: &mut TreeVisitor<'_, T>,
72 idx: UniIndex,
73 rel_pos: AccessRelatedness,
74 ) -> ContinueTraversal {
75 let args = NodeAppArgs { idx, rel_pos, nodes: this.nodes, data: this.data };
76 (self.f_continue)(&args)
77 }
78
79 fn propagate_at(
80 &mut self,
81 this: &mut TreeVisitor<'_, T>,
82 idx: UniIndex,
83 rel_pos: AccessRelatedness,
84 ) -> Result<(), Err> {
85 (self.f_propagate)(NodeAppArgs { idx, rel_pos, nodes: this.nodes, data: this.data })
86 }
87
88 /// Returns the root of this tree.
89 fn go_upwards_from_accessed(
90 &mut self,
91 this: &mut TreeVisitor<'_, T>,
92 accessed_node: UniIndex,
93 visit_children: ChildrenVisitMode,
94 ) -> Result<UniIndex, Err> {
95 // We want to visit the accessed node's children first.
96 // However, we will below walk up our parents and push their children (our cousins)
97 // onto the stack. To ensure correct iteration order, this method thus finishes
98 // by reversing the stack. This only works if the stack is empty initially.
99 assert!(self.stack.is_empty());
100 // First, handle accessed node. A bunch of things need to
101 // be handled differently here compared to the further parents
102 // of `accesssed_node`.
103 {
104 self.propagate_at(this, accessed_node, AccessRelatedness::LocalAccess)?;
105 if matches!(visit_children, ChildrenVisitMode::VisitChildrenOfAccessed) {
106 let accessed_node = this.nodes.get(accessed_node).unwrap();
107 // We `rev()` here because we reverse the entire stack later.
108 for &child in accessed_node.children.iter().rev() {
109 self.stack.push((
110 child,
111 AccessRelatedness::ForeignAccess,
112 RecursionState::BeforeChildren,
113 ));
114 }
115 }
116 }
117 // Then, handle the accessed node's parents. Here, we need to
118 // make sure we only mark the "cousin" subtrees for later visitation,
119 // not the subtree that contains the accessed node.
120 let mut last_node = accessed_node;
121 while let Some(current) = this.nodes.get(last_node).unwrap().parent {
122 self.propagate_at(this, current, AccessRelatedness::LocalAccess)?;
123 let node = this.nodes.get(current).unwrap();
124 // We `rev()` here because we reverse the entire stack later.
125 for &child in node.children.iter().rev() {
126 if last_node == child {
127 continue;
128 }
129 self.stack.push((
130 child,
131 AccessRelatedness::ForeignAccess,
132 RecursionState::BeforeChildren,
133 ));
134 }
135 last_node = current;
136 }
137 // Reverse the stack, as discussed above.
138 self.stack.reverse();
139 Ok(last_node)
140 }
141
142 fn finish_foreign_accesses(&mut self, this: &mut TreeVisitor<'_, T>) -> Result<(), Err> {
143 while let Some((idx, rel_pos, step)) = self.stack.last_mut() {
144 let idx = *idx;
145 let rel_pos = *rel_pos;
146 match *step {
147 // How to do bottom-up traversal, 101: Before you handle a node, you handle all children.
148 // For this, you must first find the children, which is what this code here does.
149 RecursionState::BeforeChildren => {
150 // Next time we come back will be when all the children are handled.
151 *step = RecursionState::AfterChildren;
152 // Now push the children, except if we are told to skip this subtree.
153 let handle_children = self.should_continue_at(this, idx, rel_pos);
154 match handle_children {
155 ContinueTraversal::Recurse => {
156 let node = this.nodes.get(idx).unwrap();
157 for &child in node.children.iter() {
158 self.stack.push((child, rel_pos, RecursionState::BeforeChildren));
159 }
160 }
161 ContinueTraversal::SkipSelfAndChildren => {
162 // skip self
163 self.stack.pop();
164 continue;
165 }
166 }
167 }
168 // All the children are handled, let's actually visit this node
169 RecursionState::AfterChildren => {
170 self.stack.pop();
171 self.propagate_at(this, idx, rel_pos)?;
172 }
173 }
174 }
175 Ok(())
176 }
177
178 fn new(f_continue: NodeContinue, f_propagate: NodeApp) -> Self {
179 Self { f_continue, f_propagate, stack: Vec::new(), phantom: PhantomData }
180 }
181}
182
183impl<'tree, T> TreeVisitor<'tree, T> {
184 /// Applies `f_propagate` to every vertex of the tree in a piecewise bottom-up way: First, visit
185 /// all ancestors of `start_idx` (starting with `start_idx` itself), then children of `start_idx`, then the rest,
186 /// going bottom-up in each of these two "pieces" / sections.
187 /// This ensures that errors are triggered in the following order
188 /// - first invalid accesses with insufficient permissions, closest to the accessed node first,
189 /// - then protector violations, bottom-up, starting with the children of the accessed node, and then
190 /// going upwards and outwards.
191 ///
192 /// The following graphic visualizes it, with numbers indicating visitation order and `start_idx` being
193 /// the node that is visited first ("1"):
194 ///
195 /// ```text
196 /// 3
197 /// /|
198 /// / |
199 /// 9 2
200 /// | |\
201 /// | | \
202 /// 8 1 7
203 /// / \
204 /// 4 6
205 /// |
206 /// 5
207 /// ```
208 ///
209 /// `f_propagate` should follow the following format: for a given `Node` it updates its
210 /// `Permission` depending on the position relative to `start_idx` (given by an
211 /// `AccessRelatedness`).
212 /// `f_continue` is called earlier on foreign nodes, and describes whether to even start
213 /// visiting the subtree at that node. If it e.g. returns `SkipSelfAndChildren` on node 6
214 /// above, then nodes 5 _and_ 6 would not be visited by `f_propagate`. It is not used for
215 /// notes having a child access (nodes 1, 2, 3).
216 ///
217 /// Finally, remember that the iteration order is not relevant for UB, it only affects
218 /// diagnostics. It also affects tree traversal optimizations built on top of this, so
219 /// those need to be reviewed carefully as well whenever this changes.
220 ///
221 /// Returns the index of the root of the accessed tree.
222 pub fn traverse_this_parents_children_other<Err>(
223 mut self,
224 start_idx: UniIndex,
225 f_continue: impl Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
226 f_propagate: impl FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
227 ) -> Result<UniIndex, Err> {
228 let mut stack = TreeVisitorStack::new(f_continue, f_propagate);
229 // Visits the accessed node itself, and all its parents, i.e. all nodes
230 // undergoing a child access. Also pushes the children and the other
231 // cousin nodes (i.e. all nodes undergoing a foreign access) to the stack
232 // to be processed later.
233 let root = stack.go_upwards_from_accessed(
234 &mut self,
235 start_idx,
236 ChildrenVisitMode::VisitChildrenOfAccessed,
237 )?;
238 // Now visit all the foreign nodes we remembered earlier.
239 // For this we go bottom-up, but also allow f_continue to skip entire
240 // subtrees from being visited if it would be a NOP.
241 stack.finish_foreign_accesses(&mut self)?;
242 Ok(root)
243 }
244
245 /// Like `traverse_this_parents_children_other`, but skips the children of `start_idx`.
246 ///
247 /// Returns the index of the root of the accessed tree.
248 pub fn traverse_nonchildren<Err>(
249 mut self,
250 start_idx: UniIndex,
251 f_continue: impl Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
252 f_propagate: impl FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
253 ) -> Result<UniIndex, Err> {
254 let mut stack = TreeVisitorStack::new(f_continue, f_propagate);
255 // Visits the accessed node itself, and all its parents, i.e. all nodes
256 // undergoing a child access. Also pushes the other cousin nodes to the
257 // stack, but not the children of the accessed node.
258 let root = stack.go_upwards_from_accessed(
259 &mut self,
260 start_idx,
261 ChildrenVisitMode::SkipChildrenOfAccessed,
262 )?;
263 // Now visit all the foreign nodes we remembered earlier.
264 // For this we go bottom-up, but also allow f_continue to skip entire
265 // subtrees from being visited if it would be a NOP.
266 stack.finish_foreign_accesses(&mut self)?;
267 Ok(root)
268 }
269
270 /// Traverses all children of `start_idx` including `start_idx` itself.
271 /// Uses `f_continue` to filter out subtrees and then processes each node
272 /// with `f_propagate` so that the children get processed before their
273 /// parents.
274 pub fn traverse_children_this<Err>(
275 mut self,
276 start_idx: UniIndex,
277 f_continue: impl Fn(&NodeAppArgs<'_, T>) -> ContinueTraversal,
278 f_propagate: impl FnMut(NodeAppArgs<'_, T>) -> Result<(), Err>,
279 ) -> Result<(), Err> {
280 let mut stack = TreeVisitorStack::new(f_continue, f_propagate);
281
282 stack.stack.push((
283 start_idx,
284 AccessRelatedness::ForeignAccess,
285 RecursionState::BeforeChildren,
286 ));
287 stack.finish_foreign_accesses(&mut self)
288 }
289}