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

1 files changed, 19 insertions, 6 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);
     }