tree_iter/
iter.rs

1use std::{collections::VecDeque, marker::PhantomData};
2
3use crate::traversal_order::{BreadthFirst, DepthFirst, TraversalOrder};
4
5/// Trait for immutable tree traversal.
6///
7/// This trait defines the interface required to iterate over a tree structure
8/// in an immutable fashion. Any type that implements this trait can be traversed
9/// using the provided iterators.
10///
11/// # Examples
12///
13/// ```rust
14/// use tree_iter::prelude::*;
15/// use tree_iter::tree::Node;
16///
17/// // Create a custom tree structure
18/// struct MyTree<T> {
19///     value: T,
20///     children: Vec<MyTree<T>>,
21/// }
22///
23/// // Implement TreeNode for custom tree structure
24/// impl<T> TreeNode for MyTree<T> {
25///     fn children(&self) -> impl DoubleEndedIterator<Item = &Self> {
26///         self.children.iter()
27///     }
28/// }
29///
30/// // Now you can use the tree iterator
31/// let my_tree = MyTree {
32///     value: 1,
33///     children: vec![
34///         MyTree { value: 2, children: vec![] },
35///         MyTree { value: 3, children: vec![] },
36///     ],
37/// };
38///
39/// let values: Vec<i32> = my_tree.iter::<DepthFirst>()
40///                             .map(|node| node.value)
41///                             .collect();
42/// ```
43pub trait TreeNode {
44    /// Returns an iterator over the children of this node.
45    ///
46    /// This method must be implemented by all types implementing `TreeNode`.
47    fn children(&self) -> impl DoubleEndedIterator<Item = &Self>;
48
49    /// Creates an iterator that traverses the tree starting from this node.
50    ///
51    /// # Type Parameters
52    ///
53    /// * `T` - The traversal order strategy to use (e.g., `DepthFirst` or `BreadthFirst`).
54    fn iter<T: TraversalOrder>(&self) -> TreeIter<'_, Self, T>
55    where
56        Self: Sized,
57    {
58        TreeIter::new([self])
59    }
60}
61
62/// An iterator over tree nodes in a specified traversal order.
63///
64/// This struct provides the implementation for traversing a tree structure
65/// with nodes implementing the `TreeNode` trait.
66///
67/// # Type Parameters
68///
69/// * `'a` - The lifetime of the tree nodes being traversed.
70/// * `N` - The type of tree node.
71/// * `T` - The traversal order strategy (e.g., `DepthFirst` or `BreadthFirst`).
72#[derive(Debug)]
73pub struct TreeIter<'a, N, T> {
74    /// Queue of nodes to be visited.
75    nodes: VecDeque<&'a N>,
76    /// Phantom data to track the traversal order type.
77    _order: PhantomData<T>,
78}
79
80impl<'a, N, T: TraversalOrder> TreeIter<'a, N, T> {
81    /// Creates a new tree iterator from a collection of root nodes.
82    ///
83    /// # Parameters
84    ///
85    /// * `roots` - An iterator yielding the root nodes to start traversal from.
86    pub fn new(roots: impl IntoIterator<Item = &'a N>) -> Self {
87        Self {
88            nodes: roots.into_iter().collect(),
89            _order: PhantomData,
90        }
91    }
92}
93
94/// Implementation of `Iterator` for breadth-first traversal.
95impl<'a, N: TreeNode> Iterator for TreeIter<'a, N, BreadthFirst> {
96    type Item = &'a N;
97
98    /// Returns the next node in breadth-first order.
99    ///
100    /// Breadth-first traversal visits all nodes at a given depth before moving to the next level.
101    fn next(&mut self) -> Option<Self::Item> {
102        let node = self.nodes.pop_front()?;
103        self.nodes.extend(node.children());
104        Some(node)
105    }
106}
107
108/// Implementation of `Iterator` for depth-first traversal.
109impl<'a, N: TreeNode> Iterator for TreeIter<'a, N, DepthFirst> {
110    type Item = &'a N;
111
112    /// Returns the next node in depth-first order.
113    ///
114    /// Depth-first traversal explores as far down a branch as possible before backtracking.
115    fn next(&mut self) -> Option<Self::Item> {
116        let node = self.nodes.pop_front()?;
117        for child in node.children().rev() {
118            self.nodes.push_front(child);
119        }
120        Some(node)
121    }
122}