aatree/node/
mod.rs

1//! Low-level implementation of an AA tree. You shouldn't have to use this directly; instead, use
2//! the implementations in [`AATreeSet`](crate::AATreeSet) and [`AATreeMap`](crate::AATreeMap).
3
4use alloc::boxed::Box;
5use core::mem;
6
7mod insert;
8mod remove;
9mod traverse;
10
11pub use traverse::*;
12
13#[derive(Clone, Debug, PartialEq)]
14pub struct AANode<T>(Option<Box<Node<T>>>);
15
16#[derive(Clone, Debug, PartialEq)]
17pub(super) struct Node<T> {
18	pub(super) level: u8,
19	pub(super) content: T,
20	pub(super) left_child: AANode<T>,
21	pub(super) right_child: AANode<T>
22}
23
24impl<T> From<Node<T>> for AANode<T> {
25	fn from(node: Node<T>) -> Self {
26		Self(Some(Box::new(node)))
27	}
28}
29
30impl<T> Default for AANode<T> {
31	fn default() -> Self {
32		Self::new()
33	}
34}
35
36impl<T> From<T> for AANode<T> {
37	fn from(content: T) -> Self {
38		Node {
39			level: 1,
40			content,
41			left_child: Self(None),
42			right_child: Self(None)
43		}
44		.into()
45	}
46}
47
48impl<T> AANode<T> {
49	pub(super) fn unbox(self) -> Option<Node<T>> {
50		self.0.map(|this| *this)
51	}
52
53	pub(super) fn as_ref(&self) -> Option<&Node<T>> {
54		self.0.as_ref().map(Box::as_ref)
55	}
56
57	fn as_mut(&mut self) -> Option<&mut Node<T>> {
58		self.0.as_mut().map(Box::as_mut)
59	}
60
61	fn take(&mut self) -> Self {
62		Self(self.0.take())
63	}
64
65	/// Create a new `Nil` node.
66	pub const fn new() -> Self {
67		Self(None)
68	}
69
70	/// Return true if this node is `Nil`.
71	pub const fn is_nil(&self) -> bool {
72		self.0.is_none()
73	}
74
75	/// Return true if this node is a leaf.
76	pub fn is_leaf(&self) -> bool {
77		match self.as_ref() {
78			None => false,
79			Some(Node {
80				left_child,
81				right_child,
82				..
83			}) => left_child.is_nil() && right_child.is_nil()
84		}
85	}
86
87	/// Return true if this node has a left child.
88	pub fn has_left_child(&self) -> bool {
89		match self.as_ref() {
90			None => false,
91			Some(Node { left_child, .. }) => !left_child.is_nil()
92		}
93	}
94
95	/// Return true if this node has a right child.
96	pub fn has_right_child(&self) -> bool {
97		match self.as_ref() {
98			None => false,
99			Some(Node { right_child, .. }) => !right_child.is_nil()
100		}
101	}
102
103	fn left_child_mut(&mut self) -> Option<&mut Self> {
104		self.as_mut().and_then(|Node { left_child, .. }| {
105			(!left_child.is_nil()).then(|| left_child)
106		})
107	}
108
109	fn right_child_mut(&mut self) -> Option<&mut Self> {
110		self.as_mut().and_then(|Node { right_child, .. }| {
111			(!right_child.is_nil()).then(|| right_child)
112		})
113	}
114
115	fn set_right_child(&mut self, child: Self) {
116		match self.as_mut() {
117			None => panic!("I don't have a right child"),
118			Some(Node { right_child, .. }) => {
119				*right_child = child;
120			}
121		}
122	}
123
124	pub(super) fn level(&self) -> u8 {
125		match self.as_ref() {
126			None => 0,
127			Some(Node { level, .. }) => *level
128		}
129	}
130
131	fn content_mut(&mut self) -> Option<&mut T> {
132		self.as_mut().map(|Node { content, .. }| content)
133	}
134
135	/// Update the level of this node. **Panic** if the node is [`Nil`](Self::Nil).
136	fn set_level(&mut self, level: u8) {
137		match self.as_mut() {
138			None => panic!("Cannot change level of Nil"),
139			Some(Node { level: l, .. }) => mem::replace(l, level)
140		};
141	}
142
143	/// ```none
144	///   L <--- S           S ---> T
145	///  / \      \     =>  /      / \
146	/// A   B      R       A      B   R
147	/// ```
148	fn skew(mut self) -> Self {
149		match self.as_mut() {
150			None => self,
151			Some(Node {
152				level,
153				left_child: l,
154				..
155			}) => {
156				// if level = l.level, remove the B node from L
157				let b_node = match l.as_mut() {
158					Some(Node {
159						level: l_level,
160						right_child: b,
161						..
162					}) if level == l_level => b.take(),
163					_ => return self
164				};
165
166				// add the B node as our left child, removing L
167				let mut l_node = mem::replace(l, b_node);
168
169				// add our node T as the right child of L
170				l_node
171					.as_mut()
172					.unwrap_or_else(|| unreachable!())
173					.right_child = self;
174
175				// L is our new node
176				l_node
177			}
178		}
179	}
180
181	/// ```none
182	///   S --> R --> X          R
183	///  /     /          =>    / \
184	/// A     B                T   X
185	///                       / \
186	///                      A   B
187	/// ```
188	fn split(mut self) -> Self {
189		match self.as_mut() {
190			None => self,
191			Some(Node {
192				level,
193				right_child: r,
194				..
195			}) => {
196				// remove the B node if R and X are not Nil
197				let b_node = match r.as_mut() {
198					Some(Node {
199						left_child: b,
200						right_child: x,
201						..
202					}) if &x.level() == level => b.take(),
203					_ => return self
204				};
205
206				// attach the B node to our node, removing R
207				let mut r_node = mem::replace(r, b_node);
208
209				// attach our node to R and increment its level
210				let r_node_mut = r_node.as_mut().unwrap_or_else(|| unreachable!());
211				r_node_mut.level += 1;
212				r_node_mut.left_child = self;
213
214				// R is our new node
215				r_node
216			}
217		}
218	}
219}
220
221#[cfg(all(test, not(feature = "benchmark")))]
222mod test {
223	use super::*;
224
225	macro_rules! tree {
226		() => {
227			AANode::new()
228		};
229		(Nil) => {
230			AANode::new()
231		};
232		($content:expr) => {
233			AANode::from($content)
234		};
235		($content:expr => [$level:expr, $left:tt, $right:tt]) => {
236			{
237				let _left = tree!(@internal $left);
238				let _right = tree!(@internal $right);
239				AANode(Some(Box::new(Node {
240					level: $level,
241					content: $content,
242					left_child: _left,
243					right_child: _right
244				})))
245			}
246		};
247		(@internal ($content:expr => [$level:expr, $left:tt, $right:tt])) => {
248			tree!($content => [$level, $left, $right])
249		};
250		(@internal $inner:tt) => {
251			tree!($inner)
252		};
253	}
254
255	// ### TEST SKEW ###
256
257	#[test]
258	fn test_skew_nil() {
259		let root: AANode<char> = tree!();
260		println!("Input: {:?}", root);
261		let skewed = root.skew();
262		let expected = tree!();
263		assert_eq!(skewed, expected);
264	}
265
266	#[test]
267	fn test_skew_leaf() {
268		let root = tree!('T');
269		println!("Input: {:?}", root);
270		let skewed = root.skew();
271		let expected = tree!('T');
272		assert_eq!(skewed, expected);
273	}
274
275	#[test]
276	fn test_skew_simple() {
277		let root = tree!('T' => [2, ('L' => [2, Nil, Nil]), 'R']);
278		println!("Input: {:?}", root);
279		let skewed = root.skew();
280		let expected = tree!('L' => [2, Nil, ('T' => [2, Nil, 'R'])]);
281		assert_eq!(skewed, expected);
282	}
283
284	#[test]
285	fn test_skew_full() {
286		let root = tree!('T' => [2, ('L' => [2, 'A', 'B']), 'R']);
287		println!("Input: {:?}", root);
288		let skewed = root.skew();
289		let expected = tree!('L' => [2, 'A', ('T' => [2, 'B', 'R'])]);
290		assert_eq!(skewed, expected);
291	}
292
293	// ### TEST SPLIT ###
294
295	#[test]
296	fn test_split_nil() {
297		let root: AANode<char> = tree!();
298		println!("Input: {:?}", root);
299		let splitted = root.split();
300		let expected = tree!();
301		assert_eq!(splitted, expected);
302	}
303
304	#[test]
305	fn test_split_leaf() {
306		let root = tree!('T');
307		println!("Input: {:?}", root);
308		let splitted = root.split();
309		let expected = tree!('T');
310		assert_eq!(splitted, expected);
311	}
312
313	#[test]
314	fn test_split_good_tree() {
315		let root = tree!('T' => [2, 'A', ('R' => [2, 'B', 'X'])]);
316		println!("Input: {:?}", root);
317		let splitted = root.split();
318		let expected = tree!('T' => [2, 'A', ('R' => [2, 'B', 'X'])]);
319		assert_eq!(splitted, expected);
320	}
321
322	#[test]
323	fn test_split_bad_tree() {
324		let root = tree!('T' => [2, 'A', ('R' => [2, 'B', ('X' => [2, 'Y', 'Z'])])]);
325		println!("Input: {:?}", root);
326		let splitted = root.split();
327		let expected = tree!('R' => [3, ('T' => [2, 'A', 'B']), ('X' => [2, 'Y', 'Z'])]);
328		assert_eq!(splitted, expected);
329	}
330
331	// ### TEST INSERT ###
332
333	#[test]
334	fn test_insert_greater() {
335		let mut root = tree!();
336		for content in ['A', 'B', 'C', 'D', 'E', 'F', 'G'].iter() {
337			assert!(root.insert(*content));
338		}
339		let expected = tree!('D' => [3, ('B' => [2, 'A', 'C']), ('F' => [2, 'E', 'G'])]);
340		assert_eq!(root, expected);
341	}
342
343	#[test]
344	fn test_insert_smaller() {
345		let mut root = tree!();
346		for content in ['Z', 'Y', 'X', 'W', 'V'].iter() {
347			assert!(root.insert(*content));
348		}
349		let expected = tree!('W' => [2, 'V', ('Y' => [2, 'X', 'Z'])]);
350		assert_eq!(root, expected);
351	}
352
353	#[test]
354	fn test_insert_multiple() {
355		let mut root = tree!();
356		for content in ['A', 'A'].iter() {
357			root.insert(*content);
358		}
359		let expected = tree!('A');
360		assert_eq!(root, expected);
361	}
362
363	// ### TEST REMOVE ###
364
365	#[test]
366	fn test_remove_successor() {
367		let mut root = tree!('B' => [1, Nil, 'C']);
368		println!("Input:  `{:?}`", root);
369		let removed = root.remove(&'B');
370		let expected = tree!('C');
371		assert_eq!(removed, Some('B'));
372		assert_eq!(root, expected);
373	}
374
375	#[test]
376	fn test_remove_predecessor() {
377		let mut root = tree!('B' => [2, 'A', 'C']);
378		println!("Input:  `{:?}`", root);
379		let removed = root.remove(&'B');
380		let expected = tree!('A' => [1, Nil, 'C']);
381		assert_eq!(removed, Some('B'));
382		assert_eq!(root, expected);
383	}
384
385	#[test]
386	fn test_remove_complex() {
387		// example taken from https://web.eecs.umich.edu/~sugih/courses/eecs281/f11/lectures/12-AAtrees+Treaps.pdf
388		let mut root = tree!(30 => [3, (15 => [2, 5, 20]), (70 => [3, (50 => [2, 35, (60 => [2, 55, 65])]), (85 => [2, 80, 90])])]);
389		println!("Input:  `{:?}`", root);
390		let removed = root.remove(&5);
391		let expected = tree!(50 => [3, (30 => [2, (15 => [1, Nil, 20]), 35]), (70 => [3, (60 => [2, 55, 65]), (85 => [2, 80, 90])])]);
392		assert_eq!(removed, Some(5));
393		assert_eq!(root, expected);
394	}
395}
aatree/node/mod.rs

aatree/node/
mod.rs
