Merge pull request #297 from o1-labs/feature/composability-rebased

zkApp composability

src/examples/zkapps/composability.ts

@@ -0,0 +1,118 @@
+/**
+ * zkApps composability
+ */
+import {
+  Field,
+  isReady,
+  method,
+  Mina,
+  Party,
+  Permissions,
+  PrivateKey,
+  SmartContract,
+  state,
+  State,
+} from 'snarkyjs';
+
+const doProofs = true;
+
+await isReady;
+
+// contract which can add 1 to a number
+class Incrementer extends SmartContract {
+  @method increment(x: Field): Field {
+    return x.add(1);
+  }
+}
+
+// contract which can add two numbers, plus 1, and return the result
+// incrementing by one is outsourced to another contract (it's cleaner that way, we want to stick to the single responsibility principle)
+class Adder extends SmartContract {
+  @method addPlus1(x: Field, y: Field): Field {
+    // compute result
+    let sum = x.add(y);
+    // call the other contract to increment
+    let incrementer = new Incrementer(incrementerAddress);
+    return incrementer.increment(sum);
+  }
+}
+
+// contract which calls the Adder, stores the result on chain & emits an event
+class Caller extends SmartContract {
+  @state(Field) sum = State<Field>();
+  events = { sum: Field };
+
+  @method callAddAndEmit(x: Field, y: Field) {
+    let adder = new Adder(adderAddress);
+    let sum = adder.addPlus1(x, y);
+    this.emitEvent('sum', sum);
+    this.sum.set(sum);
+  }
+}
+
+// script to deploy zkapps and do interactions
+
+let Local = Mina.LocalBlockchain();
+Mina.setActiveInstance(Local);
+
+// a test account that pays all the fees, and puts additional funds into the zkapp
+let feePayer = Local.testAccounts[0].privateKey;
+
+// the first contract's address
+let incrementerKey = PrivateKey.random();
+let incrementerAddress = incrementerKey.toPublicKey();
+// the second contract's address
+let adderKey = PrivateKey.random();
+let adderAddress = adderKey.toPublicKey();
+// the third contract's address
+let zkappKey = PrivateKey.random();
+let zkappAddress = zkappKey.toPublicKey();
+
+let zkapp = new Caller(zkappAddress);
+let adderZkapp = new Adder(adderAddress);
+let incrementerZkapp = new Incrementer(incrementerAddress);
+
+if (doProofs) {
+  console.log('compile (incrementer)');
+  await Incrementer.compile(incrementerAddress);
+  console.log('compile (adder)');
+  await Adder.compile(adderAddress);
+  console.log('compile (caller)');
+  await Caller.compile(zkappAddress);
+}
+
+console.log('deploy');
+let tx = await Mina.transaction(feePayer, () => {
+  // TODO: enable funding multiple accounts properly
+  Party.fundNewAccount(feePayer, {
+    initialBalance: Mina.accountCreationFee().add(Mina.accountCreationFee()),
+  });
+  zkapp.deploy({ zkappKey });
+  if (!doProofs) {
+    zkapp.setPermissions({
+      ...Permissions.default(),
+      editState: Permissions.proofOrSignature(),
+    });
+  }
+  adderZkapp.deploy({ zkappKey: adderKey });
+  incrementerZkapp.deploy({ zkappKey: incrementerKey });
+});
+tx.send();
+
+console.log('call interaction');
+tx = await Mina.transaction(feePayer, () => {
+  // we just call one contract here, nothing special to do
+  zkapp.callAddAndEmit(Field(5), Field(6));
+  if (!doProofs) zkapp.sign(zkappKey);
+});
+if (doProofs) {
+  console.log('proving (3 proofs.. can take a bit!)');
+  await tx.prove();
+}
+
+console.dir(JSON.parse(tx.toJSON()), { depth: 5 });
+
+tx.send();
+
+// should hopefully be 12 since we added 5 + 6 + 1
+console.log('state: ' + zkapp.sum.get());

src/index.ts

@@ -58,10 +58,11 @@ export { Character, CircuitString } from './lib/string';
 // experimental APIs
 import { Reducer } from './lib/zkapp';
 import { createChildParty } from './lib/party';
+import { memoizeWitness } from './lib/circuit_value';
 export { Experimental };
 
 /**
  * This module exposes APIs that are unstable, in the sense that the API surface is expected to change.
  * (Not unstable in the sense that they are less functional or tested than other parts.)
  */
-const Experimental = { Reducer, createChildParty };
+const Experimental = { Reducer, createChildParty, memoizeWitness };

src/lib/circuit_value.ts

@@ -1,5 +1,6 @@
 import 'reflect-metadata';
 import { Circuit, Field, Bool, JSONValue, AsFieldElements } from '../snarky';
+import { Context } from './global-context';
 import { snarkContext } from './proof_system';
 
 // external API
@@ -15,7 +16,15 @@ export {
 };
 
 // internal API
