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        tree.children[0].children.push(Node::new(20));
190        tree.children[1].children.push(Node::new(30));
191
192        // Verify the modified structure
193        let final_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
194        assert_eq!(final_values, vec![1, 2, 20, 3, 30]);
195    }
196
197    #[test]
198    fn test_forest_traversal() {
199        // Create a forest with two trees
200        let mut forest = vec![
201            Node {
202                value: 1,
203                children: vec![Node::new(2)],
204            },
205            Node {
206                value: 3,
207                children: vec![Node::new(4)],
208            },
209        ];
210        // Create a mutable iterator over the forest
211        let mut iter = TreeIterMut::<'_, _, BreadthFirst>::new(forest.iter_mut());
212        while let Some(mut node) = iter.next() {
213            node.value += 10;
214        }
215        // Verify the modified structure
216        let value = TreeIter::<'_, _, BreadthFirst>::new(forest.iter())
217            .map(|node| node.value)
218            .collect::<Vec<_>>();
219        assert_eq!(value, vec![11, 13, 12, 14]);
220    }
221
222    #[test]
223    fn test_complex_tree_traversal() {
224        // Create a more complex tree
225        let mut tree = Node {
226            value: 1,
227            children: vec![
228                Node {
229                    value: 2,
230                    children: vec![
231                        Node {
232                            value: 4,
233                            children: vec![Node::new(7)],
234                        },
235                        Node::new(5),
236                    ],
237                },
238                Node {
239                    value: 3,
240                    children: vec![Node {
241                        value: 6,
242                        children: vec![Node::new(8), Node::new(9)],
243                    }],
244                },
245            ],
246        };
247
248        // Test depth-first traversal
249        let df_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
250        assert_eq!(df_values, vec![1, 2, 4, 7, 5, 3, 6, 8, 9]);
251
252        // Test mutable depth-first traversal
253        let mut df_values_mut = vec![];
254        let mut iter = tree.iter_mut::<DepthFirst>();
255        while let Some(node) = iter.next() {
256            df_values_mut.push(node.value);
257        }
258        assert_eq!(df_values_mut, vec![1, 2, 4, 7, 5, 3, 6, 8, 9]);
259
260        // Test breadth-first traversal
261        let bf_values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
262        assert_eq!(bf_values, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
263
264        // Test mutable breadth-first traversal
265        let mut bf_values_mut = vec![];
266        let mut iter = tree.iter_mut::<BreadthFirst>();
267        while let Some(node) = iter.next() {
268            bf_values_mut.push(node.value);
269        }
270        assert_eq!(bf_values_mut, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
271    }
272}