Add select and next_generation

examples/generate.rs

@@ -51,13 +51,27 @@ fn next_level(value: &Coordinates) -> Vec<(Direction, Coordinates)> {
 }
 
 fn main() {
+    println!("Generating");
     let mut tree = Tree::generations(Coordinates::default(), 3, next_level);
 
     println!("{:#?}", tree);
     println!("{:?}", tree.leaves());
 
+    println!("Pruning");
     tree.retain(|item| item.x > 0 && item.y > 0);
 
     println!("{:#?}", tree);
     println!("{:?}", tree.leaves());
+
+    println!("Selecting");
+    let mut tree = tree.select(&Direction::Up).unwrap();
+
+    println!("{:#?}", tree);
+    println!("{:?}", tree.leaves());
+
+    println!("Next generation");
+    tree.next_generation(next_level);
+
+    println!("{:#?}", tree);
+    println!("{:?}", tree.leaves());
 }

src/lib.rs

@@ -37,6 +37,32 @@ where
         }
     }
 
+    pub fn next_generation<F>(&mut self, callback: F)
+    where
+        F: Fn(&T) -> Vec<(K, T)> + Copy,
+    {
+        if self.children.is_empty() {
+            let children = (callback)(&self.value);
+            self.children.extend(children.into_iter().map(|(k, v)| {
+                (
+                    k,
+                    Tree {
+                        value: v,
+                        children: HashMap::new(),
+                    },
+                )
+            }));
+        } else {
+            for child in self.children.values_mut() {
+                child.next_generation(callback);
+            }
+        }
+    }
+
+    pub fn select(mut self, path: &K) -> Option<Self> {
+        self.children.remove(path)
+    }
+
     pub fn leaves(&self) -> Vec<(Vec<&K>, &T)> {
         self.do_leaves(vec![])
     }