espalier/
lib.rs

1use std::{fmt::Debug, iter::Iterator, marker::PhantomData};
2
3#[cfg(test)]
4mod tests;
5
6/// An element in the vector. It has the user value, a level starting from 0
7/// for the root element, and a parent. The root's parent is 0.
8pub struct Node<K, V> {
9    /// The stored value
10    pub value: V,
11    /// Index of the parent node.
12    parent: usize,
13    /// The number of descendents, not including this node. This allows
14    /// fast iteration of children.
15    num_descendants: usize,
16    /// This just exists because we didn't use K, but we want it to be part
17    /// of the type.
18    _key_type: PhantomData<K>,
19}
20
21impl<K, V> Node<K, V>
22where
23    usize: Into<K>,
24{
25    /// Get the ID of the parent node. The ID of the root node is equal to
26    /// the root ID (so check for loops!).
27    pub fn parent(&self) -> K {
28        self.parent.into()
29    }
30
31    /// The number of descendents, not including this node.
32    pub fn num_descendants(&self) -> usize {
33        self.num_descendants
34    }
35}
36
37impl<K, V: Debug> Debug for Node<K, V> {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("Node")
40            .field("value", &self.value)
41            .field("parent", &self.parent)
42            .field("num_descendants", &self.num_descendants)
43            .finish()
44    }
45}
46
47impl<K, V: PartialEq> PartialEq for Node<K, V> {
48    fn eq(&self, other: &Self) -> bool {
49        self.value == other.value
50            && self.parent == other.parent
51            && self.num_descendants == other.num_descendants
52    }
53}
54
55impl<K, V: Clone> Clone for Node<K, V> {
56    fn clone(&self) -> Self {
57        Self {
58            value: self.value.clone(),
59            parent: self.parent.clone(),
60            num_descendants: self.num_descendants.clone(),
61            _key_type: PhantomData,
62        }
63    }
64}
65
66impl<K, V: Eq> Eq for Node<K, V> {}
67
68impl<K, V: Copy> Copy for Node<K, V> {}
69
70/// A flattened tree. The nodes are stored in pre-order (depth first order).
71pub struct Tree<K, V> {
72    nodes: Vec<Node<K, V>>,
73    parent_stack: Vec<usize>,
74}
75
76impl<K, V> Default for Tree<K, V> {
77    fn default() -> Self {
78        Self {
79            nodes: Default::default(),
80            parent_stack: Default::default(),
81        }
82    }
83}
84
85impl<K, V> Tree<K, V>
86where
87    usize: Into<K>,
88    K: Into<usize>,
89{
90    /// Create a new empty tree containing no nodes.
91    pub fn new() -> Self {
92        Self::default()
93    }
94
95    /// Create an empty tree with reserved capacity for a certain number of nodes.
96    pub fn with_capacity(capacity: usize) -> Self {
97        Self {
98            nodes: Vec::with_capacity(capacity),
99            parent_stack: Vec::new(),
100        }
101    }
102
103    /// Return the total number of nodes in the tree.
104    pub fn len(&self) -> usize {
105        self.nodes.len()
106    }
107
108    /// Return true if the tree is empty. In this case it has no root node.
109    pub fn is_empty(&self) -> bool {
110        self.nodes.is_empty()
111    }
112
113    /// Push a child of the current node. If this is the first node it will
114    /// become the root. The new node becomes the current node.
115    ///
116    /// # WARNING
117    ///
118    /// If you don't push with the correct values then iteration may give
119    /// unexpected results.
120    pub fn push(&mut self, value: V) -> K {
121        let id = self.len();
122
123        self.nodes.push(Node {
124            value,
125            parent: *self.parent_stack.last().unwrap_or(&id),
126            num_descendants: 0,
127            _key_type: PhantomData,
128        });
129
130        // Increment the descendent counts of all parent nodes.
131        for &parent in self.parent_stack.iter() {
132            self.nodes[parent].num_descendants += 1;
133        }
134
135        self.parent_stack.push(id);
136
137        id.into()
138    }
139
140    /// Set the current node to its parent. It's safe to call this if the
141    /// the tree is empty, in which case nothing will change.
142    ///
143    /// It returns the ID of the new "current" node or None if we have `up()`d
144    /// all the way to the top.
145    ///
146    /// It is ok to call this if the current node is the root node. If you then
147    /// add more nodes you will end up with a tree with multiple root nodes.
148    /// this should work fine but might be confusing!
149    pub fn up(&mut self) -> Option<K> {
150        self.parent_stack.pop();
151        self.parent_stack.last().map(|&id| id.into())
152    }
153
154    /// Get a reference to a node. Returns `None` for invalid IDs.
155    pub fn get(&self, id: K) -> Option<&Node<K, V>> {
156        self.nodes.get(id.into())
157    }
158
159    /// Get a mutable reference to a node. Returns `None` for invalid IDs.
160    pub fn get_mut(&mut self, id: K) -> Option<&mut Node<K, V>> {
161        self.nodes.get_mut(id.into())
162    }
163
164    /// Get a reference to the first node (or `None` if the tree is empty).
165    /// This will normally be the tree's only root node but it is possible
166    /// to have trees with multiple roots.
167    pub fn first(&self) -> Option<&Node<K, V>> {
168        self.nodes.first()
169    }
170
171    /// Get a mutable reference to the first node (or `None` if the tree is empty).
172    /// This will normally be the tree's only root node but it is possible
173    /// to have trees with multiple roots.
174    pub fn first_mut(&mut self) -> Option<&mut Node<K, V>> {
175        self.nodes.first_mut()
176    }
177
178    /// Get a reference to the last node (or `None` if the tree is empty).
179    /// This is the "current" node. If you call push() it will add a child
180    /// to this node.
181    pub fn last(&self) -> Option<&Node<K, V>> {
182        self.nodes.last()
183    }
184
185    /// Get a mutable reference to the last node (or `None` if the tree is empty).
186    /// This is the "current" node. If you call push() it will add a child
187    /// to this node.
188    pub fn last_mut(&mut self) -> Option<&mut Node<K, V>> {
189        self.nodes.last_mut()
190    }
191
192    /// Iterate through all the tree nodes in the order they were added (which
193    /// must be pre-order / depth first).
194    pub fn iter(&self) -> impl Iterator<Item = &Node<K, V>> {
195        self.nodes.iter()
196    }
197
198    /// Convert the tree into an iterator through all the tree nodes in the
199    /// order they were added (which must be pre-order / depth first).
200    pub fn into_iter(self) -> impl Iterator<Item = Node<K, V>> {
201        self.nodes.into_iter()
202    }
203
204    /// Get a slice of all nodes in the tree in the order they were added
205    /// (which must be pre-order / depth-first).
206    pub fn all(&self) -> &[Node<K, V>] {
207        self.nodes.as_slice()
208    }
209
210    /// Get a slice of all the descendents of a node.
211    pub fn descendents(&self, id: K) -> &[Node<K, V>] {
212        let id = id.into();
213        let num_descendants = self
214            .nodes
215            .get(id)
216            .map(|node| node.num_descendants)
217            .unwrap_or_default();
218        &self.nodes[id + 1..id + 1 + num_descendants]
219    }
220
221    /// Get an iterator over the parents of a node (not including the node itself).
222    pub fn parents(&self, id: K) -> ParentIter<'_, K, V> {
223        let id = id.into();
224        ParentIter { id, tree: self }
225    }
226
227    /// Get an iterator over the immediate children of a node.
228    pub fn children(&self, id: K) -> ChildrenIter<'_, K, V> {
229        let id = id.into();
230        ChildrenIter {
231            current_id: id + 1,
232            max_id: id
233                + self
234                    .nodes
235                    .get(id)
236                    .map(|node| node.num_descendants)
237                    .unwrap_or_default(),
238            tree: self,
239        }
240    }
241}
242
243impl<K, V: Debug> Debug for Tree<K, V> {
244    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245        f.debug_struct("Tree")
246            .field("nodes", &self.nodes)
247            .field("parent_stack", &self.parent_stack)
248            .finish()
249    }
250}
251
252impl<K, V: PartialEq> PartialEq for Tree<K, V> {
253    fn eq(&self, other: &Self) -> bool {
254        self.nodes == other.nodes && self.parent_stack == other.parent_stack
255    }
256}
257
258impl<K, V: Clone> Clone for Tree<K, V> {
259    fn clone(&self) -> Self {
260        Self {
261            nodes: self.nodes.clone(),
262            parent_stack: self.parent_stack.clone(),
263        }
264    }
265}
266
267impl<K, V: Eq> Eq for Tree<K, V> {}
268
269pub struct ParentIter<'a, K, V> {
270    id: usize,
271    tree: &'a Tree<K, V>,
272}
273
274impl<'a, K, V> Iterator for ParentIter<'a, K, V>
275where
276    K: Into<usize>,
277    usize: Into<K>,
278{
279    type Item = (K, &'a Node<K, V>);
280
281    fn next(&mut self) -> Option<Self::Item> {
282        self.tree.nodes.get(self.id).and_then(|node| {
283            if node.parent == self.id {
284                None
285            } else {
286                self.id = node.parent;
287                self.tree
288                    .nodes
289                    .get(self.id)
290                    .map(|node| (self.id.into(), node))
291            }
292        })
293    }
294}
295
296pub struct ChildrenIter<'a, K, V> {
297    current_id: usize,
298    max_id: usize,
299    tree: &'a Tree<K, V>,
300}
301
302impl<'a, K, V> Iterator for ChildrenIter<'a, K, V>
303where
304    usize: Into<K>,
305{
306    type Item = (K, &'a Node<K, V>);
307
308    fn next(&mut self) -> Option<Self::Item> {
309        if self.current_id <= self.max_id {
310            let node = self.tree.nodes.get(self.current_id);
311            let id_and_node = node.map(|node| (self.current_id.into(), node));
312            self.current_id += self
313                .tree
314                .nodes
315                .get(self.current_id)
316                .map(|node| node.num_descendants)
317                .unwrap_or_default()
318                + 1;
319            id_and_node
320        } else {
321            None
322        }
323    }
324}
espalier/lib.rs

espalier/
lib.rs
