Skip to main content

treelike/example/
lintree.rs

1use crate::Treelike;
2
3/// A tree whose nodes are stored in a backing slice.
4///
5/// The root node is at index 0 and the child
6/// nodes at `index*2 + 1` and `index*2 + 2`.
7///
8/// Used in most examples, as it is easy to initialize.
9///
10/// Also shows a case where Treelike is implemented on the node itself, and not a reference to one.
11/// An important pitfall for that is that you may need to manually implement [Copy] and [Clone] as
12/// deriving them places [Copy]/[Clone] bounds on all type parameters (`T` in this case), even
13/// though that might not be necessary due to the content not being stored in-line, and therefore
14/// not being [Copy]/[Clone]d.
15#[derive(Debug)]
16pub struct LinTree<'a, T> {
17	index: usize,
18	slice: &'a [T],
19}
20
21impl<'a, T> LinTree<'a, T> {
22	pub fn new(index: usize, slice: &'a [T]) -> Self { LinTree { index, slice } }
23
24	fn tuple_new((index, slice): (usize, &'a [T])) -> Option<Self> {
25		slice.get(index).map(|_| Self::new(index, slice))
26	}
27}
28
29impl<'a, T> Copy for LinTree<'a, T> {}
30
31impl<'a, T> Clone for LinTree<'a, T> {
32	fn clone(&self) -> Self {
33		Self {
34			index: self.index,
35			slice: self.slice,
36		}
37	}
38}
39
40impl<'a, T: core::fmt::Debug> Treelike for LinTree<'a, T> {
41	type Content = &'a T;
42
43	type ChildIterator = core::iter::FlatMap<
44		core::iter::Zip<
45			core::iter::Chain<core::iter::Once<usize>, core::iter::Once<usize>>,
46			core::iter::Repeat<&'a [T]>,
47		>,
48		Option<LinTree<'a, T>>,
49		fn((usize, &'a [T])) -> Option<LinTree<'a, T>>,
50	>;
51
52	fn content(self) -> Self::Content { &self.slice[self.index] }
53
54	fn children(self) -> Self::ChildIterator {
55		use core::iter::{once, repeat};
56		let left = 2 * self.index + 1;
57		let right = 2 * self.index + 2;
58		once(left)
59			.chain(once(right))
60			.zip(repeat(self.slice))
61			.flat_map(Self::tuple_new)
62	}
63
64	/// This is also an example of overriding the [Treelike]s default implementations where
65	/// necessary. LinTree can provide breadth-first traversal with a simple iteration
66	fn callback_bft<CB: FnMut(Self::Content, usize)>(self, mut callback: CB) {
67		let usize_bits = (core::mem::size_of::<usize>() * 8) as u32;
68		for (i, content) in self.slice.iter().enumerate().skip(self.index) {
69			// this is flooring log2 for integers
70			// find the first one by subtracting the bit-length from the leading_zeros
71			// first one - 1 is already floored log2
72			let depth = usize_bits - (i + 1).leading_zeros() - 1;
73			callback(content, depth as usize);
74		}
75	}
76
77	//TODO could also implement filtered bft, but that requires itertools group_by. put that in
78	//as an optional dependency maybe
79}
80
81#[cfg(feature = "test")]
82#[test]
83fn depth_test() {
84	extern crate alloc;
85	use alloc::vec::Vec;
86	let base = [0, (1), 2, (3), 4, 5, 6, (7), 8, 9, 10, 11, 12, 13, 14, (15)];
87
88	let root = LinTree::new(0, &base);
89	let mut state = Vec::new();
90	root.callback_bft(|_content, depth| state.push(depth));
91	assert_eq!(&state, &[0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4]);
92
93	let deeper = LinTree::new(4, &base);
94	let mut state = Vec::new();
95	deeper.callback_bft(|_content, depth| state.push(depth));
96	assert_eq!(&state, &[2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4]);
97}
98
99#[test]
100fn basic_test() {
101	extern crate alloc;
102	use alloc::vec::Vec;
103	let base = [3, 4, 5, 6, 7];
104	let root = LinTree::new(0, &base);
105
106	let mut state = Vec::new();
107	root.callback_dft(|val, _depth| state.push(*val), ());
108	assert_eq!(alloc::vec![6, 7, 4, 5, 3], state);
109}
110
111#[cfg(feature = "alloc")]
112#[test]
113fn iter_test() {
114	use alloc::vec::Vec;
115	let base = [0, (1), 2, (3), 4, 5, 6, (7), 8, 9, 10, 11, 12, 13, 14, (15)];
116	let root = LinTree::new(0, &base);
117
118	let mut state = Vec::new();
119	root.callback_dft(|val, _depth| state.push(*val), ());
120
121	let iter_state: Vec<_> = root.iter_dft(()).cloned().collect();
122	assert_eq!(iter_state, state);
123}
124
125#[cfg(feature = "alloc")]
126#[test]
127fn dft_pre_test() {
128	use alloc::vec::Vec;
129	let base = [0, (1), 2, (3), 4, 5, 6, (7), 8, 9, 10, 11, 12, 13, 14, (15)];
130	let root = LinTree::new(0, &base);
131
132	let mut state = Vec::new();
133	root.callback_dft_pre(|val, _depth| state.push(*val), ());
134
135	let iter_state: Vec<_> = root.iter_dft_pre(()).cloned().collect();
136	assert_eq!(iter_state, state);
137}
138
139#[cfg(feature = "alloc")]
140#[test]
141fn bft_pre_test() {
142	use alloc::vec::Vec;
143	let base = [0, (1), 2, (3), 4, 5, 6, (7), 8, 9, 10, 11, 12, 13, 14, (15)];
144	let root = LinTree::new(0, &base);
145
146	let mut state = Vec::new();
147	root.callback_bft(|val, _depth| state.push(*val));
148
149	let iter_state: Vec<_> = root.iter_bft(()).cloned().collect();
150	assert_eq!(iter_state, state);
151}