tree_iter/
tree.rs

1use crate::{iter::TreeNode, prelude::TreeNodeMut};
2
3/// A generic tree node implementation.
4///
5/// This struct provides a simple implementation of a tree node with a value and a vector of children.
6/// It implements both `TreeNode` and `TreeNodeMut` traits, allowing it to be used with both
7/// immutable and mutable iterators.
8///
9/// # Type Parameters
10///
11/// * `T` - The type of value stored in each node.
12///
13/// # Examples
14///
15/// ```rust
16/// use tree_iter::tree::Node;
17/// use tree_iter::prelude::*;
18///
19/// // Create a simple tree
20/// let tree = Node {
21///     value: 1,
22///     children: vec![Node::new(2), Node::new(3)],
23/// };
24///
25/// // Traverse in depth-first order
26/// let values: Vec<i32> = tree.iter::<DepthFirst>()
27///                           .map(|node| node.value)
28///                           .collect();
29/// assert_eq!(values, vec![1, 2, 3]);
30/// ```
31#[derive(Debug, Clone, PartialEq, Eq, Default, Hash)]
32pub struct Node<T> {
33    /// The value stored in this node.
34    pub value: T,
35    /// The children of this node.
36    pub children: Vec<Node<T>>,
37}
38
39impl<T> Node<T> {
40    /// Creates a new node with the given value and no children.
41    ///
42    /// # Parameters
43    ///
44    /// * `value` - The value to store in the node.
45    ///
46    /// # Returns
47    ///
48    /// A new `Node` with the given value and an empty children vector.
49    pub fn new(value: T) -> Self {
50        Self {
51            value,
52            children: Vec::new(),
53        }
54    }
55}
56
57/// Implementation of `TreeNode` for `Node<T>`.
58///
59/// This allows immutable iteration over the tree.
60impl<T> TreeNode for Node<T> {
61    /// Returns an iterator over the children of this node.
62    fn children(&self) -> impl DoubleEndedIterator<Item = &Self> {
63        self.children.iter()
64    }
65}
66
67/// Implementation of `TreeNodeMut` for `Node<T>`.
68///
69/// This allows mutable iteration over the tree.
70impl<T> TreeNodeMut for Node<T> {
71    /// Returns a mutable iterator over the children of this node.
72    fn children_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut Self> {
73        self.children.iter_mut()
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use crate::prelude::*;
81
82    #[test]
83    fn test_empty_tree() {
84        let tree: Node<i32> = Node::new(42);
85        let values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
86        assert_eq!(values, vec![42]);
87
88        let values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
89        assert_eq!(values, vec![42]);
90    }
91
92    #[test]
93    fn test_depth_first_traversal() {
94        let tree = Node {
95            value: 1,
96            children: vec![
97                Node {
98                    value: 2,
99                    children: vec![Node::new(4), Node::new(5)],
100                },
101                Node {
102                    value: 3,
103                    children: vec![Node::new(6)],
104                },
105            ],
106        };
107
108        let values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
109        assert_eq!(values, vec![1, 2, 4, 5, 3, 6]);
110    }
111
112    #[test]
113    fn test_breadth_first_traversal() {
114        let tree = Node {
115            value: 1,
116            children: vec![
117                Node {
118                    value: 2,
119                    children: vec![Node::new(4), Node::new(5)],
120                },
121                Node {
122                    value: 3,
123                    children: vec![Node::new(6)],
124                },
125            ],
126        };
127
128        let values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
129        assert_eq!(values, vec![1, 2, 3, 4, 5, 6]);
130    }
131
132    #[test]
133    fn test_mutable_depth_first_traversal() {
134        let mut tree = Node {
135            value: 1,
136            children: vec![Node::new(2), Node::new(3)],
137        };
138
139        // Double each value in the tree
140        let mut iter = tree.iter_mut::<DepthFirst>();
141        while let Some(mut node) = iter.next() {
142            node.value *= 2;
143        }
144
145        // Verify values were changed
146        let values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
147        assert_eq!(values, vec![2, 4, 6]);
148    }
149
150    #[test]
151    fn test_mutable_breadth_first_traversal() {
152        let mut tree = Node {
153            value: 1,
154            children: vec![
155                Node {
156                    value: 2,
157                    children: vec![Node::new(4)],
158                },
159                Node::new(3),
160            ],
161        };
162
163        // Add 10 to each value in the tree, and add a child to node with value 2
164        let mut iter = tree.iter_mut::<BreadthFirst>();
165        while let Some(mut node) = iter.next() {
166            if node.value == 2 {
167                node.children.push(Node::new(10));
168            }
169            node.value += 10;
170        }
171
172        // Verify values were changed in breadth-first order
173        let values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
174        assert_eq!(values, vec![11, 12, 13, 14, 20]);
175    }
176
177    #[test]
178    fn test_tree_modification() {
179        // Create a tree manually with specific structure
180        let mut tree = Node {
181            value: 1,
182            children: vec![Node::new(2), Node::new(3)],
183        };
184
185        // Verify initial structure
186        let initial_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
187        assert_eq!(initial_values, vec![1, 2, 3]);
188
189        // Modify the tree directly, avoiding iterator with borrow issues
190        tree.children[0].children.push(Node::new(20));
191        tree.children[1].children.push(Node::new(30));
192
193        // Verify the modified structure
194        let final_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
195        assert_eq!(final_values, vec![1, 2, 20, 3, 30]);
196    }
197
198    #[test]
199    fn test_forest_traversal() {
200        // Create a forest with two trees
201        let mut forest = vec![
202            Node {
203                value: 1,
204                children: vec![Node::new(2)],
205            },
206            Node {
207                value: 3,
208                children: vec![Node::new(4)],
209            },
210        ];
211        // Create a mutable iterator over the forest
212        let mut iter = TreeMutIter::<'_, _, BreadthFirst>::new(forest.iter_mut());
213        while let Some(mut node) = iter.next() {
214            node.value += 10;
215        }
216        // Verify the modified structure
217        let value = TreeIter::<'_, _, BreadthFirst>::new(forest.iter())
218            .map(|node| node.value)
219            .collect::<Vec<_>>();
220        assert_eq!(value, vec![11, 13, 12, 14]);
221    }
222
223    #[test]
224    fn test_complex_tree_traversal() {
225        // Create a more complex tree
226        let mut tree = Node {
227            value: 1,
228            children: vec![
229                Node {
230                    value: 2,
231                    children: vec![
232                        Node {
233                            value: 4,
234                            children: vec![Node::new(7)],
235                        },
236                        Node::new(5),
237                    ],
238                },
239                Node {
240                    value: 3,
241                    children: vec![Node {
242                        value: 6,
243                        children: vec![Node::new(8), Node::new(9)],
244                    }],
245                },
246            ],
247        };
248
249        // Test depth-first traversal
250        let df_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
251        assert_eq!(df_values, vec![1, 2, 4, 7, 5, 3, 6, 8, 9]);
252
253        // Test mutable depth-first traversal
254        let mut df_values_mut = vec![];
255        let mut iter = tree.iter_mut::<DepthFirst>();
256        while let Some(node) = iter.next() {
257            df_values_mut.push(node.value);
258        }
259        assert_eq!(df_values_mut, vec![1, 2, 4, 7, 5, 3, 6, 8, 9]);
260
261        // Test breadth-first traversal
262        let bf_values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
263        assert_eq!(bf_values, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
264
265        // Test mutable breadth-first traversal
266        let mut bf_values_mut = vec![];
267        let mut iter = tree.iter_mut::<BreadthFirst>();
268        while let Some(node) = iter.next() {
269            bf_values_mut.push(node.value);
270        }
271        assert_eq!(bf_values_mut, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
272    }
273}