1 files changed, 95 insertions, 24 deletions

diff --git a/pbt/ts/src/property.ts b/pbt/ts/src/property.ts
index 7714747..7b3ee59 100644
--- a/pbt/ts/src/property.ts
+++ b/pbt/ts/src/property.ts
@@ -1,30 +1,31 @@
-import Prando from "prando";
+import Prando from 'prando';
 
 type Result = 'OK' | 'Falsified' | 'Exception';
 
-type ShrinkResult<A> = { counterexample: A, shrinks: number };
+type ShrinkResult<A> = {counterexample: A; shrinks: number};
 
 type TestResult<A> = {
-  seed: number | null,
-  result: Result,
-  counterexample: ShrinkResult<A> | null
+  seed: number | null;
+  result: Result;
+  counterexample: ShrinkResult<A> | null;
 };
 
 type Predicate<A> = (a: A) => boolean;
 
-export type Gen<A> = (rng: Prando) => ((size: number) => A);
+export type Gen<A> = (rng: Prando) => (size: number) => A;
 
 type Shrinker<A> = (a: A) => A[];
 
-type Property<A> = (seed: number,
-  size: number) => TestResult<A>;
+type Property<A> = (seed: number, size: number) => TestResult<A>;
 
 const MAX_SUCCESS = 100;
 
-function findMinimalCounterExample<A>(x: A,
+function findMinimalCounterExample<A>(
+  x: A,
   predicate: Predicate<A>,
   shrinker: Shrinker<A>,
-  depth: number = 0): ShrinkResult<A> {
+  depth: number = 0,
+): ShrinkResult<A> {
   let xs = shrinker(x);
   let counterexample = x;
   let shrinks = depth;
@@ -37,33 +38,35 @@ function findMinimalCounterExample<A>(x: A,
       }
     }
   }
-  return { counterexample, shrinks };
+  return {counterexample, shrinks};
 }
 
-
 export function generate<A>(gen: Gen<A>): A {
   let rng = new Prando(Math.random() * 1000);
   return gen(rng)(100);
 }
 
 export function property<A>(
-  predicate: Predicate<A>,
   generator: Gen<A>,
-  shrinker: Shrinker<A>): Property<A> {
-  return (seed: number,
-    size: number) => {
+  predicate: Predicate<A>,
+  shrinker?: Shrinker<A>,
+): Property<A> {
+  return (seed: number, size: number) => {
     let rng = new Prando(seed);
     let gen = generator(rng);
     let i = 0;
     for (; i < MAX_SUCCESS; i++) {
       let x = gen(size);
       if (!predicate(x)) {
-        let counterexample = findMinimalCounterExample(x, predicate, shrinker);
-        return { seed, result: 'Falsified', tests: i + 1, counterexample };
+        let counterexample = {counterexample: x, shrinks: 0};
+        if (shrinker !== undefined) {
+          counterexample = findMinimalCounterExample(x, predicate, shrinker);
+        }
+        return {seed, result: 'Falsified', tests: i + 1, counterexample};
       }
       size++;
     }
-    return { seed, result: 'OK', tests: i, counterexample: null };
+    return {seed, result: 'OK', tests: i, counterexample: null};
   };
 }
 
@@ -82,6 +85,65 @@ export function listOf<A>(gen: Gen<A>): Gen<A[]> {
   };
 }
 
+/// generates a map with exactly size elements
+export function mapOf<K, V>(genK: Gen<K>, genV: Gen<V>): Gen<Map<K, V>> {
+  return (rng: Prando) => {
+    let gk = genK(rng);
+    let gv = genV(rng);
+    return (size: number) => {
+      let result = new Map<K, V>();
+      for (let i = 0; i < size; i++) {
+        result.set(gk(size), gv(size));
+      }
+      return result;
+    };
+  };
+}
+
+/// generates a map with at least size+1 elements. As size is expected
+/// to be positive or null this guarantees at least one element hence
+/// the name
+export function mapOf1<K, V>(genK: Gen<K>, genV: Gen<V>): Gen<Map<K, V>> {
+  return (rng: Prando) => {
+    let gk = genK(rng);
+    let gv = genV(rng);
+    return (size: number) => {
+      let result = new Map<K, V>();
+      for (let i = 0; i <= size; i++) {
+        result.set(gk(size), gv(size));
+      }
+      return result;
+    };
+  };
+}
+
+export function elements<A>(arr: A[]): Gen<A> {
+  return (rng: Prando) => (_size: number) => {
+    let index = rng.nextInt(0, arr.length - 1);
+    return arr[index];
+  };
+}
+
+export function keyValueOf<K, V>(map: Map<K, V>): Gen<[K, V]> {
+  let keys = Array.from(map.keys());
+  return (rng: Prando) => {
+    let gkey = elements(keys)(rng);
+    return (size: number) => {
+      let key = gkey(size);
+      let value = map.get(key)!;
+      return [key, value];
+    };
+  };
+}
+
+export function choose(min: number, max: number): Gen<number> {
+  return (rng: Prando) => (size: number) => rng.nextInt(min, max);
+}
+
+export let someInt: Gen<number> = (rng: Prando) => (size: number) =>
+  rng.nextInt(-size, size);
+
+// basic shrinkers
 export function shrinklist<A>(arr: A[]): A[][] {
   let result = [];
 
@@ -100,13 +162,22 @@ export function shrinklist<A>(arr: A[]): A[][] {
   return result;
 }
 
-export let someInt: Gen<number> = (rng: Prando) =>
-  (size: number) =>
-    rng.nextInt(-size, size);
-
 export const Property = {
   run<A>(prop: Property<A>): TestResult<A> {
     let seed = Math.floor(Math.random() * 1000000);
     return prop(seed, 0);
-  }
+  },
 };
+
+// helpers
+
+export function forall<A, B>(
+  gen: Gen<A>,
+  k: (a: A) => Property<B>,
+): Property<B> {
+  return (seed: number, size: number) => {
+    let rng = new Prando(seed);
+    let a = gen(rng)(size);
+    return k(a)(seed, size);
+  };
+}