generic_cursors/
with_data.rs

1use std::{collections::VecDeque, marker::PhantomData};
2
3pub struct MutRefStackWithData<'root, T: ?Sized, U: 'root> {
4    /// Ensures this mutrefstack does not exceed the lifetime of its root.
5    lifetime: PhantomData<(&'root mut T, U)>,
6    /// The stack of pointers. Each one borrows from the one prior, except the first which is the `root` and may never be popped.
7    data: Vec<(*mut T, U)>,
8}
9
10pub enum MoveDecision<'root, 'this, T: ?Sized, U: 'root> {
11    Ascend,
12    Stay,
13    Descend(&'this mut T, U),
14    Inject(&'root mut T, U),
15}
16
17pub enum MoveError {
18    AscendAtRoot,
19}
20
21impl<'root, T: ?Sized, U> MutRefStackWithData<'root, T, U> {
22    /// Create a new MutRefStack from a mutable reference to the root
23    /// of a recursive data structure.
24    pub fn new(root: &'root mut T, additional_data: U) -> Self {
25        Self {
26            lifetime: PhantomData,
27            data: vec![(root, additional_data)],
28        }
29    }
30
31    /// Obtain a shared reference to the top of the stack.
32    pub fn top(&self) -> (&T, &U) {
33        let &(ptr, ref additional_data) = self
34            .data
35            .last()
36            .expect("root pointer should never be popped");
37        let ptr: *const T = ptr;
38        (unsafe { &(*ptr) }, additional_data)
39    }
40
41    /// Obtain a mutable reference to the top of the stack.
42    pub fn top_mut(&mut self) -> (&mut T, &mut U) {
43        let &mut (ptr, ref mut additional_data) = self
44            .data
45            .last_mut()
46            .expect("root pointer should never be popped");
47        (unsafe { &mut (*ptr) }, additional_data)
48    }
49
50    /// Is this MutRefStack currently at its root?
51    pub fn is_at_root(&self) -> bool {
52        self.data.len() == 1
53    }
54
55    /// Descend into the recursive data structure, returning a mutable reference to the new top element.
56    /// Rust's borrow checker enforces that the closure cannot inject any lifetime (other than `'static`),
57    /// because the closure must work for any lifetime `'node`.
58    pub fn descend_with(
59        &mut self,
60        f: impl for<'node, 'addl> FnOnce(&'node mut T, &'addl mut U) -> Option<(&'node mut T, U)>,
61    ) -> Option<(&mut T, &mut U)> {
62        let &mut (ptr, ref mut addl) = self
63            .data
64            .last_mut()
65            .expect("root pointer should never be popped");
66        {
67            let node = unsafe { &mut *ptr };
68            let (desc, new_addl) = f(node, addl)?;
69            self.data.push((desc, new_addl));
70        }
71        Some(self.top_mut())
72    }
73
74    /// Inject a new reference to the top of the stack. The reference still must live
75    /// as long as the root of the stack.
76    pub fn inject_with(
77        &mut self,
78        f: impl for<'node, 'addl> FnOnce(&'node mut T, &'addl mut U) -> Option<(&'root mut T, U)>,
79    ) -> Option<(&mut T, &mut U)> {
80        let (top, addl) = self.top_mut();
81        let (new_top, new_addl) = f(top, addl)?;
82        self.data.push((new_top, new_addl));
83        Some(self.top_mut())
84    }
85
86    /// Ascend back up from the recursive data structure, returning a mutable reference to the new top element, if it changed.
87    /// If we are not currently at the root, ascend and return a reference to the new top, a reference to the new top's additional data, and the old top's additional data.
88    /// If we are already the root, returns None (the top is the root and does not change).
89    pub fn ascend(&mut self) -> Option<((&mut T, &mut U), U)> {
90        match self.data.len() {
91            0 => unreachable!("root pointer must always exist"),
92            1 => None,
93            _ => {
94                let Some((_ptr, addl)) = self.data.pop() else { unreachable!() };
95                Some((self.top_mut(), addl))
96            }
97        }
98    }
99
100    /// Ascend back up from the recursive data structure while the given closure returns `true`, returning a mutable reference to the new top element.
101    /// If we are not currently at the root, and the predicate returns `true`, ascend and continue.
102    /// If we are already at the root, or if the predicate returned false, returns a reference to the top element.
103    pub fn ascend_while<P>(
104        &mut self,
105        mut predicate: P,
106    ) -> ((&mut T, &mut U), impl IntoIterator<Item = U>)
107    where
108        P: FnMut(&mut T, &mut U) -> bool,
109    {
110        let mut items = VecDeque::new();
111        while !self.is_at_root() {
112            let (top, addl) = self.top_mut();
113            if !predicate(top, addl) {
114                break;
115            }
116            let Some((_top, addl)) = self.ascend() else {
117                unreachable!()
118            };
119            items.push_front(addl);
120        }
121        (self.top_mut(), items)
122    }
123
124    /// Ascend from, descend from, inject a new top, or stay at the current node,
125    /// based on the return value of the closure.
126    pub fn move_with(
127        &mut self,
128        f: impl for<'node, 'addl> FnOnce(&'node mut T, &'addl mut U) -> MoveDecision<'root, 'node, T, U>,
129    ) -> Result<((&mut T, &mut U), Option<U>), MoveError> {
130        let top = self.top_mut();
131        let result = f(top.0, top.1);
132        match result {
133            MoveDecision::Ascend => {
134                let (top, old_addl) = self.ascend().ok_or(MoveError::AscendAtRoot)?;
135                Ok((top, Some(old_addl)))
136            }
137            MoveDecision::Stay => Ok((self.top_mut(), None)),
138            MoveDecision::Inject(new_top, new_addl) | MoveDecision::Descend(new_top, new_addl) => {
139                let new_top: *mut T = new_top;
140                self.data.push((new_top, new_addl));
141                Ok((self.top_mut(), None))
142            }
143        }
144    }
145
146    /// Return reference to the top element of this stack, forgetting about the stack entirely.
147    pub fn into_top(self) -> &'root mut T {
148        let ptr = self.data.last().unwrap().0;
149        unsafe { &mut *ptr }
150    }
151}