Support multiple parties (flat)

src/examples/simple_zkapp.ts

@@ -14,6 +14,7 @@ import {
   UInt32,
   Bool,
   PublicKey,
+  Circuit,
 } from 'snarkyjs';
 
 await isReady;
@@ -32,6 +33,7 @@ class SimpleZkapp extends SmartContract {
     this.setPermissions({
       ...Permissions.default(),
       editState: Permissions.proofOrSignature(),
+      send: Permissions.proofOrSignature(),
     });
     this.balance.addInPlace(UInt64.fromNumber(initialBalance));
     this.x.set(initialState);
@@ -61,7 +63,10 @@ class SimpleZkapp extends SmartContract {
     // pay out half of the zkapp balance to the caller
     let balance = this.account.balance.get();
     this.account.balance.assertEquals(balance);
-    let halfBalance = balance.div(2);
+    // FIXME UInt64.div() doesn't work on variables
+    let halfBalance = Circuit.witness(UInt64, () =>
+      balance.toConstant().div(2)
+    );
     this.balance.subInPlace(halfBalance);
     callerParty.balance.addInPlace(halfBalance);
 
@@ -71,6 +76,8 @@ class SimpleZkapp extends SmartContract {
   }
 }
 
+const doProofs = true;
+
 let Local = Mina.LocalBlockchain();
 Mina.setActiveInstance(Local);
 
@@ -89,6 +96,11 @@ let initialBalance = 10_000_000_000;
 let initialState = Field(1);
 let zkapp = new SimpleZkapp(zkappAddress);
 
+if (doProofs) {
+  console.log('compile');
+  await SimpleZkapp.compile(zkappAddress);
+}
+
 console.log('deploy');
 let tx = await Mina.transaction(feePayer, () => {
   Party.fundNewAccount(feePayer, { initialBalance });
@@ -102,28 +114,29 @@ console.log(`initial balance: ${zkapp.account.balance.get().div(1e9)} MINA`);
 console.log('update');
 tx = await Mina.transaction(feePayer, () => {
   zkapp.update(Field(3));
-  zkapp.sign(zkappKey);
+  if (!doProofs) zkapp.sign(zkappKey);
 });
+if (doProofs) await tx.prove();
 tx.send();
 
 console.log('payout');
 tx = await Mina.transaction(feePayer, () => {
   zkapp.payout(privilegedKey);
-  zkapp.sign(zkappKey);
+  if (!doProofs) zkapp.sign(zkappKey);
 });
+if (doProofs) await tx.prove();
 tx.send();
 
-console.log(tx.toJSON());
-
 console.log('final state: ' + zkapp.x.get());
 console.log(`final balance: ${zkapp.account.balance.get().div(1e9)} MINA`);
 
 console.log('try to payout a second time..');
 tx = await Mina.transaction(feePayer, () => {
   zkapp.payout(privilegedKey);
-  zkapp.sign(zkappKey);
+  if (!doProofs) zkapp.sign(zkappKey);
 });
 try {
+  if (doProofs) await tx.prove();
   tx.send();
 } catch (err: any) {
   console.log('Transaction failed with error', err.message);
@@ -133,8 +146,9 @@ console.log('try to payout to a different account..');
 try {
   tx = await Mina.transaction(feePayer, () => {
     zkapp.payout(Local.testAccounts[2].privateKey);
-    zkapp.sign(zkappKey);
+    if (!doProofs) zkapp.sign(zkappKey);
   });
+  if (doProofs) await tx.prove();
   tx.send();
 } catch (err: any) {
   console.log('Transaction failed with error', err.message);

src/index.ts

@@ -37,14 +37,12 @@ export {
 } from './lib/zkapp';
 export { state, State, declareState } from './lib/state';
 export { Proof, SelfProof, ZkProgram, verify } from './lib/proof_system';
-export * from './lib/party';
+export { Party, Permissions, ZkappPublicInput } from './lib/party';
 export {
   fetchAccount,
   fetchLastBlock,
-  parseFetchedAccount,
   addCachedAccount,
   setGraphqlEndpoint,
-  sendZkappQuery,
   sendZkapp,
 } from './lib/fetch';
 export * as Encryption from './lib/encryption';

src/lib/circuit_value.ts

@@ -1,6 +1,6 @@
 import 'reflect-metadata';
 import { Circuit, Field, Bool, JSONValue, AsFieldElements } from '../snarky';
-import { withContext } from './global-context';
+import { snarkContext } from './proof_system';
 
 // external API
 export {
@@ -69,6 +69,10 @@ abstract class CircuitValue {
     return (this.constructor as any).toJSON(this);
   }
 
+  toConstant(): this {
+    return (this.constructor as any).toConstant(this);
+  }
+
   equals(x: this) {
     return Circuit.equal(this, x);
   }
@@ -301,14 +305,17 @@ function circuitMain(
   }
 
   target.snarkyMain = (w: Array<any>, pub: Array<any>) => {
-    let [, result] = withContext({ inCheckedComputation: true }, () => {
-      let args = [];
-      for (let i = 0; i < numArgs; ++i) {
-        args.push((publicIndexSet.has(i) ? pub : w).shift());
-      }
+    let [, result] = snarkContext.runWith(
+      { inCheckedComputation: true },
+      () => {
+        let args = [];
+        for (let i = 0; i < numArgs; ++i) {
+          args.push((publicIndexSet.has(i) ? pub : w).shift());
+        }
 
-      return target[propertyName].apply(target, args);
-    });
+        return target[propertyName].apply(target, args);
+      }
+    );
     return result;
   };
 
@@ -511,12 +518,15 @@ Circuit.switch = function <T, A extends AsFieldElements<T>>(
   return type.ofFields(fields);
 };
 
-Circuit.constraintSystem = function (f) {
-  let [, result] = withContext(
+Circuit.constraintSystem = function <T>(f: () => T) {
+  let [, result] = snarkContext.runWith(
     { inAnalyze: true, inCheckedComputation: true },
     () => {
-      let { rows, digest, json } = (Circuit as any)._constraintSystem(f);
-      return { rows, digest };
+      let result: T;
+      let { rows, digest, json } = (Circuit as any)._constraintSystem(() => {
+        result = f();
+      });
+      return { rows, digest, result: result! };
     }
   );
   return result;

src/lib/global-context.ts

@@ -1,110 +1,112 @@
 import { Party } from './party';
 
-export {
-  mainContext,
-  withContext,
-  withContextAsync,
-  getContext,
-  inProver,
-  inCompile,
-  inCheckedComputation,
-  inAnalyze,
-};
+export { Context };
 
-// context for compiling / proving
-// TODO reconcile mainContext with currentTransaction
-type MainContext = {
-  witnesses?: unknown[];
-  self?: Party;
-  expectedAccesses: number | undefined;
-  actualAccesses: number;
-  inProver?: boolean;
-  inCompile?: boolean;
-  inCheckedComputation?: boolean;
-  inAnalyze?: boolean;
-  otherContext?: any;
-};
+namespace Context {
+  export type id = number;
 
-let mainContext = undefined as MainContext | undefined;
-type PartialContext = Partial<MainContext>;
+  export type t<Context> = (() => Context | undefined) & {
+    data: { context: Context; id: id }[];
+    allowsNesting: boolean;
 
-function withContext<T>(
-  {
-    witnesses = undefined,
-    expectedAccesses = undefined,
-    actualAccesses = 0,
-    self,
-    ...other
-  }: PartialContext,
-  f: () => T
-) {
-  let prevContext = mainContext;
-  mainContext = { witnesses, expectedAccesses, actualAccesses, self, ...other };
-  let result: T;
-  let resultContext: MainContext;
+    get(): Context;
+    has(): boolean;
+    runWith<Result>(context: Context, func: () => Result): [Context, Result];
+    runWithAsync<Result>(
+      context: Context,
+      func: () => Promise<Result>
+    ): Promise<[Context, Result]>;
+    enter(context: Context): id;
+    leave(id: id): Context;
+    id: () => id;
+  };
+}
+const Context = { create };
+
+function create<C>(
+  options = {
+    allowsNesting: true,
+    default: undefined,
+  } as { allowsNesting?: boolean; default?: C }
+): Context.t<C> {
+  let t: Context.t<C> = Object.assign(
+    function (): C | undefined {
+      return t.data[t.data.length - 1].context;
+    },
+    {
+      data: [],
+      allowsNesting: options.allowsNesting ?? true,
+      get: () => get(t),
+      has: () => t.data.length !== 0,
+      runWith: <R>(context: C, func: () => R) => runWith(t, context, func),
+      runWithAsync: <R>(context: C, func: () => Promise<R>) =>
+        runWithAsync(t, context, func),
+      enter: (context: C) => enter(t, context),
+      leave: (id: Context.id) => leave(t, id),
+      id: () => {
+        if (t.data.length === 0) throw Error(contextConflictMessage);
+        return t.data[t.data.length - 1].id;
+      },
+    }
+  );
+  if (options.default !== undefined) enter(t, options.default);
+  return t;
+}
+
+function enter<C>(t: Context.t<C>, context: C): Context.id {
+  if (t.data.length > 0 && !t.allowsNesting) {
+    throw Error(contextConflictMessage);
+  }
+  let id = Math.random();
+  t.data.push({ context, id });
+  return id;
+}
+
+function leave<C>(t: Context.t<C>, id: Context.id): C {
+  let current = t.data.pop();
+  if (current === undefined) throw Error(contextConflictMessage);
+  if (current.id !== id) throw Error(contextConflictMessage);
+  return current.context;
+}
+
+function get<C>(t: Context.t<C>): C {
+  if (t.data.length === 0) throw Error(contextConflictMessage);
+  let current = t.data[t.data.length - 1];
+  return current.context;
+}
+
+function runWith<C, Result>(
+  t: Context.t<C>,
+  context: C,
+  func: () => Result
+): [C, Result] {
+  let id = enter(t, context);
+  let result: Result;
+  let resultContext: C;
   try {
-    result = f();
+    result = func();
   } finally {
-    resultContext = mainContext;
-    mainContext = prevContext;
+    resultContext = leave(t, id);
   }
-  return [resultContext, result] as [MainContext, T];
+  return [resultContext, result];
 }
 
-// TODO: this is unsafe, the mainContext will be overridden if we invoke this function multiple times concurrently
-// at the moment, we solve this by detecting unsafe use and throwing an error
-async function withContextAsync<T>(
-  {
-    witnesses = undefined,
-    expectedAccesses = 1,
-    actualAccesses = 0,
-    self,
-    ...other
-  }: PartialContext,
-  f: () => Promise<T>
-) {
-  let prevContext = mainContext;
-  mainContext = { witnesses, expectedAccesses, actualAccesses, self, ...other };
-  let result: T;
-  let resultContext: MainContext;
+async function runWithAsync<C, Result>(
+  t: Context.t<C>,
+  context: C,
+  func: () => Promise<Result>
+): Promise<[C, Result]> {
+  let id = enter(t, context);
+  let result: Result;
+  let resultContext: C;
   try {
-    result = await f();
-    if (mainContext.actualAccesses !== mainContext.expectedAccesses)
-      throw Error(contextConflictMessage);
+    result = await func();
   } finally {
-    resultContext = mainContext;
-    mainContext = prevContext;
+    resultContext = leave(t, id);
   }
-  return [resultContext, result] as [MainContext, T];
+  return [resultContext, result];
 }
 
 let contextConflictMessage =
   "It seems you're running multiple provers concurrently within" +
   ' the same JavaScript thread, which, at the moment, is not supported and would lead to bugs.';
-function getContext() {
-  if (mainContext === undefined) throw Error(contextConflictMessage);
-  mainContext.actualAccesses++;
-  if (
-    mainContext.expectedAccesses !== undefined &&
-    mainContext.actualAccesses > mainContext.expectedAccesses
-  )
-    throw Error(contextConflictMessage);
-  return mainContext;
-}
-
-function inProver() {
-  return !!mainContext?.inProver;
-}
-function inCompile() {
-  return !!mainContext?.inCompile;
-}
-function inCheckedComputation() {
-  return (
-    !!mainContext?.inCompile ||
-    !!mainContext?.inProver ||
-    !!mainContext?.inCheckedComputation
-  );
-}
-function inAnalyze() {
-  return !!mainContext?.inAnalyze;
-}

src/lib/hash.ts

@@ -1,5 +1,5 @@
 import { Poseidon as Poseidon_, Field } from '../snarky';
-import { inCheckedComputation } from './global-context';
+import { inCheckedComputation } from './proof_system';
 
 // external API
 export { Poseidon };