vec_tree/
lib.rs

1/*!
2A safe tree using an arena allocator that allows deletion without suffering from
3[the ABA problem](https://en.wikipedia.org/wiki/ABA_problem) by using generational
4indices.
5
6It uses [generational-arena](https://github.com/fitzgen/generational-arena) under
7the hood, made by [fitzgen](https://github.com/fitzgen), special thanks to him.
8
9[generational-arena](https://github.com/fitzgen/generational-arena) is itself inspired
10by [Catherine West's closing keynote at RustConf
112018](http://rustconf.com/program.html#closingkeynote), where these ideas
12were presented in the context of an Entity-Component-System for games
13programming.
14
15## What? Why?
16
17When you are working with a tree and you want to add and delete individual
18nodes at a time, or you are writing a game and its world consists of many
19inter-referencing objects with dynamic lifetimes that depend on user
20input. These are situations where matching Rust's ownership and lifetime rules
21can get tricky.
22
23It doesn't make sense to use shared ownership with interior mutability (ie
24`Rc<RefCell<T>>` or `Arc<Mutex<T>>`) nor borrowed references (ie `&'a T` or `&'a
25mut T`) for structures. The cycles rule out reference counted types, and the
26required shared mutability rules out borrows. Furthermore, lifetimes are dynamic
27and don't follow the borrowed-data-outlives-the-borrower discipline.
28
29In these situations, it is tempting to store objects in a `Vec<T>` and have them
30reference each other via their indices. No more borrow checker or ownership
31problems! Often, this solution is good enough.
32
33However, now we can't delete individual items from that `Vec<T>` when we no
34longer need them, because we end up either
35
36* messing up the indices of every element that follows the deleted one, or
37
38* suffering from the [ABA
39  problem](https://en.wikipedia.org/wiki/ABA_problem). To elaborate further, if
40  we tried to replace the `Vec<T>` with a `Vec<Option<T>>`, and delete an
41  element by setting it to `None`, then we create the possibility for this buggy
42  sequence:
43
44    * `obj1` references `obj2` at index `i`
45
46    * someone else deletes `obj2` from index `i`, setting that element to `None`
47
48    * a third thing allocates `obj3`, which ends up at index `i`, because the
49      element at that index is `None` and therefore available for allocation
50
51    * `obj1` attempts to get `obj2` at index `i`, but incorrectly is given
52      `obj3`, when instead the get should fail.
53
54By introducing a monotonically increasing generation counter to the collection,
55associating each element in the collection with the generation when it was
56inserted, and getting elements from the collection with the *pair* of index and
57the generation at the time when the element was inserted, then we can solve the
58aforementioned ABA problem. When indexing into the collection, if the index
59pair's generation does not match the generation of the element at that index,
60then the operation fails.
61
62## Features
63
64* Zero `unsafe`
65* There is different iterators to traverse the tree
66* Well tested
67
68## Usage
69
70First, add `vec-tree` to your `Cargo.toml`:
71
72```toml
73[dependencies]
74vec-tree = "0.1"
75```
76
77Then, import the crate and use the `vec-tree::Tree`
78
79```rust
80extern crate vec_tree;
81use vec_tree::VecTree;
82
83let mut tree = VecTree::new();
84
85// Insert some elements into the tree.
86let root_node = tree.insert_root(1);
87let child_node_1 = tree.insert(10, root_node);
88let child_node_2 = tree.insert(11, root_node);
89let child_node_3 = tree.insert(12, root_node);
90let grandchild = tree.insert(100, child_node_3);
91
92// Inserted elements can be accessed infallibly via indexing (and missing
93// entries will panic).
94assert_eq!(tree[child_node_1], 10);
95
96// Alternatively, the `get` and `get_mut` methods provide fallible lookup.
97if let Some(node_value) = tree.get(child_node_2) {
98    println!("The node value is: {}", node_value);
99}
100if let Some(node_value) = tree.get_mut(grandchild) {
101    *node_value = 101;
102}
103
104// We can remove elements.
105tree.remove(child_node_3);
106
107// Insert a new one.
108let child_node_4 = tree.insert(13, root_node);
109
110// The tree does not contain `child_node_3` anymore, but it does contain
111// `child_node_4`, even though they are almost certainly at the same index
112// within the arena of the tree in practice. Ambiguities are resolved with
113// an associated generation tag.
114assert!(!tree.contains(child_node_3));
115assert!(tree.contains(child_node_4));
116
117// We can also move a node (and its descendants).
118tree.append_child(child_node_1, child_node_4);
119
120// Iterate over the children of a node.
121for value in tree.children(child_node_1) {
122    println!("value: {:?}", value);
123}
124
125// Or all the descendants in depth first search order.
126let descendants = tree
127    .descendants(root_node)
128    .map(|node| tree[node])
129    .collect::<Vec<i32>>();
130
131assert_eq!(descendants, [1, 10, 13, 11]);
132```
133 */
134
135#![forbid(unsafe_code)]
136
137extern crate generational_arena;
138use generational_arena::Arena;
139pub use generational_arena::Index;
140
141use core::ops;
142use std::{fmt, mem};
143
144/// The `VecTree` allows inserting and removing elements that are referred to by
145/// `Index`.
146///
147/// [See the module-level documentation for example usage and motivation.](./index.html)
148#[derive(Clone, Debug)]
149pub struct VecTree<T> {
150    nodes: Arena<Node<T>>,
151    root_index: Option<Index>,
152}
153
154#[derive(Clone, Debug)]
155struct Node<T> {
156    parent: Option<Index>,
157    previous_sibling: Option<Index>,
158    next_sibling: Option<Index>,
159    first_child: Option<Index>,
160    last_child: Option<Index>,
161    data: T,
162}
163
164const DEFAULT_CAPACITY: usize = 4;
165
166impl<T> Default for VecTree<T> {
167    fn default() -> Self {
168        VecTree::with_capacity(DEFAULT_CAPACITY)
169    }
170}
171
172impl<T> VecTree<T> {
173    /// Constructs a new, empty `VecTree`.
174    ///
175    /// # Examples
176    ///
177    /// ```
178    /// use vec_tree::VecTree;
179    ///
180    /// let mut tree = VecTree::<usize>::new();
181    /// # let _ = tree;
182    /// ```
183    pub fn new() -> VecTree<T> {
184        VecTree::with_capacity(DEFAULT_CAPACITY)
185    }
186
187    /// Constructs a new, empty `VecTree<T>` with the specified capacity.
188    ///
189    /// The `VecTree<T>` will be able to hold `n` elements without further allocation.
190    ///
191    /// # Examples
192    ///
193    /// ```
194    /// use vec_tree::VecTree;
195    ///
196    /// let mut tree = VecTree::with_capacity(10);
197    /// let root = tree.try_insert_root(0).unwrap();
198    ///
199    /// // These insertions will not require further allocation.
200    /// for i in 1..10 {
201    ///     assert!(tree.try_insert(i, root).is_ok());
202    /// }
203    ///
204    /// // But now we are at capacity, and there is no more room.
205    /// assert!(tree.try_insert(99, root).is_err());
206    /// ```
207    pub fn with_capacity(n: usize) -> VecTree<T> {
208        VecTree {
209            nodes: Arena::with_capacity(n),
210            root_index: None,
211        }
212    }
213
214
215    /// Allocate space for `additional_capacity` more elements in the tree.
216    ///
217    /// # Panics
218    ///
219    /// Panics if this causes the capacity to overflow.
220    ///
221    /// # Examples
222    ///
223    /// ```
224    /// use vec_tree::VecTree;
225    ///
226    /// let mut tree = VecTree::with_capacity(10);
227    /// tree.reserve(5);
228    /// assert_eq!(tree.capacity(), 15);
229    /// # let _: VecTree<usize> = tree;
230    /// ```
231    #[inline]
232    pub fn reserve(&mut self, additional_capacity: usize) {
233        self.nodes.reserve(additional_capacity);
234    }
235
236    /// Attempts to insert `data` into the tree using existing capacity.
237    ///
238    /// This method will never allocate new capacity in the tree.
239    ///
240    /// If insertion succeeds, then the `data`'s index is returned. If
241    /// insertion fails, then `Err(data)` is returned to give ownership of
242    /// `data` back to the caller.
243    ///
244    /// # Examples
245    ///
246    /// ```
247    /// use vec_tree::VecTree;
248    ///
249    /// let mut tree = VecTree::new();
250    /// let root = tree.insert_root(0);
251    ///
252    /// match tree.try_insert(42, root) {
253    ///     Ok(idx) => {
254    ///         // Insertion succeeded.
255    ///         assert_eq!(tree[idx], 42);
256    ///     }
257    ///     Err(x) => {
258    ///         // Insertion failed.
259    ///         assert_eq!(x, 42);
260    ///     }
261    /// };
262    /// ```
263    #[inline]
264    pub fn try_insert(&mut self, data: T, parent_id: Index) -> Result<Index, T> {
265        let node_result = self.try_create_node(data);
266
267        match node_result {
268            Ok(node) => {
269                self.append_child(parent_id, node);
270                node_result
271            }
272            Err(err) => Err(err),
273        }
274    }
275
276    /// Insert `data` into the tree, allocating more capacity if necessary.
277    ///
278    /// The `data`'s associated index in the tree is returned.
279    ///
280    /// # Examples
281    ///
282    /// ```
283    /// use vec_tree::VecTree;
284    ///
285    /// let mut tree = VecTree::with_capacity(1);
286    /// assert_eq!(tree.capacity(), 1);
287    ///
288    /// let root = tree.insert_root(0);
289    ///
290    /// let idx = tree.insert(42, root);
291    /// assert_eq!(tree[idx], 42);
292    /// assert_eq!(tree.capacity(), 2);
293    /// ```
294    #[inline]
295    pub fn insert(&mut self, data: T, parent_id: Index) -> Index {
296        let node = self.create_node(data);
297
298        self.append_child(parent_id, node);
299
300        node
301    }
302
303    /// Attempts to insert `data` into the tree as root node using existing
304    /// capacity.
305    ///
306    /// This method will never allocate new capacity in the tree.
307    ///
308    /// If insertion succeeds, then the `data`'s index is returned. If
309    /// insertion fails, then `Err(data)` is returned to give ownership of
310    /// `data` back to the caller.
311    ///
312    /// # Examples
313    ///
314    /// ```
315    /// use vec_tree::VecTree;
316    ///
317    /// let mut tree = VecTree::new();
318    ///
319    /// match tree.try_insert_root(42) {
320    ///     Ok(idx) => {
321    ///         // Insertion succeeded.
322    ///         assert_eq!(tree[idx], 42);
323    ///     }
324    ///     Err(x) => {
325    ///         // Insertion failed.
326    ///         assert_eq!(x, 42);
327    ///     }
328    /// };
329    /// ```
330    #[inline]
331    pub fn try_insert_root(&mut self, data: T) -> Result<Index, T> {
332        if self.root_index.is_some() {
333            panic!("A root node already exists");
334        }
335
336        match self.try_create_node(data) {
337            Ok(node_id) => {
338                self.root_index = Some(node_id);
339                Ok(node_id)
340            }
341            Err(error) => Err(error),
342        }
343    }
344
345    /// Insert `data` into the tree as a root node, allocating more
346    /// capacity if necessary.
347    ///
348    /// The `data`'s associated index in the tree is returned.
349    ///
350    /// # Examples
351    ///
352    /// ```
353    /// use vec_tree::VecTree;
354    ///
355    /// let mut tree = VecTree::with_capacity(1);
356    /// assert_eq!(tree.capacity(), 1);
357    ///
358    /// let root = tree.insert_root(42);
359    ///
360    /// assert_eq!(tree[root], 42);
361    /// ```
362    #[inline]
363    pub fn insert_root(&mut self, data: T) -> Index {
364        if self.root_index.is_some() {
365            panic!("A root node already exists");
366        }
367
368        let node_id = self.create_node(data);
369        self.root_index = Some(node_id);
370        node_id
371    }
372
373    #[inline]
374    fn try_create_node(&mut self, data: T) -> Result<Index, T> {
375        let new_node = Node {
376            parent: None,
377            first_child: None,
378            last_child: None,
379            previous_sibling: None,
380            next_sibling: None,
381            data,
382        };
383
384        match self.nodes.try_insert(new_node) {
385            Ok(index) => Ok(index),
386            Err(Node { data, .. }) => Err(data),
387        }
388    }
389
390    #[inline]
391    fn create_node(&mut self, data: T) -> Index {
392        let new_node = Node {
393            parent: None,
394            first_child: None,
395            last_child: None,
396            previous_sibling: None,
397            next_sibling: None,
398            data,
399        };
400
401        self.nodes.insert(new_node)
402    }
403
404    /// Remove the element at index `node_id` from the tree.
405    ///
406    /// If the element at index `node_id` is still in the tree, then it is
407    /// returned. If it is not in the tree, then `None` is returned.
408    ///
409    /// # Examples
410    ///
411    /// ```
412    /// use vec_tree::VecTree;
413    ///
414    /// let mut tree = VecTree::new();
415    /// let root = tree.insert_root(42);
416    ///
417    /// assert_eq!(tree.remove(root), Some(42));
418    /// assert_eq!(tree.remove(root), None);
419    /// ```
420    pub fn remove(&mut self, node_id: Index) -> Option<T> {
421        if !self.contains(node_id) {
422            return None;
423        }
424
425        let descendants = self.descendants(node_id).skip(1).collect::<Vec<Index>>();
426        let node = self.nodes.remove(node_id).unwrap();
427
428        let previous_sibling_opt = node.previous_sibling;
429        let next_sibling_opt = node.next_sibling;
430
431        if let Some(previous_sibling_idx) = previous_sibling_opt {
432            if let Some(next_sibling_idx) = next_sibling_opt {
433                // If has previous and next.
434                let (previous_sibling, next_sibling) =
435                    self.nodes.get2_mut(previous_sibling_idx, next_sibling_idx);
436
437                previous_sibling.unwrap().next_sibling = Some(next_sibling_idx);
438                next_sibling.unwrap().previous_sibling = Some(previous_sibling_idx);
439            } else if let Some(parent_idx) = node.parent {
440                // If has previous but no next.
441                let previous_sibling = &mut self.nodes[previous_sibling_idx];
442                previous_sibling.next_sibling = None;
443
444                let parent = &mut self.nodes[parent_idx];
445                parent.last_child = Some(previous_sibling_idx);
446            }
447        } else if let Some(next_sibling_idx) = next_sibling_opt {
448            // If has next but no previous.
449            let next_sibling = &mut self.nodes[next_sibling_idx];
450            next_sibling.previous_sibling = None;
451
452            if let Some(parent_idx) = node.parent {
453                let parent = &mut self.nodes[parent_idx];
454                parent.first_child = Some(next_sibling_idx);
455            }
456        } else if let Some(parent_idx) = node.parent {
457            // If it has no previous and no next.
458            let parent = &mut self.nodes[parent_idx];
459            parent.first_child = None;
460            parent.last_child = None;
461        }
462
463        // Remove descendants from arena.
464        for node_id in descendants {
465            self.nodes.remove(node_id);
466        }
467
468        // Set root_index to None if needed
469        if let Some(root_index) = self.root_index {
470            if root_index == node_id {
471                self.root_index = None;
472            }
473        }
474
475        Some(node.data)
476    }
477
478    /// Is the element at index `node_id` in the tree?
479    ///
480    /// Returns `true` if the element at `node_id` is in the tree, `false` otherwise.
481    ///
482    /// # Examples
483    ///
484    /// ```
485    /// use vec_tree::VecTree;
486    ///
487    /// let mut tree = VecTree::new();
488    /// let root = tree.insert_root(0);
489    ///
490    /// assert!(tree.contains(root));
491    /// tree.remove(root);
492    /// assert!(!tree.contains(root));
493    /// ```
494    pub fn contains(&self, node_id: Index) -> bool {
495        self.nodes.get(node_id).is_some()
496    }
497
498    #[inline]
499    pub fn append_child(&mut self, node_id: Index, new_child_id: Index) {
500        self.detach(new_child_id);
501
502        let last_child_opt;
503        {
504            let (node_opt, new_child_node_opt) = self.nodes.get2_mut(node_id, new_child_id);
505
506            if node_opt.is_none() {
507                panic!("The node you are trying to append to is invalid");
508            }
509
510            if new_child_node_opt.is_none() {
511                panic!("The node you are trying to append is invalid");
512            }
513
514            let node = node_opt.unwrap();
515            let new_child_node = new_child_node_opt.unwrap();
516
517            new_child_node.parent = Some(node_id);
518
519            last_child_opt = mem::replace(&mut node.last_child, Some(new_child_id));
520            if let Some(last_child) = last_child_opt {
521                new_child_node.previous_sibling = Some(last_child);
522            } else {
523                debug_assert!(node.first_child.is_none());
524                node.first_child = Some(new_child_id);
525            }
526        }
527
528        if let Some(last_child) = last_child_opt {
529            debug_assert!(self.nodes[last_child].next_sibling.is_none());
530            self.nodes[last_child].next_sibling = Some(new_child_id);
531        }
532    }
533
534    #[inline]
535    fn detach(&mut self, node_id: Index) {
536        let (parent, previous_sibling, next_sibling) = {
537            let node = &mut self.nodes[node_id];
538            (
539                node.parent.take(),
540                node.previous_sibling.take(),
541                node.next_sibling.take(),
542            )
543        };
544
545        if let Some(next_sibling) = next_sibling {
546            self.nodes[next_sibling].previous_sibling = previous_sibling;
547        } else if let Some(parent) = parent {
548            self.nodes[parent].last_child = previous_sibling;
549        }
550
551        if let Some(previous_sibling) = previous_sibling {
552            self.nodes[previous_sibling].next_sibling = next_sibling;
553        } else if let Some(parent) = parent {
554            self.nodes[parent].first_child = next_sibling;
555        }
556    }
557
558    /// Get a shared reference to the element at index `node_id` if it is in the
559    /// tree.
560    ///
561    /// If the element at index `node_id` is not in the tree, then `None` is returned.
562    ///
563    /// # Examples
564    ///
565    /// ```
566    /// use vec_tree::VecTree;
567    ///
568    /// let mut tree = VecTree::new();
569    /// let root = tree.insert_root(42);
570    ///
571    /// assert_eq!(tree.get(root), Some(&42));
572    /// tree.remove(root);
573    /// assert!(tree.get(root).is_none());
574    /// ```
575    pub fn get(&self, node_id: Index) -> Option<&T> {
576        match self.nodes.get(node_id) {
577            Some(Node { ref data, .. }) => Some(data),
578            _ => None,
579        }
580    }
581
582    /// Get an exclusive reference to the element at index `node_id` if it is in the
583    /// tree.
584    ///
585    /// If the element at index `node_id` is not in the tree, then `None` is returned.
586    ///
587    /// # Examples
588    ///
589    /// ```
590    /// use vec_tree::VecTree;
591    ///
592    /// let mut tree = VecTree::new();
593    /// let root = tree.insert_root(42);
594    ///
595    /// *tree.get_mut(root).unwrap() += 1;
596    /// assert_eq!(tree.remove(root), Some(43));
597    /// assert!(tree.get_mut(root).is_none());
598    /// ```
599    pub fn get_mut(&mut self, node_id: Index) -> Option<&mut T> {
600        match self.nodes.get_mut(node_id) {
601            Some(Node { ref mut data, .. }) => Some(data),
602            _ => None,
603        }
604    }
605    /// Get the root node index from the tree.
606    ///
607    /// If no root node is created in the tree, None is returned.
608    ///
609    /// # Examples
610    ///
611    /// ```
612    /// use vec_tree::VecTree;
613    ///
614    /// let mut tree = VecTree::new();
615    /// assert_eq!(tree.get_root_index(), None);
616    ///
617    /// tree.insert_root(42);
618    /// let root = tree.get_root_index().unwrap();
619    /// assert_eq!(tree[root], 42);
620    /// ```
621    pub fn get_root_index(&self) -> Option<Index> {
622        self.root_index
623    }
624
625    /// Get the capacity of this tree.
626    ///
627    /// The capacity is the maximum number of elements the tree can hold
628    /// without further allocation, including however many it currently
629    /// contains.
630    ///
631    /// # Examples
632    ///
633    /// ```
634    /// use vec_tree::VecTree;
635    ///
636    /// let mut tree = VecTree::with_capacity(10);
637    /// let root = tree.insert_root(0);
638    ///
639    /// // `try_insert` does not allocate new capacity.
640    /// for i in 1..10 {
641    ///     assert!(tree.try_insert(i, root).is_ok());
642    ///     assert_eq!(tree.capacity(), 10);
643    /// }
644    ///
645    /// // But `insert` will if the root is already at capacity.
646    /// tree.insert(11, root);
647    /// assert!(tree.capacity() > 10);
648    /// ```
649    pub fn capacity(&self) -> usize {
650        self.nodes.capacity()
651    }
652
653    /// Clear all the items inside the tree, but keep its allocation.
654    ///
655    /// # Examples
656    ///
657    /// ```
658    /// use vec_tree::VecTree;
659    ///
660    /// let mut tree = VecTree::with_capacity(1);
661    /// let root = tree.insert_root(42);
662    /// tree.insert(43, root); // The capacity is doubled when reached.
663    ///
664    /// tree.clear();
665    /// assert_eq!(tree.capacity(), 2);
666    /// ```
667    pub fn clear(&mut self) {
668        self.nodes.clear();
669        self.root_index = None;
670    }
671
672    /// Return an iterator of references to this node’s parent.
673    pub fn parent(&self, node_id: Index) -> Option<Index> {
674        match self.nodes.get(node_id) {
675            Some(node) => node.parent,
676            _ => None,
677        }
678    }
679
680    /// Return an iterator of references to this node’s children.
681    pub fn children(&self, node_id: Index) -> ChildrenIter<T> {
682        ChildrenIter {
683            tree: self,
684            node_id: self.nodes[node_id].first_child,
685        }
686    }
687
688    /// Return an iterator of references to this node and the siblings before it.
689    ///
690    /// Call `.next().unwrap()` once on the iterator to skip the node itself.
691    pub fn preceding_siblings(&self, node_id: Index) -> PrecedingSiblingsIter<T> {
692        PrecedingSiblingsIter {
693            tree: self,
694            node_id: Some(node_id),
695        }
696    }
697
698    /// Return an iterator of references to this node and the siblings after it.
699    ///
700    /// Call `.next().unwrap()` once on the iterator to skip the node itself.
701    pub fn following_siblings(&self, node_id: Index) -> FollowingSiblingsIter<T> {
702        FollowingSiblingsIter {
703            tree: self,
704            node_id: Some(node_id),
705        }
706    }
707
708    /// Return an iterator of references to this node and its ancestors.
709    ///
710    /// Call `.next().unwrap()` once on the iterator to skip the node itself.
711    pub fn ancestors(&self, node_id: Index) -> AncestorsIter<T> {
712        AncestorsIter {
713            tree: self,
714            node_id: Some(node_id),
715        }
716    }
717
718    /// Return an iterator of references to this node and its descendants, in tree order.
719    fn traverse(&self, node_id: Index) -> TraverseIter<T> {
720        TraverseIter {
721            tree: self,
722            root: node_id,
723            next: Some(NodeEdge::Start(node_id)),
724        }
725    }
726
727    /// Return an iterator of references to this node and its descendants, with deoth in the tree,
728    /// in tree order.
729    fn traverse_with_depth(&self, node_id: Index) -> TraverseWithDepthIter<T> {
730        TraverseWithDepthIter {
731            tree: self,
732            root: node_id,
733            next: Some(NodeEdgeWithDepth::Start(node_id, 0)),
734        }
735    }
736
737    /// Return an iterator of references to this node and its descendants, in tree order.
738    ///
739    /// Parent nodes appear before the descendants.
740    /// Call `.next().unwrap()` once on the iterator to skip the node itself.
741    pub fn descendants(&self, node_id: Index) -> DescendantsIter<T> {
742        DescendantsIter(self.traverse(node_id))
743    }
744
745    /// Return an iterator of references to this node and its descendants, with deoth in the tree,
746    /// in tree order.
747    ///
748    /// Parent nodes appear before the descendants.
749    /// Call `.next().unwrap()` once on the iterator to skip the node itself.
750    pub fn descendants_with_depth(&self, node_id: Index) -> DescendantsWithDepthIter<T> {
751        DescendantsWithDepthIter(self.traverse_with_depth(node_id))
752    }
753}
754
755impl<T> fmt::Display for Node<T> {
756    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
757        write!(f, "Parent: {:?}, ", self.parent)?;
758        write!(f, "Previous sibling: {:?}, ", self.previous_sibling)?;
759        write!(f, "Next sibling: {:?}, ", self.next_sibling)?;
760        write!(f, "First child: {:?}, ", self.first_child)?;
761        write!(f, "Last child: {:?}", self.last_child)
762    }
763}
764
765impl<T> ops::Index<Index> for VecTree<T> {
766    type Output = T;
767
768    fn index(&self, index: Index) -> &Self::Output {
769        self.get(index).unwrap()
770    }
771}
772
773impl<T> ops::IndexMut<Index> for VecTree<T> {
774    fn index_mut(&mut self, index: Index) -> &mut Self::Output {
775        self.get_mut(index).unwrap()
776    }
777}
778
779macro_rules! impl_node_iterator {
780    ($name:ident, $next:expr) => {
781        impl<'a, T> Iterator for $name<'a, T> {
782            type Item = Index;
783
784            fn next(&mut self) -> Option<Index> {
785                match self.node_id.take() {
786                    Some(node_id) => {
787                        self.node_id = $next(&self.tree.nodes[node_id]);
788                        Some(node_id)
789                    }
790                    None => None,
791                }
792            }
793        }
794    };
795}
796
797/// An iterator of references to the children of a given node.
798pub struct ChildrenIter<'a, T: 'a> {
799    tree: &'a VecTree<T>,
800    node_id: Option<Index>,
801}
802impl_node_iterator!(ChildrenIter, |node: &Node<T>| node.next_sibling);
803
804/// An iterator of references to the siblings before a given node.
805pub struct PrecedingSiblingsIter<'a, T: 'a> {
806    tree: &'a VecTree<T>,
807    node_id: Option<Index>,
808}
809impl_node_iterator!(PrecedingSiblingsIter, |node: &Node<T>| node
810    .previous_sibling);
811
812/// An iterator of references to the siblings after a given node.
813pub struct FollowingSiblingsIter<'a, T: 'a> {
814    tree: &'a VecTree<T>,
815    node_id: Option<Index>,
816}
817impl_node_iterator!(FollowingSiblingsIter, |node: &Node<T>| node.next_sibling);
818
819/// An iterator of references to the ancestors a given node.
820pub struct AncestorsIter<'a, T: 'a> {
821    tree: &'a VecTree<T>,
822    node_id: Option<Index>,
823}
824impl_node_iterator!(AncestorsIter, |node: &Node<T>| node.parent);
825
826#[derive(Debug, Clone)]
827/// Indicator if the node is at a start or endpoint of the tree
828pub enum NodeEdge<T> {
829    /// Indicates that start of a node that has children. Yielded by `TraverseIter::next` before the
830    /// node’s descendants.
831    Start(T),
832
833    /// Indicates that end of a node that has children. Yielded by `TraverseIter::next` after the
834    /// node’s descendants.
835    End(T),
836}
837
838/// An iterator of references to a given node and its descendants, in depth-first search pre-order
839/// NLR traversal.
840/// https://en.wikipedia.org/wiki/Tree_traversal#Pre-order_(NLR)
841pub struct TraverseIter<'a, T: 'a> {
842    tree: &'a VecTree<T>,
843    root: Index,
844    next: Option<NodeEdge<Index>>,
845}
846
847impl<'a, T> Iterator for TraverseIter<'a, T> {
848    type Item = NodeEdge<Index>;
849
850    fn next(&mut self) -> Option<NodeEdge<Index>> {
851        match self.next.take() {
852            Some(item) => {
853                self.next = match item {
854                    NodeEdge::Start(node_id) => match self.tree.nodes[node_id].first_child {
855                        Some(first_child) => Some(NodeEdge::Start(first_child)),
856                        None => Some(NodeEdge::End(node_id)),
857                    },
858                    NodeEdge::End(node_id) => {
859                        if node_id == self.root {
860                            None
861                        } else {
862                            match self.tree.nodes[node_id].next_sibling {
863                                Some(next_sibling) => Some(NodeEdge::Start(next_sibling)),
864                                None => {
865                                    match self.tree.nodes[node_id].parent {
866                                        Some(parent) => Some(NodeEdge::End(parent)),
867
868                                        // `self.tree.nodes[node_id].parent` here can only be `None`
869                                        // if the tree has been modified during iteration, but
870                                        // silently stoping iteration seems a more sensible behavior
871                                        // than panicking.
872                                        None => None,
873                                    }
874                                }
875                            }
876                        }
877                    }
878                };
879                Some(item)
880            }
881            None => None,
882        }
883    }
884}
885
886/// An iterator of references to a given node and its descendants, in tree order.
887pub struct DescendantsIter<'a, T: 'a>(pub TraverseIter<'a, T>);
888
889impl<'a, T> Iterator for DescendantsIter<'a, T> {
890    type Item = Index;
891
892    fn next(&mut self) -> Option<Index> {
893        loop {
894            match self.0.next() {
895                Some(NodeEdge::Start(node_id)) => return Some(node_id),
896                Some(NodeEdge::End(_)) => {}
897                None => return None,
898            }
899        }
900    }
901}
902
903#[derive(Debug, Clone)]
904/// Indicator if the node is at a start or endpoint of the tree
905pub enum NodeEdgeWithDepth<T> {
906    /// Indicates that start of a node that has children. Yielded by `TraverseIter::next` before the
907    /// node’s descendants.
908    Start(T, u32),
909
910    /// Indicates that end of a node that has children. Yielded by `TraverseIter::next` after the
911    /// node’s descendants.
912    End(T, u32),
913}
914
915/// An iterator of references to a given node and its descendants, with depth, in depth-first
916/// search pre-order NLR traversal.
917/// https://en.wikipedia.org/wiki/Tree_traversal#Pre-order_(NLR)
918pub struct TraverseWithDepthIter<'a, T: 'a> {
919    tree: &'a VecTree<T>,
920    root: Index,
921    next: Option<NodeEdgeWithDepth<Index>>,
922}
923
924impl<'a, T> Iterator for TraverseWithDepthIter<'a, T> {
925    type Item = NodeEdgeWithDepth<Index>;
926
927    fn next(&mut self) -> Option<NodeEdgeWithDepth<Index>> {
928        match self.next.take() {
929            Some(item) => {
930                self.next = match item {
931                    NodeEdgeWithDepth::Start(node_id, depth) => {
932                        match self.tree.nodes[node_id].first_child {
933                            Some(first_child) => {
934                                Some(NodeEdgeWithDepth::Start(first_child, depth + 1))
935                            }
936                            None => Some(NodeEdgeWithDepth::End(node_id, depth)),
937                        }
938                    }
939                    NodeEdgeWithDepth::End(node_id, depth) => {
940                        if node_id == self.root {
941                            None
942                        } else {
943                            match self.tree.nodes[node_id].next_sibling {
944                                Some(next_sibling) => {
945                                    Some(NodeEdgeWithDepth::Start(next_sibling, depth))
946                                }
947                                None => {
948                                    match self.tree.nodes[node_id].parent {
949                                        Some(parent) => {
950                                            Some(NodeEdgeWithDepth::End(parent, depth - 1))
951                                        }
952
953                                        // `self.tree.nodes[node_id].parent` here can only be `None`
954                                        // if the tree has been modified during iteration, but
955                                        // silently stoping iteration seems a more sensible behavior
956                                        // than panicking.
957                                        None => None,
958                                    }
959                                }
960                            }
961                        }
962                    }
963                };
964                Some(item)
965            }
966            None => None,
967        }
968    }
969}
970
971/// An iterator of references to a given node and its descendants, with depth, in tree order.
972pub struct DescendantsWithDepthIter<'a, T: 'a>(pub TraverseWithDepthIter<'a, T>);
973
974impl<'a, T> Iterator for DescendantsWithDepthIter<'a, T> {
975    type Item = (Index, u32);
976
977    fn next(&mut self) -> Option<(Index, u32)> {
978        loop {
979            match self.0.next() {
980                Some(NodeEdgeWithDepth::Start(node_id, depth)) => return Some((node_id, depth)),
981                Some(NodeEdgeWithDepth::End(_, _)) => {}
982                None => return None,
983            }
984        }
985    }
986}