Add a tree map operation

tree/src/lib.rs

@@ -9,7 +9,7 @@ use std::{
     rc::Rc,
 };
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub enum Tree<T> {
     Empty,
     Root(Node<T>),
@@ -22,8 +22,16 @@ impl<T> Default for Tree<T> {
 }
 
 impl<T> Tree<T> {
-    pub fn set_value(&mut self, value: T) {
+    pub fn new(value: T) -> (Tree<T>, Node<T>) {
-        *self = Tree::Root(Node::new(value));
+        let node = Node::new(value);
+        let tree = Tree::Root(node.clone());
+        (tree, node)
+    }
+
+    pub fn set_value(&mut self, value: T) -> Node<T> {
+        let node = Node::new(value);
+        *self = Tree::Root(node.clone());
+        node
     }
 
     /// Use a breadth-first-search pattern to find a node, returning the node if found.
@@ -51,6 +59,33 @@ impl<T> Tree<T> {
         }
         None
     }
+
+    pub fn map<F, U>(&self, op: F) -> Tree<U>
+    where
+        F: FnOnce(&Node<T>) -> Node<U> + Copy,
+    {
+        match self {
+            Tree::Empty => Tree::Empty,
+            Tree::Root(root) => {
+                let new_root = op(root);
+                let mut queue: VecDeque<(Node<T>, Node<U>)> = root
+                    .children()
+                    .iter()
+                    .map(|child| (child.clone(), new_root.clone()))
+                    .collect();
+
+                while let Some((source, dest)) = queue.pop_front() {
+                    let res = op(&source);
+                    dest.add_child_node(res.clone());
+
+                    for child in source.children().iter() {
+                        queue.push_back((child.clone(), res.clone()));
+                    }
+                }
+                Tree::Root(new_root)
+            }
+        }
+    }
 }
 
 // By using the Rc<RefCell> container here, I'm able to make Node easily clonable while still
@@ -100,8 +135,10 @@ impl<T> Node<T> {
         self.0.borrow_mut().children.push(child)
     }
 
-    pub fn add_child_value(&self, value: T) {
+    pub fn add_child_value(&self, value: T) -> Node<T> {
-        self.0.borrow_mut().children.push(Node::new(value))
+        let node = Node::new(value);
+        self.0.borrow_mut().children.push(node.clone());
+        node
     }
 }
 
@@ -130,4 +167,17 @@ mod tests {
         assert_eq!(*n.value(), 15);
         // assert_eq!(n.children(), vec![Rc::new(RefCell::new(Node::new(20)))]);
     }
+
+    #[test]
+    fn it_can_map_one_tree_to_another() {
+        let (tree, n) = Tree::new(15);
+        let n = n.add_child_value(16);
+        let _ = n.add_child_value(17);
+
+        let tree2 = tree.map(|v| Node::new(v.value().to_string()));
+
+        assert!(tree2.find_bfs(|val| *val == "15").is_some());
+        assert!(tree2.find_bfs(|val| *val == "16").is_some());
+        assert!(tree2.find_bfs(|val| *val == "17").is_some());
+    }
 }