1 files changed, 22 insertions, 8 deletions
diff --git a/src/main/java/au/id/zancanaro/RoseTree.java b/src/main/java/au/id/zancanaro/RoseTree.java
index 2a83ad2..bbcb949 100644
--- a/src/main/java/au/id/zancanaro/RoseTree.java
+++ b/src/main/java/au/id/zancanaro/RoseTree.java
@@ -4,6 +4,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 public class RoseTree<T> {
     private final T value;
@@ -36,14 +37,14 @@ public class RoseTree<T> {
 
     public static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
         return Iterators.flatten(Iterators.rangeIterator(trees.length, index ->
-                Iterators.mappingIterator(child -> {
-                    @SuppressWarnings("unchecked")
-                    RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length];
-                    for (int i = 0; i < trees.length; ++i) {
-                        result[i] = (i == index ? child : trees[i]);
-                    }
-                    return result;
-                }, trees[index].getChildren())
+                        Iterators.mappingIterator(child -> {
+                            @SuppressWarnings("unchecked")
+                            RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length];
+                            for (int i = 0; i < trees.length; ++i) {
+                                result[i] = (i == index ? child : trees[i]);
+                            }
+                            return result;
+                        }, trees[index].getChildren())
         ));
     }
 
@@ -69,4 +70,17 @@ public class RoseTree<T> {
     public <R> RoseTree<R> flatmap(Function<T, RoseTree<R>> f) {
         return RoseTree.join(this.fmap(f));
     }
+
+    public RoseTree<T> filter(Predicate<T> predicate) {
+        if (predicate.test(this.getValue())) {
+            return new RoseTree<>(
+                    this.getValue(),
+                    () -> Iterators.mappingIterator(tree -> tree.filter(predicate),
+                            Iterators.filteringIterator(
+                                    tree -> predicate.test(tree.getValue()),
+                                    this.getChildren())));
+        } else {
+            throw new IllegalArgumentException("Current value doesn't match predicate: whoops!");
+        }
+    }
 }