Add @DataSource, allow printing of RoseTrees, other small changes to generators

1 files changed, 19 insertions, 6 deletions

diff --git a/src/main/java/au/id/zancanaro/Generators.java b/src/main/java/au/id/zancanaro/Generators.java
index e0927ff..01ab64a 100644
--- a/src/main/java/au/id/zancanaro/Generators.java
+++ b/src/main/java/au/id/zancanaro/Generators.java
@@ -1,11 +1,17 @@
 package au.id.zancanaro;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class Generators {
+    public static <T> Generator<T> sized(Function<Integer, Generator<T>> makeGenerator) {
+        return (random, size) -> makeGenerator.apply(size).generate(random, size);
+    }
+
     public static <T> Generator<T> suchThat(Generator<T> gen, Predicate<T> pred) {
         return (random, size) -> {
             RoseTree<T> result = gen.generate(random, size);
@@ -22,6 +28,13 @@ public class Generators {
         return integer(0, gens.length).flatMap(index -> gens[index]);
     }
 
+    public static Generator<Boolean> bool() {
+        return (random, size) ->
+                random.nextBoolean() ?
+                        new RoseTree<>(true, Collections.singletonList(new RoseTree<>(false, Collections.emptyList()))) :
+                        new RoseTree<>(false, Collections.emptyList());
+    }
+
     public static Generator<Integer> integer(int lower, int upper) {
         return (random, size) -> {
             int value = lower + random.nextInt(upper - lower);
@@ -31,11 +44,11 @@ public class Generators {
     }
 
     public static Generator<Integer> integer() {
-        return (random, size) -> integer(-size, size).generate(random, size);
+        return sized(size -> integer(-size, size));
     }
 
     public static Generator<Integer> natural() {
-        return (random, size) -> integer(0, size).generate(random, size);
+        return sized(size -> integer(0, size));
     }
 
     private static Iterable<RoseTree<Integer>> intShrinkingIterable(final int value, final int bound) {
@@ -57,11 +70,11 @@ public class Generators {
     }
 
     public static <T> Generator<List<T>> listOf(Generator<T> gen) {
-        return (random, size) -> integer(0, size).flatMap(count -> {
+        return (Generator<List<T>>) sized(size -> {
             @SuppressWarnings("unchecked")
-            Generator<T>[] gens = (Generator<T>[]) new Generator[count];
+            Generator<T>[] gens = (Generator<T>[]) new Generator[size];
             Arrays.fill(gens, gen);
-            return Generator.tuple(gens).fmap(Arrays::asList);
-        }).generate(random, size);
+            return Generator.tuple(gens).fmap(Arrays::asList).fmap(Collections::unmodifiableList);
+        });
     }
 }