From 4af0a27a9fbb89e81e760028c486be2c6df0ec45 Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Fri, 1 Jun 2018 15:44:51 +0200
Subject: [PATCH] jetCheck: add length-aware generator for strings

---
 .../src/org/jetbrains/jetCheck/Generator.java | 26 ++++++++++++++-----
 .../org/jetbrains/jetCheck/GeneratorTest.java |  5 ++++
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/jetCheck/src/org/jetbrains/jetCheck/Generator.java b/jetCheck/src/org/jetbrains/jetCheck/Generator.java
index 6d6e1eefd6c9..3cbfa46425fb 100644
--- a/jetCheck/src/org/jetbrains/jetCheck/Generator.java
+++ b/jetCheck/src/org/jetbrains/jetCheck/Generator.java
@@ -181,6 +181,11 @@ public class Generator<T> {
     return integers(min, max).map(i -> (char)i.intValue());
   }
 
+  /** Generates characters that occur in the given string */
+  public static Generator<Character> charsFrom(String possibleChars) {
+    return sampledFrom(IntStream.range(0, possibleChars.length()).mapToObj(possibleChars::charAt).collect(Collectors.toList()));
+  }
+
   /** Generates ASCII characters excluding the system ones (lower than 32) */
   public static Generator<Character> asciiPrintableChars() {
     return charsInRange((char)32, (char)126);
@@ -210,17 +215,24 @@ public class Generator<T> {
 
   /** Generates (possibly empty) random strings consisting of the given characters (provided as a string) */
   public static Generator<String> stringsOf(@NotNull String possibleChars) {
-    List<Character> chars = IntStream.range(0, possibleChars.length()).mapToObj(possibleChars::charAt).collect(Collectors.toList());
-    return stringsOf(sampledFrom(chars));
+    return stringsOf(charsFrom(possibleChars));
   }
 
   /** Generates (possibly empty) random strings consisting of characters provided by the given generator */
   public static Generator<String> stringsOf(@NotNull Generator<Character> charGen) {
-    return listsOf(charGen).map(chars -> {
-      StringBuilder sb = new StringBuilder();
-      chars.forEach(sb::append);
-      return sb.toString();
-    });
+    return listsOf(charGen).map(Generator::charsToString);
+  }
+
+  /** Generates (possibly empty) random strings of length in the given distribution, consisting of characters provided by the given generator */
+  public static Generator<String> stringsOf(@NotNull IntDistribution length, @NotNull Generator<Character> charGen) {
+    return listsOf(length, charGen).map(Generator::charsToString);
+  }
+
+  @NotNull
+  private static String charsToString(List<Character> chars) {
+    StringBuilder sb = new StringBuilder();
+    chars.forEach(sb::append);
+    return sb.toString();
   }
 
   /** Generates random strings consisting ASCII letters and digit and starting with a letter */
diff --git a/jetCheck/test/org/jetbrains/jetCheck/GeneratorTest.java b/jetCheck/test/org/jetbrains/jetCheck/GeneratorTest.java
index f14f20d1e319..03a408caf28e 100644
--- a/jetCheck/test/org/jetbrains/jetCheck/GeneratorTest.java
+++ b/jetCheck/test/org/jetbrains/jetCheck/GeneratorTest.java
@@ -36,6 +36,11 @@ public class GeneratorTest extends PropertyCheckerTestCase {
     assertEquals("a", checkGeneratesExample(stringsOf(asciiPrintableChars()),
                                             s -> s.contains("a"),
                                             10));
+
+    String aWithB = checkGeneratesExample(stringsOf(IntDistribution.uniform(2, 100), asciiPrintableChars()),
+                                     s -> s.contains("a") && s.contains("b"),
+                                     27);
+    assertTrue(aWithB, "ab".equals(aWithB) || "ba".equals(aWithB));
   }
 
   public void testLetterStringContains() {