1 files changed, 49 insertions, 52 deletions
diff --git a/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java b/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
index bdc7604..4337c14 100644
--- a/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
+++ b/src/main/java/au/id/zancanaro/javacheck/ShrinkTree.java
@@ -3,20 +3,20 @@ package au.id.zancanaro.javacheck;
 import java.io.IOException;
 import java.io.Writer;
 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.*;
+import java.util.function.Supplier;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 @SuppressWarnings("unused")
 public class ShrinkTree<T> {
     private final T value;
-    private final Iterable<ShrinkTree<T>> children;
+    private final Supplier<Stream<ShrinkTree<T>>> children;
 
-    private ShrinkTree(T value, Iterable<ShrinkTree<T>> children) {
+    private ShrinkTree(T value, Supplier<Stream<ShrinkTree<T>>> children) {
         this.value = value;
         this.children = children;
     }
@@ -25,20 +25,23 @@ public class ShrinkTree<T> {
         return value;
     }
 
-    public Iterator<ShrinkTree<T>> getChildren() {
-        return children.iterator();
+    public Stream<ShrinkTree<T>> getChildren() {
+        return children.get();
     }
 
+    @SuppressWarnings("Convert2MethodRef")
     public static <T> ShrinkTree<T> pure(T value) {
-        return new ShrinkTree<>(value, Collections.emptyList());
+        // Converting the () -> Stream.empty() into Stream::empty actually
+        // breaks Java's generic type inference. Who would have thought?
+        return new ShrinkTree<>(value, () -> Stream.empty());
     }
 
     public static <T> ShrinkTree<T> join(ShrinkTree<ShrinkTree<T>> tree) {
         return new ShrinkTree<>(
                 tree.getValue().getValue(),
-                () -> concat(
-                        mappingIterator(ShrinkTree::join, tree.children.iterator()),
-                        tree.getValue().children.iterator()));
+                () -> Stream.concat(
+                        tree.getChildren().map(ShrinkTree::join),
+                        tree.getValue().getChildren()));
     }
 
     private static <T> List<T> makeHeadList(ShrinkTree<T>[] trees) {
@@ -49,48 +52,42 @@ public class ShrinkTree<T> {
         return heads;
     }
 
-    public static <T> Iterator<ShrinkTree<T>[]> promoteChildren(ShrinkTree<T>[] trees) {
-        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) {
-                                result[i] = (i == index ? child : trees[i]);
-                            }
-                            return result;
-                        }, trees[index].getChildren())));
-    }
-
-    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;
-        });
+    @SuppressWarnings("unchecked")
+    public static <T> Stream<ShrinkTree<T>[]> promoteChildren(ShrinkTree<T>[] trees) {
+        return IntStream.range(0, trees.length)
+                .mapToObj(index -> trees[index].getChildren().map(child ->
+                        IntStream.range(0, trees.length)
+                                .mapToObj(i -> (i == index ? child : trees[i]))
+                                .toArray(ShrinkTree[]::new)))
+                .flatMap(x -> x)
+                .map(x -> (ShrinkTree<T>[]) x);
     }
 
-    public static <T> Iterator<ShrinkTree<T>[]> removeAndPromoteChildren(ShrinkTree<T>[] trees) {
-        return concat(removeChildren(trees), promoteChildren(trees));
+    public static <T> Stream<ShrinkTree<T>[]> removeChildren(ShrinkTree<T>[] trees) {
+        return IntStream.range(0, trees.length)
+                .mapToObj(index -> IntStream.range(0, trees.length)
+                        .filter(i -> i != index)
+                        .mapToObj(i -> trees[i])
+                        .toArray(ShrinkTree[]::new));
+    }
+
+    public static <T> Stream<ShrinkTree<T>[]> removeAndPromoteChildren(ShrinkTree<T>[] trees) {
+        return Stream.concat(removeChildren(trees), promoteChildren(trees));
     }
 
     public static <T> ShrinkTree<List<T>> combine(
             ShrinkTree<T>[] trees,
-            Function<ShrinkTree<T>[], Iterator<ShrinkTree<T>[]>> processChildren) {
+            Function<ShrinkTree<T>[], Stream<ShrinkTree<T>[]>> processChildren) {
         return new ShrinkTree<>(
                 makeHeadList(trees),
-                () -> mappingIterator(
-                        shrinks -> ShrinkTree.combine(shrinks, processChildren),
-                        processChildren.apply(trees)));
+                () -> processChildren.apply(trees)
+                        .map(shrinks -> combine(shrinks, processChildren)));
     }
 
     public <R> ShrinkTree<R> map(Function<T, R> f) {
         return new ShrinkTree<>(
                 f.apply(this.value),
-                () -> mappingIterator(tree -> tree.map(f), this.children.iterator()));
+                () -> this.getChildren().map(tree -> tree.map(f)));
     }
 
     public <R> ShrinkTree<R> flatMap(Function<T, ShrinkTree<R>> f) {
@@ -101,24 +98,20 @@ public class ShrinkTree<T> {
         if (predicate.test(this.getValue())) {
             return new ShrinkTree<>(
                     this.getValue(),
-                    () -> mappingIterator(tree -> tree.filter(predicate),
-                            filteringIterator(
-                                    tree -> predicate.test(tree.getValue()),
-                                    this.getChildren())));
+                    () -> this.getChildren()
+                            .filter(tree -> predicate.test(tree.getValue()))
+                            .map(tree -> tree.filter(predicate)));
         } else {
             throw new IllegalArgumentException("Current value doesn't match predicate: whoops!");
         }
     }
 
     public ShrinkTree<T> withShrinkStrategy(ShrinkStrategy<T> strategy) {
-        return new ShrinkTree<>(this.getValue(), strategyIterable(this.getValue(), strategy));
+        return new ShrinkTree<>(this.getValue(), () -> strategyStream(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));
+    private static <T> Stream<ShrinkTree<T>> strategyStream(final T value, final ShrinkStrategy<T> strategy) {
+        return strategy.shrink(value).map(v -> new ShrinkTree<>(v, () -> strategyStream(v, strategy)));
     }
 
     public void print(Writer output) throws IOException {
@@ -128,8 +121,12 @@ public class ShrinkTree<T> {
     public void print(Writer output, Function<T, String> toString) throws IOException {
         output.write(toString.apply(this.getValue()));
         output.write('[');
-        for (ShrinkTree<T> child : children) {
-            child.print(output, toString);
+        Iterator<ShrinkTree<T>> iterator = children.get().iterator();
+        while (iterator.hasNext()) {
+            iterator.next().print(output, toString);
+            if (iterator.hasNext()) {
+                output.write(' ');
+            }
         }
         output.write(']');
         output.flush();