bitstring_trees/
walk_mut.rs

1//! Walk tree structures without call stack
2
3use core::{
4	marker::PhantomData,
5	ptr::NonNull,
6};
7
8use alloc::vec::Vec;
9
10/// Allows different node and tree types in [`WalkMut`].
11pub enum NodeOrTree<T, N> {
12	/// [`WalkMut`] is currently at a node
13	Node(N),
14	/// [`WalkMut`] is currently at tree
15	Tree(T),
16}
17
18impl<T, N> NodeOrTree<T, N> {
19	/// Map tree value
20	pub fn map_tree<F, U>(self, f: F) -> NodeOrTree<U, N>
21	where
22		F: FnOnce(T) -> U,
23	{
24		match self {
25			Self::Tree(r) => NodeOrTree::Tree(f(r)),
26			Self::Node(n) => NodeOrTree::Node(n),
27		}
28	}
29
30	/// Map node value
31	pub fn map_node<F, U>(self, f: F) -> NodeOrTree<T, U>
32	where
33		F: FnOnce(N) -> U,
34	{
35		match self {
36			Self::Tree(r) => NodeOrTree::Tree(r),
37			Self::Node(n) => NodeOrTree::Node(f(n)),
38		}
39	}
40
41	/// Return node
42	pub fn node(self) -> Option<N> {
43		match self {
44			Self::Tree(_) => None,
45			Self::Node(n) => Some(n),
46		}
47	}
48}
49
50impl<N> NodeOrTree<N, N> {
51	/// If tree and node type are equivalent extract inner type.
52	#[inline]
53	pub fn flatten(self) -> N {
54		match self {
55			Self::Tree(r) => r,
56			Self::Node(n) => n,
57		}
58	}
59}
60
61impl<N> NodeOrTree<Option<N>, N> {
62	/// If tree and node type are equivalent extract inner type.
63	#[inline]
64	pub fn flatten_optional(self) -> Option<N> {
65		match self {
66			Self::Tree(r) => r,
67			Self::Node(n) => Some(n),
68		}
69	}
70}
71
72/// Walk tree structures without call stack
73///
74/// Walking tree structures with mutable references usually
75/// requires a recursive call stack to make the borrow-checker
76/// happy.
77///
78/// This uses a stack ([`Vec`]) to keep track of the "current"
79/// mutable reference (and hiding the previous ones).
80///
81/// (There is no way to implement this without `unsafe`, but the
82/// abstraction should be safe.)
83///
84/// Each nested level can also track additional value of type `A`.
85pub struct WalkMut<'r, T: ?Sized, N: ?Sized, A = ()> {
86	_lifetime: PhantomData<&'r mut T>,
87	tree: NonNull<T>,
88	stack: Vec<(NonNull<N>, A)>,
89}
90
91impl<'r, T: ?Sized, N: ?Sized, A> WalkMut<'r, T, N, A> {
92	/// Start a new tree walk at a tree
93	pub fn new(tree: &'r mut T) -> Self {
94		Self {
95			_lifetime: PhantomData,
96			tree: tree.into(),
97			stack: Vec::new(),
98		}
99	}
100
101	/// Walk down the tree one step
102	///
103	/// The step can fail by returning [`Err`].
104	pub fn try_walk<F, E>(&mut self, with: F) -> Result<(), E>
105	where
106		F: for<'n> FnOnce(NodeOrTree<&'n mut T, &'n mut N>) -> Result<(&'n mut N, A), E>,
107	{
108		match with(self.current_mut()) {
109			Err(e) => Err(e),
110			Ok((next, add)) => {
111				let next: NonNull<N> = next.into();
112				self.stack.push((next, add));
113				Ok(())
114			},
115		}
116	}
117
118	/// Walk up to the previous level.
119	///
120	/// Returns the associated data stored with the step,
121	/// or [`None`] if already at the initial tree.
122	pub fn pop(&mut self) -> Option<A> {
123		Some(self.stack.pop()?.1)
124	}
125
126	/// Walk up to tree
127	pub fn pop_all(&mut self) -> &mut T {
128		self.stack.clear();
129		unsafe { self.tree.as_mut() }
130	}
131
132	/// Get mutable reference to current node or tree
133	///
134	/// If you need the result to outlive the destruction of the [`WalkMut`] value, see [`into_current_mut`].
135	///
136	/// [`into_current_mut`]: WalkMut::into_current_mut
137	pub fn current_mut(&mut self) -> NodeOrTree<&mut T, &mut N> {
138		if let Some((cur, _)) = self.stack.last_mut() {
139			NodeOrTree::Node(unsafe { cur.as_mut() })
140		} else {
141			NodeOrTree::Tree(unsafe { self.tree.as_mut() })
142		}
143	}
144
145	/// Extract mutable reference to current node or tree
146	///
147	/// Also see [`current_mut`]
148	///
149	/// [`current_mut`]: WalkMut::current_mut
150	pub fn into_current_mut(mut self) -> NodeOrTree<&'r mut T, &'r mut N> {
151		// safety: dropping stack of references means nothing else can create
152		// new references to them while 'r still blocks new references to the root.
153		if let Some((cur, _)) = self.stack.last_mut() {
154			NodeOrTree::Node(unsafe { cur.as_mut() })
155		} else {
156			NodeOrTree::Tree(unsafe { self.tree.as_mut() })
157		}
158	}
159
160	/// Extract mutable reference to tree
161	pub fn into_tree_mut(mut self) -> &'r mut T {
162		self.stack.clear();
163		// safety: same as pop_all() + into_current_mut() (which must return the tree)
164		unsafe { self.tree.as_mut() }
165	}
166
167	/// Get reference to current node or tree
168	pub fn current(&self) -> NodeOrTree<&T, &N> {
169		if let Some((cur, _)) = self.stack.last() {
170			NodeOrTree::Node(unsafe { cur.as_ref() })
171		} else {
172			NodeOrTree::Tree(unsafe { self.tree.as_ref() })
173		}
174	}
175}
bitstring_trees/walk_mut.rs

bitstring_trees/
walk_mut.rs
