1use std::{collections::VecDeque, marker::PhantomData};
2
3pub struct MutRefStackWithData<'root, T: ?Sized, U: 'root> {
4 lifetime: PhantomData<(&'root mut T, U)>,
6 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 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 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 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 pub fn is_at_root(&self) -> bool {
52 self.data.len() == 1
53 }
54
55 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 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 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 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 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 pub fn into_top(self) -> &'root mut T {
148 let ptr = self.data.last().unwrap().0;
149 unsafe { &mut *ptr }
150 }
151}