-export { cloneCircuitValue, circuitValueEquals, circuitArray };
+export {
+  cloneCircuitValue,
+  circuitValueEquals,
+  circuitArray,
+  memoizationContext,
+  memoizeWitness,
+  getBlindingValue,
+  toConstant,
+};
 
 type AnyConstructor = new (...args: any) => any;
 
@@ -485,6 +494,10 @@ function circuitValueEquals<T>(a: T, b: T): boolean {
   );
 }
 
+function toConstant<T>(type: AsFieldElements<T>, value: T): T {
+  return type.ofFields(type.toFields(value).map((x) => x.toConstant()));
+}
+
 // TODO: move `Circuit` to JS entirely, this patching harms code discoverability
 Circuit.switch = function <T, A extends AsFieldElements<T>>(
   mask: Bool[],
@@ -533,3 +546,38 @@ Circuit.constraintSystem = function <T>(f: () => T) {
   );
   return result;
 };
+
+let memoizationContext = Context.create<{
+  memoized: Field[][];
+  currentIndex: number;
+  blindingValue: Field;
+}>();
+
+/**
+ * Like Circuit.witness, but memoizes the witness during transaction construction
+ * for reuse by the prover. This is needed to witness non-deterministic values.
+ */
+function memoizeWitness<T>(type: AsFieldElements<T>, compute: () => T) {
+  return Circuit.witness(type, () => {
+    if (!memoizationContext.has()) return compute();
+    let context = memoizationContext.get();
+    let { memoized, currentIndex } = context;
+    let currentValue = memoized[currentIndex];
+    if (currentValue === undefined) {
+      let value = compute();
+      currentValue = type.toFields(value).map((x) => x.toConstant());
+      memoized[currentIndex] = currentValue;
+    }
+    context.currentIndex += 1;
+    return type.ofFields(currentValue);
+  });
+}
+
+function getBlindingValue() {
+  if (!memoizationContext.has()) return Field.random();
+  let context = memoizationContext.get();
+  if (context.blindingValue === undefined) {
+    context.blindingValue = Field.random();
+  }
+  return context.blindingValue;
+}

src/lib/party.test.ts

@@ -7,7 +7,6 @@ import {
   shutdown,
   Field,
   PublicKey,
-  CircuitValue,
   Mina,
   Experimental,
 } from '../../dist/server';
