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}