1 files changed, 30 insertions, 14 deletions

diff --git a/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java b/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
index a424806..7dff917 100644
--- a/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
+++ b/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
@@ -2,16 +2,21 @@ package au.id.zancanaro.javacheck;
 
 import java.io.IOException;
 import java.io.Writer;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
 import java.util.function.Function;
 import java.util.function.Predicate;
 
+import static au.id.zancanaro.javacheck.Iterators.*;
+
 @SuppressWarnings("unused")
 public class ShrinkTree<T> {
     private final T value;
     private final Iterable<ShrinkTree<T>> children;
 
-    public ShrinkTree(T value, Iterable<ShrinkTree<T>> children) {
+    private ShrinkTree(T value, Iterable<ShrinkTree<T>> children) {
         this.value = value;
         this.children = children;
     }
@@ -31,15 +36,15 @@ public class ShrinkTree<T> {
     public static <T> ShrinkTree<T> join(ShrinkTree<ShrinkTree<T>> tree) {
         return new ShrinkTree<>(
                 tree.getValue().getValue(),
-                () -> Iterators.concat(
-                        Iterators.mappingIterator(ShrinkTree::join, tree.children.iterator()),
+                () -> concat(
+                        mappingIterator(ShrinkTree::join, tree.children.iterator()),
                         tree.getValue().children.iterator()));
     }
 
     private static <T> Iterator<ShrinkTree<T>[]> permutations(ShrinkTree<T>[] trees) {
-        return Iterators.flatten(
-                Iterators.rangeIterator(trees.length,
-                        index -> Iterators.mappingIterator(child -> {
+        return flatten(
+                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) {
@@ -61,14 +66,14 @@ public class ShrinkTree<T> {
     public static <T, R> ShrinkTree<R> zip(Function<List<T>, R> fn, ShrinkTree<T>[] trees) {
         return new ShrinkTree<>(
                 fn.apply(makeHeadList(trees)),
-                () -> Iterators.mappingIterator(
+                () -> mappingIterator(
                         shrinks -> ShrinkTree.zip(fn, shrinks),
                         ShrinkTree.permutations(trees)));
     }
 
     private static <T> Iterator<ShrinkTree<T>[]> removeEach(ShrinkTree<T>[] trees) {
-        return Iterators.concat(
-                Iterators.rangeIterator(trees.length, index -> {
+        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) {
@@ -82,7 +87,7 @@ public class ShrinkTree<T> {
     public static <T, R> ShrinkTree<R> shrink(Function<List<T>, R> fn, ShrinkTree<T>[] trees) {
         return new ShrinkTree<>(
                 fn.apply(makeHeadList(trees)),
-                () -> Iterators.mappingIterator(
+                () -> mappingIterator(
                         shrinks -> ShrinkTree.shrink(fn, shrinks),
                         ShrinkTree.removeEach(trees)));
     }
@@ -90,7 +95,7 @@ public class ShrinkTree<T> {
     public <R> ShrinkTree<R> map(Function<T, R> f) {
         return new ShrinkTree<>(
                 f.apply(this.value),
-                () -> Iterators.mappingIterator(tree -> tree.map(f), this.children.iterator()));
+                () -> mappingIterator(tree -> tree.map(f), this.children.iterator()));
     }
 
     public <R> ShrinkTree<R> flatMap(Function<T, ShrinkTree<R>> f) {
@@ -101,8 +106,8 @@ public class ShrinkTree<T> {
         if (predicate.test(this.getValue())) {
             return new ShrinkTree<>(
                     this.getValue(),
-                    () -> Iterators.mappingIterator(tree -> tree.filter(predicate),
-                            Iterators.filteringIterator(
+                    () -> mappingIterator(tree -> tree.filter(predicate),
+                            filteringIterator(
                                     tree -> predicate.test(tree.getValue()),
                                     this.getChildren())));
         } else {
@@ -110,6 +115,17 @@ public class ShrinkTree<T> {
         }
     }
 
+    public ShrinkTree<T> withShrinkStrategy(ShrinkStrategy<T> strategy) {
+        return new ShrinkTree<>(this.getValue(), strategyIterable(this.getValue(), strategy));
+    }
+
+    private static <T> Iterable<ShrinkTree<T>> strategyIterable(final T value, final ShrinkStrategy<T> strategy) {
+        return () ->
+                mappingIterator(
+                        v -> new ShrinkTree<>(v, strategyIterable(v, strategy)),
+                        strategy.shrink(value));
+    }
+
     @SuppressWarnings("unused")
     public void print(Writer output) throws IOException {
         print(output, Object::toString);