Clean up a bit of the ShrinkTree stuff, particularly for shrinking

3 files changed, 49 insertions, 44 deletions

diff --git a/src/main/java/au/id/zancanaro/javacheck/Generator.java b/src/main/java/au/id/zancanaro/javacheck/Generator.java
index 103bc69..0679f10 100644
--- a/src/main/java/au/id/zancanaro/javacheck/Generator.java
+++ b/src/main/java/au/id/zancanaro/javacheck/Generator.java
@@ -72,7 +72,7 @@ public interface Generator<T> {
             for (Generator<? extends T> generator : generators) {
                 result[index++] = generator.generate(random, size).map(x -> (T) x);
             }
-            return ShrinkTree.zip(Function.identity(), result);
+            return ShrinkTree.combine(result, ShrinkTree::promoteChildren);
         };
     }
 
@@ -81,9 +81,7 @@ public interface Generator<T> {
      * element taken from the provided generator.
      *
      * Shrinking for this type involves attempting to remove terms and shrink
-     * each subtree in turn, recursively. (If the length must remain fixed then
-     * the {@link ShrinkTree} produced by this generator should be filtered with
-     * {@link ShrinkTree#filter(java.util.function.Predicate)}.
+     * each subtree in turn, recursively.
      *
      * @param count     The length of the list to generate
      * @param generator The generator to use for each term in the generated
@@ -98,7 +96,7 @@ public interface Generator<T> {
             for (int i = 0; i < count; ++i) {
                 result[i] = generator.generate(random, size);
             }
-            return ShrinkTree.shrink(Function.identity(), result);
+            return ShrinkTree.combine(result, ShrinkTree::removeAndPromoteChildren);
         };
     }
 
@@ -128,7 +126,10 @@ public interface Generator<T> {
      * result of this
      */
     default <R> Generator<R> flatMap(Function<T, Generator<R>> action) {
-        return (random, size) -> ShrinkTree.join(this.generate(random, size).map(action).map(g -> g.generate(random, size)));
+        return (random, size) -> ShrinkTree.join(
+                this.generate(random, size)
+                        .map(action
+                                .andThen(g -> g.generate(random, size))));
     }
 
     /**
@@ -156,6 +157,18 @@ public interface Generator<T> {
         };
     }
 
+    /**
+     * Create a new generator which generates values with a shrink tree
+     * determined by the provided {@link ShrinkStrategy}.
+     *
+     * @param strategy The shrink strategy by which to shrink generated values
+     * @return A new generator which will shrink values from this according to
+     * the provided strategy
+     */
+    default Generator<T> withShrinkStrategy(ShrinkStrategy<T> strategy) {
+        return (random, size) -> this.generate(random, size).withShrinkStrategy(strategy);
+    }
+
     default Iterator<T> sample(Random random, int maxSize) {
         return new GeneratorSampleIterator<>(this, random, maxSize);
     }
diff --git a/src/main/java/au/id/zancanaro/javacheck/Generators.java b/src/main/java/au/id/zancanaro/javacheck/Generators.java
index e5699ac..9730bac 100644
--- a/src/main/java/au/id/zancanaro/javacheck/Generators.java
+++ b/src/main/java/au/id/zancanaro/javacheck/Generators.java
@@ -29,7 +29,7 @@ public final class Generators {
      * removed
      */
     public static <T> Generator<T> noShrink(Generator<T> gen) {
-        return (random, size) -> ShrinkTree.pure(gen.generate(random, size).getValue());
+        return gen.withShrinkStrategy(value -> Iterators.emptyIterator());
     }
 
     @SafeVarargs
@@ -230,7 +230,6 @@ public final class Generators {
                 chars[i++] = c;
             }
             return String.valueOf(chars);
-//            return new String(chars);
         });
     }
 }
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('[');