1 files changed, 54 insertions, 8 deletions

diff --git a/pbt/ts/src/gen.ts b/pbt/ts/src/gen.ts
index c27cfec..7d8c6f1 100644
--- a/pbt/ts/src/gen.ts
+++ b/pbt/ts/src/gen.ts
@@ -27,7 +27,7 @@ export class BytesGen implements GenBytes {
   }
 
   generate(n: number): number[] {
-    if (n <= 0) {
+    if (n < 0) {
       throw new Error('n must be greater than 0');
     }
     return this.generateBytes(n);
@@ -42,6 +42,13 @@ export class BytesGen implements GenBytes {
     return gens;
   }
 
+  reject() {
+    if (this.bytes.length === 0) {
+      throw new Error('No more bytes to shrink');
+    }
+    this.bytes.pop();
+  }
+
   replay(): GenBytes {
     return new ReplayBytes(this.bytes);
   }
@@ -106,12 +113,27 @@ class GenIntWithin implements Gen<number> {
   generate(): number {
     const range = this.max - this.min;
     const numBytes = Math.ceil(Math.log2(range) / 8);
-    const bytes = this.bytes.generate(numBytes);
-    let value = 0;
-    for (let i = 0; i < bytes.length; i++) {
-      value = (value << 8) | bytes[i];
+    const maxBytes = Math.pow(2, numBytes * 8);
+    const upperBound = maxBytes - (maxBytes % range);
+    let bytes = this.bytes.generate(numBytes);
+    let rejection = 0;
+    while (true) {
+      let value = 0;
+      for (let i = 0; i < bytes.length; i++) {
+        value = (value << 8) | bytes[i];
+      }
+
+      if (value >= upperBound) {
+        this.bytes.reject();
+        bytes = this.bytes.generate(numBytes);
+        continue;
+      }
+
+      if (rejection > 0) {
+        console.log(`Rejection sampling: rejected ${rejection} values before accepting ${value}`);
+      }
+      return this.min + (value % range);
     }
-    return this.min + (value % range);
   }
 
   generateN(n: number): number[] {
@@ -125,6 +147,30 @@ class GenIntWithin implements Gen<number> {
 
 }
 
-export const genInt = (seed : number | undefined = undefined) => new GenInt(seed);
+class GenASCIIString implements Gen<string> {
+  private bytes: BytesGen;
+
+  constructor(seed: number | undefined = undefined) {
+    this.bytes = new BytesGen(new Prando(seed ? seed : Math.random() * 1000));
+  }
+
+  generate(): string {
+    const length = this.bytes.generate(1)[0]; // limit to 20 characters
+    const bytes = this.bytes.generate(length);
+    return String.fromCharCode(...bytes.map(b => b % 128)); // limit to ASCII
+  }
+
+  generateN(n: number): string[] {
+    const result = [];
+    for (let i = 0; i < n; i++) {
+      result.push(this.generate());
+    }
+    return result;
+  }
+}
+
+export const genInt = (seed: number | undefined = undefined) => new GenInt(seed);
+
+export const genIntWithin = (min: number, max: number, seed: number | undefined = undefined) => new GenIntWithin(min, max, seed);
 
-export const genIntWithin = (min: number, max: number, seed : number | undefined = undefined) => new GenIntWithin(min, max, seed);
+export const genASCIIString = (seed: number | undefined = undefined) => new GenASCIIString(seed);