1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
use std::collections::HashMap;
use std::ops::Deref;
use std::ops::DerefMut;
use std::usize;
use colored::Colorize;
use crate::node::TreeNode;
use crate::TreeNodeRef;
pub mod leaf;
#[derive(Debug, Hash, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct NodePosition {
// Vertical depth
pub depth: usize,
// Horizontal index at each depth.
pub index: usize,
// Index of the child relative to the parent
pub child_index: usize,
}
impl std::fmt::Display for NodePosition {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}: {} {}: {} {}: {}",
"depth".bright_cyan(),
self.depth().to_string().bright_purple(),
"index".bright_cyan(),
self.index().to_string().bright_purple(),
"child_index".bright_cyan(),
self.child_index().to_string().bright_purple(),
)
}
}
impl NodePosition {
pub fn zero() -> Self {
NodePosition {
depth: 0,
index: 0,
child_index: 0,
}
}
// Return a NodePosition with all positions set to usize::MAX
pub fn max() -> Self {
NodePosition {
depth: usize::MAX,
index: usize::MAX,
child_index: usize::MAX,
}
}
/// Get the depth of this node
pub fn depth(&self) -> usize {
self.depth
}
/// Get the horizontal position of this node relative to all siblings at this depth
pub fn index(&self) -> usize {
self.index
}
/// Get the child position relative to the parent
pub fn child_index(&self) -> usize {
self.child_index
}
}
pub struct IterNode<R>
where
R: TreeNodeRef,
{
position: NodePosition,
node: R,
}
impl<R> IterNode<R>
where
R: TreeNodeRef,
{
/// The index along the horizontal at the current depth
pub fn index(&self) -> usize {
self.position.index
}
/// The vertical depth of this node
pub fn depth(&self) -> usize {
self.position.depth
}
pub fn child_index(&self) -> usize {
self.position.child_index
}
pub fn position(&self) -> &NodePosition {
&self.position
}
}
impl<R> Deref for IterNode<R>
where
R: TreeNodeRef,
{
type Target = R;
fn deref(&self) -> &Self::Target {
&self.node
}
}
impl<R> DerefMut for IterNode<R>
where
R: TreeNodeRef,
{
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.node
}
}
pub struct NodeRefIter<R>
where
R: TreeNodeRef,
{
stack: Vec<(usize, usize, usize, R)>,
index: HashMap<usize, usize>,
}
impl<R> NodeRefIter<R>
where
R: TreeNodeRef,
{
pub fn new(node: R) -> Self {
Self {
stack: Vec::from([(0, 0, 0, node)]),
index: HashMap::new(),
}
}
}
impl<R> Iterator for NodeRefIter<R>
where
R: TreeNodeRef,
{
type Item = IterNode<R>;
fn next(&mut self) -> Option<Self::Item> {
let current = self.stack.pop();
current.map(|(child_index, index, depth, node)| {
node.node().children().map(|children| {
let index = self.index.entry(depth).or_insert(0);
// Increment the horizontal index in the iterator state by the number of children we have.
// This increases the index by how many children we're about to push onto the stack,
// and will be retained as the iterator moves between different depths
*index += children.len();
children
.iter()
// Enumerate to get the child index
.enumerate()
// Reverse the iterator as we're pushing onto a stack
// that will pop the nodes back off in reverse order
.rev()
.for_each(|(child_index, child)| {
self.stack.push((
child_index,
// *index is positive offset by the number of children we're adding to the Vec,
// so we need to decrement the next index by 1 as we're iterating backwards
// The child index is also decreasing, so we can subtract the difference from the
// total nodes to get an index offset
*index - (children.len() - child_index),
// Each child being pushed will have it's depth set to the current node depth + 1
depth + 1,
(*child).clone(),
));
})
});
IterNode {
position: NodePosition {
depth,
index,
child_index,
},
node,
}
})
}
}