Fix up list shrinking, remove minor superfluous things

5 files changed, 51 insertions, 15 deletions

diff --git a/src/main/java/au/id/zancanaro/javacheck/Generator.java b/src/main/java/au/id/zancanaro/javacheck/Generator.java
index 66ebf1b..55d06f8 100644
--- a/src/main/java/au/id/zancanaro/javacheck/Generator.java
+++ b/src/main/java/au/id/zancanaro/javacheck/Generator.java
@@ -1,5 +1,6 @@
 package au.id.zancanaro.javacheck;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Random;
 import java.util.function.Function;
@@ -24,6 +25,17 @@ public interface Generator<T> {
         };
     }
 
+    static <T> Generator<List<T>> list(int count, Generator<T> gen) {
+        return (random, size) -> {
+            @SuppressWarnings("unchecked")
+            RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[count];
+            for (int i = 0; i < count; ++i) {
+                result[i] = gen.generate(random, size);
+            }
+            return RoseTree.shrink(Function.identity(), result);
+        };
+    }
+
     default <R> Generator<R> fmap(Function<T, R> f) {
         return (random, size) -> this.generate(random, size).fmap(f);
     }
diff --git a/src/main/java/au/id/zancanaro/javacheck/Generators.java b/src/main/java/au/id/zancanaro/javacheck/Generators.java
index ec6a329..9ced818 100644
--- a/src/main/java/au/id/zancanaro/javacheck/Generators.java
+++ b/src/main/java/au/id/zancanaro/javacheck/Generators.java
@@ -81,13 +81,11 @@ public final class Generators {
     }
 
     public static <T> Generator<List<T>> listOf(Generator<T> gen) {
-        return sized(size ->
-            integer(0, size).flatMap(count -> {
-                @SuppressWarnings("unchecked")
-                Generator<T>[] gens = (Generator<T>[]) new Generator[count];
-                Arrays.fill(gens, gen);
-                return Generator.tuple(gens);
-            })).fmap(Collections::unmodifiableList);
+        return (random, size) -> {
+            Generator<Integer> countGen = sized(s -> integer(0, s));
+            int count = countGen.generate(random, size).getValue();
+            return Generator.list(count, gen).generate(random, size);
+        };
     }
 
     public static Generator<Character> character() {
@@ -124,12 +122,14 @@ public final class Generators {
     }
 
     public static Generator<String> stringOf(Generator<Character> charGen) {
-        return Generators.listOf(charGen).fmap(list -> {
-            StringBuilder builder = new StringBuilder(list.size());
-            for (Object c : list) {
-                builder.append(c);
+        return listOf(charGen).fmap(list -> {
+            char[] chars = new char[list.size()];
+            int i = 0;
+            for (Character c : list) {
+                chars[i++] = c;
             }
-            return builder.toString();
+            return String.valueOf(chars);
+//            return new String(chars);
         });
     }
 }
diff --git a/src/main/java/au/id/zancanaro/javacheck/RoseTree.java b/src/main/java/au/id/zancanaro/javacheck/RoseTree.java
index d9fb508..4d7b770 100644
--- a/src/main/java/au/id/zancanaro/javacheck/RoseTree.java
+++ b/src/main/java/au/id/zancanaro/javacheck/RoseTree.java
@@ -38,7 +38,7 @@ public class RoseTree<T> {
                         tree.getValue().children.iterator()));
     }
 
-    public static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
+    private static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
         return Iterators.flatten(Iterators.rangeIterator(trees.length, index ->
                         Iterators.mappingIterator(child -> {
                             @SuppressWarnings("unchecked")
@@ -64,6 +64,32 @@ public class RoseTree<T> {
                         RoseTree.permutations(trees)));
     }
 
+    private static <T> Iterator<RoseTree<T>[]> removeEach(RoseTree<T>[] trees) {
+        return Iterators.concat(
+                Iterators.rangeIterator(trees.length, index -> {
+                    @SuppressWarnings("unchecked")
+                    RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[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, R> RoseTree<R> shrink(Function<List<T>, R> fn, RoseTree<T>[] trees) {
+        @SuppressWarnings("unchecked")
+        List<T> heads = new ArrayList(trees.length);
+        for (RoseTree<T> tree : trees) {
+            heads.add(tree.getValue());
+        }
+        return new RoseTree<>(
+                fn.apply(heads),
+                () -> Iterators.mappingIterator(
+                        roses -> RoseTree.shrink(fn, roses),
+                        RoseTree.removeEach(trees)));
+    }
+
     public <R> RoseTree<R> fmap(Function<T, R> f) {
         return new RoseTree<>(
                 f.apply(this.value),
diff --git a/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java b/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java
index 271e763..5ec0c36 100644
--- a/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java
+++ b/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java
@@ -190,7 +190,6 @@ public class Properties extends BlockJUnit4ClassRunner {
                 Generator<?>[] generators = (Generator<?>[]) new Generator[method.getParameterCount()];
                 int index = 0;
                 for (Type type : method.getGenericParameterTypes()) {
-                    // TODO: validate ahead of time that this generator will exist (ideally in the constructor validation time)
                     generators[index++] = this.generators.get(type);
                 }
                 @SuppressWarnings("unchecked")
diff --git a/src/test/java/au/id/zancanaro/javacheck/junit/SimpleListOperationsTest.java b/src/test/java/au/id/zancanaro/javacheck/junit/SimpleListOperationsTest.java
index 8a298be..089451e 100644
--- a/src/test/java/au/id/zancanaro/javacheck/junit/SimpleListOperationsTest.java
+++ b/src/test/java/au/id/zancanaro/javacheck/junit/SimpleListOperationsTest.java
@@ -3,7 +3,6 @@ package au.id.zancanaro.javacheck.junit;
 import au.id.zancanaro.javacheck.Generator;
 import au.id.zancanaro.javacheck.annotations.DataSource;
 import au.id.zancanaro.javacheck.annotations.Property;
-import au.id.zancanaro.javacheck.junit.Properties;
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;