@@ -52,7 +51,7 @@ describe('party', () => {
     // TODO remove restriction "This function can't be run outside of a checked computation."
     Circuit.runAndCheck(() => {
       let hash = party.hash();
-      expect(isLikeField(hash)).toBeTruthy();
+      expect(isField(hash)).toBeTruthy();
 
       // if we clone the party, hash should be the same
       let party2 = Party.clone(party);
@@ -95,13 +94,6 @@ describe('party', () => {
 
 // to check that we got something that looks like a Field
 // note: `instanceof Field` doesn't work
-function isLikeField(x: any) {
-  return (
-    !(x instanceof CircuitValue) &&
-    'equals' in x &&
-    'toFields' in x &&
-    'add' in x &&
-    'mul' in x &&
-    Array.isArray((x as any).value)
-  );
+function isField(x: any): x is Field {
+  return x?.constructor === Field.one.constructor;
 }

src/lib/party.ts

@@ -1,4 +1,9 @@
-import { circuitArray, circuitValue, cloneCircuitValue } from './circuit_value';
+import {
+  circuitArray,
+  circuitValue,
+  cloneCircuitValue,
+  memoizationContext,
+} from './circuit_value';
 import {
   Field,
   Bool,
@@ -493,10 +498,12 @@ type Control = Types.Party['authorization'];
 type LazySignature = { kind: 'lazy-signature'; privateKey?: PrivateKey };
 type LazyProof = {
   kind: 'lazy-proof';
-  method: Function;
+  methodName: string;
   args: any[];
   previousProofs: { publicInput: Field[]; proof: Pickles.Proof }[];
   ZkappClass: typeof SmartContract;
+  memoized: Field[][];
+  blindingValue: Field;
 };
 
 class Token {
@@ -549,7 +556,7 @@ class Party implements Types.Party {
   lazyAuthorization: LazySignature | LazyProof | undefined = undefined;
   account: Precondition.Account;
   network: Precondition.Network;
-  children: Party[] = [];
+  children: { party: Party; calls?: Field }[] = [];
   parent: Party | undefined = undefined;
 
   private isSelf: boolean;
@@ -907,7 +914,8 @@ class Party implements Types.Party {
 
   static witness<T>(
     type: AsFieldElements<T>,
-    compute: () => { party: Party; result: T }
+    compute: () => { party: Party; result: T },
+    skipCheck = false
   ) {
     // construct the circuit type for a party + other result
     let partyType = circuitArray(Field, Types.Party.sizeInFields());
@@ -929,7 +937,10 @@ class Party implements Types.Party {
     // get back a Types.Party from the fields + aux (where aux is just the default in compile)
     let aux = Types.Party.toAuxiliary(proverParty);
     let rawParty = Types.Party.fromFields(fieldsAndResult.party, aux);
-    Types.Party.check(rawParty);
+    // usually when we introduce witnesses, we add checks for their type-specific properties (e.g., booleanness).
+    // a party, however, might already be forced to be valid by the on-chain transaction logic,
+    // allowing us to skip expensive checks in user proofs.
+    if (!skipCheck) Types.Party.check(rawParty);
 
     // construct the full Party instance from the raw party + (maybe) the prover party
     let party = new Party(rawParty.body, rawParty.authorization);
@@ -949,7 +960,8 @@ const CallForest = {
     let parties = [];
     for (let party of forest) {
       party.body.callDepth = depth;
-      parties.push(party, ...CallForest.toFlatList(party.children, depth + 1));
+      let children = party.children.map((c) => c.party);
+      parties.push(party, ...CallForest.toFlatList(children, depth + 1));
     }
     return parties;
   },
@@ -963,8 +975,10 @@ const CallForest = {
   // hashes a party's children (and their children, and ...) to compute the `calls` field of ZkappPublicInput
   hashChildren(parent: Party) {
     let stackHash = CallForest.emptyHash();
-    for (let party of parent.children.reverse()) {
-      let calls = CallForest.hashChildren(party);
+    for (let { party, calls } of parent.children.reverse()) {
+      // only compute calls if it's not there yet --
+      // this gives us the flexibility to witness only direct children of a zkApp
+      calls ??= CallForest.hashChildren(party);
       let nodeHash = hashWithPrefix(prefixes.partyNode, [party.hash(), calls]);
       stackHash = hashWithPrefix(prefixes.partyCons, [nodeHash, stackHash]);
     }
@@ -989,7 +1003,7 @@ function createChildParty(
   child.body.useFullCommitment =
     useFullCommitment ?? child.body.useFullCommitment;
   child.parent = parent;
-  parent.children.push(child);
+  parent.children.push({ party: child, calls: undefined });
   return child;
 }
 
@@ -1133,24 +1147,35 @@ async function addMissingProofs(parties: Parties): Promise<{
     if (party.lazyAuthorization?.kind !== 'lazy-proof') {
       return { partyProved: party as PartyProved, proof: undefined };
     }
-    let { method, args, previousProofs, ZkappClass } = party.lazyAuthorization;
+    let {
+      methodName,
+      args,
+      previousProofs,
+      ZkappClass,
+      memoized,
+      blindingValue,
+    } = party.lazyAuthorization;
     let publicInput = partyToPublicInput(party);
     let publicInputFields = ZkappPublicInput.toFields(publicInput);
     if (ZkappClass._provers === undefined)
       throw Error(
-        `Cannot prove execution of ${method.name}(), no prover found. ` +
+        `Cannot prove execution of ${methodName}(), no prover found. ` +
           `Try calling \`await ${ZkappClass.name}.compile(address)\` first, this will cache provers in the background.`
       );
     let provers = ZkappClass._provers;
     let methodError =
-      `Error when computing proofs: Method ${method.name} not found. ` +
-      `Make sure your environment supports decorators, and annotate with \`@method ${method.name}\`.`;
+      `Error when computing proofs: Method ${methodName} not found. ` +
+      `Make sure your environment supports decorators, and annotate with \`@method ${methodName}\`.`;
     if (ZkappClass._methods === undefined) throw Error(methodError);
-    let i = ZkappClass._methods.findIndex((m) => m.methodName === method.name);
+    let i = ZkappClass._methods.findIndex((m) => m.methodName === methodName);
     if (i === -1) throw Error(methodError);
-    let [, proof] = await snarkContext.runWithAsync(
+    let [, [, proof]] = await snarkContext.runWithAsync(
       { inProver: true, witnesses: args },
-      () => provers[i](publicInputFields, previousProofs)
+      () =>
+        memoizationContext.runWithAsync(
+          { memoized, currentIndex: 0, blindingValue },
+          () => provers[i](publicInputFields, previousProofs)
+        )
     );
     Authorization.setProof(party, Pickles.proofToBase64Transaction(proof));
     let maxProofsVerified = ZkappClass._maxProofsVerified!;

src/lib/precondition.test.ts

@@ -195,6 +195,7 @@ let implementedWithRange = [
 ];
 let unimplemented = [
   () => zkapp.account.provedState,
+  () => zkapp.account.isNew,
   () => zkapp.account.delegate,
   () => zkapp.account.receiptChainHash,
   () => zkapp.network.timestamp,

src/lib/precondition.ts

@@ -54,6 +54,7 @@ let unimplementedPreconditions: LongKey[] = [
   'network.stakingEpochData.lockCheckpoint',
   // this is unimplemented because the field is missing on the account endpoint
   'account.provedState',
+  'account.isNew',
   // this is partially unimplemented because the field is not returned by the local blockchain
   'account.delegate',
   // this is unimplemented because setting this precondition made the integration test fail