use crate::{iter::TreeNode, prelude::TreeNodeMut};
#[derive(Debug, Clone, PartialEq, Eq, Default, Hash)]
pub struct Node<T> {
pub value: T,
pub children: Vec<Node<T>>,
}
impl<T> Node<T> {
pub fn new(value: T) -> Self {
Self {
value,
children: Vec::new(),
}
}
}
impl<T> TreeNode for Node<T> {
fn children(&self) -> impl DoubleEndedIterator<Item = &Self> {
self.children.iter()
}
}
impl<T> TreeNodeMut for Node<T> {
fn children_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut Self> {
self.children.iter_mut()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::prelude::*;
#[test]
fn test_empty_tree() {
let tree: Node<i32> = Node::new(42);
let values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
assert_eq!(values, vec![42]);
let values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
assert_eq!(values, vec![42]);
}
#[test]
fn test_depth_first_traversal() {
let tree = Node {
value: 1,
children: vec![
Node {
value: 2,
children: vec![Node::new(4), Node::new(5)],
},
Node {
value: 3,
children: vec![Node::new(6)],
},
],
};
let values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
assert_eq!(values, vec![1, 2, 4, 5, 3, 6]);
}
#[test]
fn test_breadth_first_traversal() {
let tree = Node {
value: 1,
children: vec![
Node {
value: 2,
children: vec![Node::new(4), Node::new(5)],
},
Node {
value: 3,
children: vec![Node::new(6)],
},
],
};
let values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
assert_eq!(values, vec![1, 2, 3, 4, 5, 6]);
}
#[test]
fn test_mutable_depth_first_traversal() {
let mut tree = Node {
value: 1,
children: vec![Node::new(2), Node::new(3)],
};
let mut iter = tree.iter_mut::<DepthFirst>();
while let Some(mut node) = iter.next() {
node.value *= 2;
}
let values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
assert_eq!(values, vec![2, 4, 6]);
}
#[test]
fn test_mutable_breadth_first_traversal() {
let mut tree = Node {
value: 1,
children: vec![
Node {
value: 2,
children: vec![Node::new(4)],
},
Node::new(3),
],
};
let mut iter = tree.iter_mut::<BreadthFirst>();
while let Some(mut node) = iter.next() {
if node.value == 2 {
node.children.push(Node::new(10));
}
node.value += 10;
}
let values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
assert_eq!(values, vec![11, 12, 13, 14, 20]);
}
#[test]
fn test_tree_modification() {
let mut tree = Node {
value: 1,
children: vec![Node::new(2), Node::new(3)],
};
let initial_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
assert_eq!(initial_values, vec![1, 2, 3]);
tree.children[0].children.push(Node::new(20));
tree.children[1].children.push(Node::new(30));
let final_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
assert_eq!(final_values, vec![1, 2, 20, 3, 30]);
}
#[test]
fn test_forest_traversal() {
let mut forest = vec![
Node {
value: 1,
children: vec![Node::new(2)],
},
Node {
value: 3,
children: vec![Node::new(4)],
},
];
let mut iter = TreeIterMut::<'_, _, BreadthFirst>::new(forest.iter_mut());
while let Some(mut node) = iter.next() {
node.value += 10;
}
let value = TreeIter::<'_, _, BreadthFirst>::new(forest.iter())
.map(|node| node.value)
.collect::<Vec<_>>();
assert_eq!(value, vec![11, 13, 12, 14]);
}
#[test]
fn test_complex_tree_traversal() {
let mut tree = Node {
value: 1,
children: vec![
Node {
value: 2,
children: vec![
Node {
value: 4,
children: vec![Node::new(7)],
},
Node::new(5),
],
},
Node {
value: 3,
children: vec![Node {
value: 6,
children: vec![Node::new(8), Node::new(9)],
}],
},
],
};
let df_values: Vec<i32> = tree.iter::<DepthFirst>().map(|n| n.value).collect();
assert_eq!(df_values, vec![1, 2, 4, 7, 5, 3, 6, 8, 9]);
let mut df_values_mut = vec![];
let mut iter = tree.iter_mut::<DepthFirst>();
while let Some(node) = iter.next() {
df_values_mut.push(node.value);
}
assert_eq!(df_values_mut, vec![1, 2, 4, 7, 5, 3, 6, 8, 9]);
let bf_values: Vec<i32> = tree.iter::<BreadthFirst>().map(|n| n.value).collect();
assert_eq!(bf_values, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
let mut bf_values_mut = vec![];
let mut iter = tree.iter_mut::<BreadthFirst>();
while let Some(node) = iter.next() {
bf_values_mut.push(node.value);
}
assert_eq!(bf_values_mut, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);
}
}