3 files changed, 45 insertions, 21 deletions

diff --git a/pbt/ts/src/bank.test.ts b/pbt/ts/src/bank.test.ts
index c448242..bbb6fcf 100644
--- a/pbt/ts/src/bank.test.ts
+++ b/pbt/ts/src/bank.test.ts
@@ -60,9 +60,10 @@ const genOperations = (accounts: Map<Account, number>) =>
   listOf(genOneOperation(accounts));
 
 const correctBalance =
-  (api: SimpleBank, initialAccounts: Map<Account, number>) =>
+  (initialAccounts: Map<Account, number>) =>
   (ops: Operation[]): boolean => {
     const expectedBalances = new Map(initialAccounts);
+    const api = new SimpleBank(initialAccounts);
 
     for (const op of ops) {
       switch (op.op) {
@@ -70,6 +71,9 @@ const correctBalance =
           const retrieved = toNullable(api.balance(op.account));
           const expected = expectedBalances.get(op.account);
           if (retrieved !== expected) {
+            console.log(
+              `Balance mismatch for account ${op.account}: expected ${expected}, got ${retrieved}`,
+            );
             return false;
           }
           break;
@@ -93,6 +97,9 @@ const correctBalance =
     for (const [account, expectedBal] of expectedBalances.entries()) {
       const retrieved = toNullable(api.balance(account));
       if (retrieved !== expectedBal) {
+        console.log(
+          `Final balance mismatch for account ${account}: expected ${expectedBal}, got ${retrieved}`,
+        );
         return false;
       }
     }
@@ -116,16 +123,12 @@ describe('Bank Simulator', () => {
   });
 
   test('correctly update balance given a sequence of deposits and withdrawals', () => {
-    const genTransactionAmount: Gen<number> = choose(1, 100000);
-
     const prop = forall(genAccounts, accounts => {
-      const api = new SimpleBank(accounts);
-      return property(genOperations(accounts), correctBalance(api, accounts));
+      return property(genOperations(accounts), correctBalance(accounts));
     });
 
     const result = Property.run(prop);
 
-    console.log(result);
     expect(result.result).toBe('OK');
   });
 });
diff --git a/pbt/ts/src/bank.ts b/pbt/ts/src/bank.ts
index e5619f4..e388354 100644
--- a/pbt/ts/src/bank.ts
+++ b/pbt/ts/src/bank.ts
@@ -1,4 +1,4 @@
-import {Option, fromNullable} from 'fp-ts/lib/Option';
+import {Option, fromNullable, isSome, matchW} from 'fp-ts/lib/Option';
 
 // keep it simple
 export type Account = string;
@@ -32,11 +32,27 @@ export class SimpleBank implements Bank {
   }
 
   deposit(account: Account, amt: number): void {
-    throw new Error('Method not implemented.');
+    const current = this.balance(account);
+    matchW(
+      () => {
+        throw new Error(`Unknown account ${account}`);
+      },
+      (bal: number) => {
+        this.accounts.set(account, bal + amt);
+      },
+    )(current);
   }
 
   withdraw(account: Account, amt: number): void {
-    throw new Error('Method not implemented.');
+    const current = this.balance(account);
+    matchW(
+      () => {
+        throw new Error(`Unknown account ${account}`);
+      },
+      (bal: number) => {
+        this.accounts.set(account, bal - amt);
+      },
+    )(current);
   }
 }
 
diff --git a/pbt/ts/src/property.ts b/pbt/ts/src/property.ts
index 04f2a62..da05328 100644
--- a/pbt/ts/src/property.ts
+++ b/pbt/ts/src/property.ts
@@ -131,21 +131,26 @@ export function frequency<A>(distribution: [number, Gen<A>][]): Gen<A> {
   const total: number = distribution.reduce((acc, [freq, _]) => acc + freq, 0);
 
   return (rng: Prando) => (size: number) => {
-    const idx = rng.next(0, total);
-    let selectedGen: Option<Gen<A>> = none;
-    selectedGen = distribution.reduce((found, [freq, gen]) => {
-      if (isSome(found)) {
-        return found;
-      } else if (idx <= freq) {
-        return some(gen);
-      } else {
-        return none;
-      }
-    }, selectedGen);
+    const idx = Math.floor(rng.next(0, total));
+    const [selectedGen, _] = distribution.reduce(
+      ([found, sum], [freq, gen]) => {
+        const current = sum + freq;
+        if (isSome(found)) {
+          return [found, current];
+        } else if (idx <= current) {
+          return [some(gen), current];
+        } else {
+          return [none, current];
+        }
+      },
+      [none as Option<Gen<A>>, 0],
+    );
     const generator: Gen<A> = pipe(
       selectedGen,
       getOrElse((() => {
-        throw new Error(`invalid frequency distribution ${distribution}`);
+        throw new Error(
+          `invalid frequency distribution ${distribution}, idx=${idx}`,
+        );
       }) as LazyArg<Gen<A>>),
     );
     return generator(rng)(size);