1 files changed, 27 insertions, 1 deletions
diff --git a/src/main/java/au/id/zancanaro/javacheck/RoseTree.java b/src/main/java/au/id/zancanaro/javacheck/RoseTree.java
index d9fb508..4d7b770 100644
--- a/src/main/java/au/id/zancanaro/javacheck/RoseTree.java
+++ b/src/main/java/au/id/zancanaro/javacheck/RoseTree.java
@@ -38,7 +38,7 @@ public class RoseTree<T> {
                         tree.getValue().children.iterator()));
     }
 
-    public static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
+    private static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
         return Iterators.flatten(Iterators.rangeIterator(trees.length, index ->
                         Iterators.mappingIterator(child -> {
                             @SuppressWarnings("unchecked")
@@ -64,6 +64,32 @@ public class RoseTree<T> {
                         RoseTree.permutations(trees)));
     }
 
+    private static <T> Iterator<RoseTree<T>[]> removeEach(RoseTree<T>[] trees) {
+        return Iterators.concat(
+                Iterators.rangeIterator(trees.length, index -> {
+                    @SuppressWarnings("unchecked")
+                    RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length - 1];
+                    for (int i = 0; i < trees.length - 1; ++i) {
+                        result[i] = trees[(i >= index ? i + 1 : i)];
+                    }
+                    return result;
+                }),
+                permutations(trees));
+    }
+
+    public static <T, R> RoseTree<R> shrink(Function<List<T>, R> fn, RoseTree<T>[] trees) {
+        @SuppressWarnings("unchecked")
+        List<T> heads = new ArrayList(trees.length);
+        for (RoseTree<T> tree : trees) {
+            heads.add(tree.getValue());
+        }
+        return new RoseTree<>(
+                fn.apply(heads),
+                () -> Iterators.mappingIterator(
+                        roses -> RoseTree.shrink(fn, roses),
+                        RoseTree.removeEach(trees)));
+    }
+
     public <R> RoseTree<R> fmap(Function<T, R> f) {
         return new RoseTree<>(
                 f.apply(this.value),