1 files changed, 29 insertions, 36 deletions

diff --git a/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java b/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
index 7dff917..bdc7604 100644
--- a/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
+++ b/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
@@ -41,55 +41,50 @@ public class ShrinkTree<T> {
                         tree.getValue().children.iterator()));
     }
 
-    private static <T> Iterator<ShrinkTree<T>[]> permutations(ShrinkTree<T>[] trees) {
+    private static <T> List<T> makeHeadList(ShrinkTree<T>[] trees) {
+        List<T> heads = new ArrayList<>(trees.length);
+        for (ShrinkTree<T> tree : trees) {
+            heads.add(tree.getValue());
+        }
+        return heads;
+    }
+
+    public static <T> Iterator<ShrinkTree<T>[]> promoteChildren(ShrinkTree<T>[] trees) {
         return flatten(
-                rangeIterator(trees.length,
-                        index -> mappingIterator(child -> {
+                rangeIterator(trees.length, index ->
+                        mappingIterator(child -> {
                             @SuppressWarnings("unchecked")
                             ShrinkTree<T>[] result = (ShrinkTree<T>[]) new ShrinkTree[trees.length];
                             for (int i = 0; i < trees.length; ++i) {
                                 result[i] = (i == index ? child : trees[i]);
                             }
                             return result;
-                        }, trees[index].getChildren())
-                ));
-    }
-
-    private static <T> List<T> makeHeadList(ShrinkTree<T>[] trees) {
-        List<T> heads = new ArrayList<>(trees.length);
-        for (ShrinkTree<T> tree : trees) {
-            heads.add(tree.getValue());
-        }
-        return heads;
+                        }, trees[index].getChildren())));
     }
 
-    public static <T, R> ShrinkTree<R> zip(Function<List<T>, R> fn, ShrinkTree<T>[] trees) {
-        return new ShrinkTree<>(
-                fn.apply(makeHeadList(trees)),
-                () -> mappingIterator(
-                        shrinks -> ShrinkTree.zip(fn, shrinks),
-                        ShrinkTree.permutations(trees)));
+    public static <T> Iterator<ShrinkTree<T>[]> removeChildren(ShrinkTree<T>[] trees) {
+        return rangeIterator(trees.length, index -> {
+            @SuppressWarnings("unchecked")
+            ShrinkTree<T>[] result = (ShrinkTree<T>[]) new ShrinkTree[trees.length - 1];
+            for (int i = 0; i < trees.length - 1; ++i) {
+                result[i] = trees[(i >= index ? i + 1 : i)];
+            }
+            return result;
+        });
     }
 
-    private static <T> Iterator<ShrinkTree<T>[]> removeEach(ShrinkTree<T>[] trees) {
-        return concat(
-                rangeIterator(trees.length, index -> {
-                    @SuppressWarnings("unchecked")
-                    ShrinkTree<T>[] result = (ShrinkTree<T>[]) new ShrinkTree[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> Iterator<ShrinkTree<T>[]> removeAndPromoteChildren(ShrinkTree<T>[] trees) {
+        return concat(removeChildren(trees), promoteChildren(trees));
     }
 
-    public static <T, R> ShrinkTree<R> shrink(Function<List<T>, R> fn, ShrinkTree<T>[] trees) {
+    public static <T> ShrinkTree<List<T>> combine(
+            ShrinkTree<T>[] trees,
+            Function<ShrinkTree<T>[], Iterator<ShrinkTree<T>[]>> processChildren) {
         return new ShrinkTree<>(
-                fn.apply(makeHeadList(trees)),
+                makeHeadList(trees),
                 () -> mappingIterator(
-                        shrinks -> ShrinkTree.shrink(fn, shrinks),
-                        ShrinkTree.removeEach(trees)));
+                        shrinks -> ShrinkTree.combine(shrinks, processChildren),
+                        processChildren.apply(trees)));
     }
 
     public <R> ShrinkTree<R> map(Function<T, R> f) {
@@ -126,12 +121,10 @@ public class ShrinkTree<T> {
                         strategy.shrink(value));
     }
 
-    @SuppressWarnings("unused")
     public void print(Writer output) throws IOException {
         print(output, Object::toString);
     }
 
-    @SuppressWarnings("unused")
     public void print(Writer output, Function<T, String> toString) throws IOException {
         output.write(toString.apply(this.getValue()));
         output.write('[');