Merkle list and token account update iterator

CHANGELOG.md

@@ -17,6 +17,12 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/834a44002...HEAD)
 
+### Added
+
+- `MerkleList<T>` to enable provable operations on a dynamically-sized list https://github.com/o1-labs/o1js/pull/1398
+  - including `MerkleListIterator<T>` to iterate over a merkle list
+- `TokenAccountUpdateIterator`, a primitive for token contracts to iterate over all token account updates in a transaction. https://github.com/o1-labs/o1js/pull/1398
+
 ## [0.16.0](https://github.com/o1-labs/o1js/compare/e5d1e0f...834a44002)
 
 ### Breaking changes

src/examples/zkapps/reducer/actions-as-merkle-list.ts

@@ -0,0 +1,148 @@
+/**
+ * This example shows how to iterate through incoming actions, not using `Reducer.reduce` but by
+ * treating the actions as a merkle list.
+ *
+ * This is mainly intended as an example for using `MerkleList`, but it might also be useful as
+ * a blueprint for processing actions in a custom and more explicit way.
+ */
+import {
+  AccountUpdate,
+  Bool,
+  Field,
+  MerkleList,
+  Mina,
+  Provable,
+  PublicKey,
+  Reducer,
+  SmartContract,
+  method,
+  assert,
+} from 'o1js';
+
+const { Actions } = AccountUpdate;
+
+// in this example, an action is just a public key
+type Action = PublicKey;
+const Action = PublicKey;
+
+// the actions within one account update are a Merkle list with a custom hash
+const emptyHash = Actions.empty().hash;
+const nextHash = (hash: Field, action: Action) =>
+  Actions.pushEvent({ hash, data: [] }, action.toFields()).hash;
+
+class MerkleActions extends MerkleList.create(Action, nextHash, emptyHash) {}
+
+// the "action state" / actions from many account updates is a Merkle list
+// of the above Merkle list, with another custom hash
+let emptyActionsHash = Actions.emptyActionState();
+const nextActionsHash = (hash: Field, actions: MerkleActions) =>
+  Actions.updateSequenceState(hash, actions.hash);
+
+class MerkleActionss extends MerkleList.create(
+  MerkleActions.provable,
+  nextActionsHash,
+  emptyActionsHash
+) {}
+
+// constants for our static-sized provable code
+const MAX_UPDATES_WITH_ACTIONS = 100;
+const MAX_ACTIONS_PER_UPDATE = 2;
+
+/**
+ * This contract allows you to push either 1 or 2 public keys as actions,
+ * and has a reducer-like method which checks whether a given public key is contained in those actions.
+ */
+class ActionsContract extends SmartContract {
+  reducer = Reducer({ actionType: Action });
+
+  @method
+  postAddress(address: PublicKey) {
+    this.reducer.dispatch(address);
+  }
+
+  // to exhibit the generality of reducer: can dispatch more than 1 action per account update
+  @method postTwoAddresses(a1: PublicKey, a2: PublicKey) {
+    this.reducer.dispatch(a1);
+    this.reducer.dispatch(a2);
+  }
+
+  @method
+  assertContainsAddress(address: PublicKey) {
+    // get actions and, in a witness block, wrap them in a Merkle list of lists
+
+    // note: need to reverse here because `getActions()` returns the last pushed action last,
+    // but MerkleList.from() wants it to be first to match the natural iteration order
+    let actionss = this.reducer.getActions().reverse();
+
+    let merkleActionss = Provable.witness(MerkleActionss.provable, () =>
+      MerkleActionss.from(actionss.map((as) => MerkleActions.from(as)))
+    );
+
+    // prove that we know the correct action state
+    this.account.actionState.requireEquals(merkleActionss.hash);
+
+    // now our provable code to process the actions is very straight-forward
+    // (note: if we're past the actual sizes, `.pop()` returns a dummy Action -- in this case, the "empty" public key which is not equal to any real address)
+    let hasAddress = Bool(false);
+
+    for (let i = 0; i < MAX_UPDATES_WITH_ACTIONS; i++) {
+      let merkleActions = merkleActionss.pop();
+
+      for (let j = 0; j < MAX_ACTIONS_PER_UPDATE; j++) {
+        let action = merkleActions.pop();
+        hasAddress = hasAddress.or(action.equals(address));
+      }
+    }
+
+    assert(hasAddress);
+  }
+}
+
+// TESTS
+
+// set up a local blockchain
+
+let Local = Mina.LocalBlockchain({ proofsEnabled: false });
+Mina.setActiveInstance(Local);
+
+let [
+  { publicKey: sender, privateKey: senderKey },
+  { publicKey: zkappAddress, privateKey: zkappKey },
+  { publicKey: otherAddress },
+  { publicKey: anotherAddress },
+] = Local.testAccounts;
+
+let zkapp = new ActionsContract(zkappAddress);
+
+// deploy the contract
+
+await ActionsContract.compile();
+console.log(
+  `rows for ${MAX_UPDATES_WITH_ACTIONS} updates with actions`,
+  ActionsContract.analyzeMethods().assertContainsAddress.rows
+);
+let deployTx = await Mina.transaction(sender, () => zkapp.deploy());
+await deployTx.sign([senderKey, zkappKey]).send();
+
+// push some actions
+
+let dispatchTx = await Mina.transaction(sender, () => {
+  zkapp.postAddress(otherAddress);
+  zkapp.postAddress(zkappAddress);
+  zkapp.postTwoAddresses(anotherAddress, sender);
+  zkapp.postAddress(anotherAddress);
+  zkapp.postTwoAddresses(zkappAddress, otherAddress);
+});
+await dispatchTx.prove();
+await dispatchTx.sign([senderKey]).send();
+
+assert(zkapp.reducer.getActions().length === 5);
+
+// check if the actions contain the `sender` address
+
+Local.setProofsEnabled(true);
+let containsTx = await Mina.transaction(sender, () =>
+  zkapp.assertContainsAddress(sender)
+);
+await containsTx.prove();
+await containsTx.sign([senderKey]).send();

src/index.ts

@@ -9,7 +9,7 @@ export {
 } from './lib/foreign-field.js';
 export { createForeignCurve, ForeignCurve } from './lib/foreign-curve.js';
 export { createEcdsa, EcdsaSignature } from './lib/foreign-ecdsa.js';
-export { Poseidon, TokenSymbol } from './lib/hash.js';
+export { Poseidon, TokenSymbol, ProvableHashable } from './lib/hash.js';
 export { Keccak } from './lib/keccak.js';
 export { Hash } from './lib/hashes-combined.js';
 
@@ -21,7 +21,6 @@ export type {
   FlexibleProvable,
   FlexibleProvablePure,
   InferProvable,
-  Unconstrained,
 } from './lib/circuit_value.js';
 export {
   CircuitValue,
@@ -31,6 +30,7 @@ export {
   provable,
   provablePure,
   Struct,
+  Unconstrained,
 } from './lib/circuit_value.js';
 export { Provable } from './lib/provable.js';
 export { Circuit, Keypair, public_, circuitMain } from './lib/circuit.js';
@@ -40,6 +40,11 @@ export { Packed, Hashed } from './lib/provable-types/packed.js';
 export { Gadgets } from './lib/gadgets/gadgets.js';
 export { Types } from './bindings/mina-transaction/types.js';
 
+export {
+  MerkleList,
+  MerkleListIterator,
+} from './lib/provable-types/merkle-list.js';
+
 export * as Mina from './lib/mina.js';
 export type { DeployArgs } from './lib/zkapp.js';
 export {
@@ -70,8 +75,11 @@ export {
   Permissions,
   ZkappPublicInput,
   TransactionVersion,
+  AccountUpdateForest,
 } from './lib/account_update.js';
 
+export { TokenAccountUpdateIterator } from './lib/mina/token/forest-iterator.js';
+
 export type { TransactionStatus } from './lib/fetch.js';
 export {
   fetchAccount,

src/lib/account_update.ts

@@ -3,6 +3,7 @@ import {
   FlexibleProvable,
   provable,
   provablePure,
+  Struct,
 } from './circuit_value.js';
 import { memoizationContext, memoizeWitness, Provable } from './provable.js';
 import { Field, Bool } from './core.js';
@@ -33,7 +34,12 @@ import {
   Actions,
 } from '../bindings/mina-transaction/transaction-leaves.js';
 import { TokenId as Base58TokenId } from './base58-encodings.js';
-import { hashWithPrefix, packToFields } from './hash.js';
+import {
+  hashWithPrefix,
+  packToFields,
+  Poseidon,
+  ProvableHashable,
+} from './hash.js';
 import {
   mocks,
   prefixes,
@@ -47,9 +53,21 @@ import { transactionCommitments } from '../mina-signer/src/sign-zkapp-command.js
 import { currentTransaction } from './mina/transaction-context.js';
 import { isSmartContract } from './mina/smart-contract-base.js';
 import { activeInstance } from './mina/mina-instance.js';
+import {
+  genericHash,
+  MerkleList,
+  MerkleListBase,
+} from './provable-types/merkle-list.js';
+import { Hashed } from './provable-types/packed.js';
 
 // external API
-export { AccountUpdate, Permissions, ZkappPublicInput, TransactionVersion };
+export {
+  AccountUpdate,
+  Permissions,
+  ZkappPublicInput,
+  TransactionVersion,
+  AccountUpdateForest,
+};
 // internal API
 export {
   smartContractContext,
@@ -74,6 +92,8 @@ export {
   SmartContractContext,
   dummySignature,
   LazyProof,
+  AccountUpdateTree,
+  hashAccountUpdate,
 };
 
 const ZkappStateLength = 8;
@@ -1496,6 +1516,72 @@ type WithCallers = {
   children: WithCallers[];
 };
 
+// call forest stuff
+
+function hashAccountUpdate(update: AccountUpdate) {
+  return genericHash(AccountUpdate, prefixes.body, update);
+}
+
+class HashedAccountUpdate extends Hashed.create(
+  AccountUpdate,
+  hashAccountUpdate
+) {}
+
+type AccountUpdateTree = {
+  accountUpdate: Hashed<AccountUpdate>;
+  calls: MerkleListBase<AccountUpdateTree>;
+};
+const AccountUpdateTree: ProvableHashable<AccountUpdateTree> = Struct({
+  accountUpdate: HashedAccountUpdate.provable,
+  calls: MerkleListBase<AccountUpdateTree>(),
+});
+
+/**
+ * Class which represents a forest (list of trees) of account updates,
+ * in a compressed way which allows iterating and selectively witnessing the account updates.
+ *
+ * The (recursive) type signature is:
+ * ```
+ * type AccountUpdateForest = MerkleList<AccountUpdateTree>;
+ * type AccountUpdateTree = {
+ *   accountUpdate: Hashed<AccountUpdate>;
+ *   calls: AccountUpdateForest;
+ * };
+ * ```
+ */
+class AccountUpdateForest extends MerkleList.create(
+  AccountUpdateTree,
+  merkleListHash
+) {
+  static fromArray(updates: AccountUpdate[]): AccountUpdateForest {
+    let nodes = updates.map((update) => {
+      let accountUpdate = HashedAccountUpdate.hash(update);
+      let calls = AccountUpdateForest.fromArray(update.children.accountUpdates);
+      return { accountUpdate, calls };
+    });
+
+    return AccountUpdateForest.from(nodes);
+  }
+}
+
+// how to hash a forest
+
+function merkleListHash(forestHash: Field, tree: AccountUpdateTree) {
+  return hashCons(forestHash, hashNode(tree));
+}
+function hashNode(tree: AccountUpdateTree) {
+  return Poseidon.hashWithPrefix(prefixes.accountUpdateNode, [
+    tree.accountUpdate.hash,
+    tree.calls.hash,
+  ]);
+}
+function hashCons(forestHash: Field, nodeHash: Field) {
+  return Poseidon.hashWithPrefix(prefixes.accountUpdateCons, [
+    nodeHash,
+    forestHash,
+  ]);
+}
+
 const CallForest = {
   // similar to Mina_base.ZkappCommand.Call_forest.to_account_updates_list
   // takes a list of accountUpdates, which each can have children, so they form a "forest" (list of trees)
@@ -1528,14 +1614,18 @@ const CallForest = {
   // hashes a accountUpdate's children (and their children, and ...) to compute
   // the `calls` field of ZkappPublicInput
   hashChildren(update: AccountUpdate): Field {
+    if (!Provable.inCheckedComputation()) {
+      return CallForest.hashChildrenBase(update);
+    }
+
     let { callsType } = update.children;
     // compute hash outside the circuit if callsType is "Witness"
     // i.e., allowing accountUpdates with arbitrary children
     if (callsType.type === 'Witness') {
       return Provable.witness(Field, () => CallForest.hashChildrenBase(update));
     }
     let calls = CallForest.hashChildrenBase(update);
-    if (callsType.type === 'Equals' && Provable.inCheckedComputation()) {
+    if (callsType.type === 'Equals') {
       calls.assertEquals(callsType.value);
     }
     return calls;

src/lib/circuit_value.ts

@@ -538,6 +538,13 @@ and Provable.asProver() blocks, which execute outside the proof.
     this.option = { isSome: true, value };
   }
 
+  /**
+   * Set the unconstrained value to the same as another `Unconstrained`.
+   */
+  setTo(value: Unconstrained<T>) {
+    this.option = value.option;
+  }
+
   /**
    * Create an `Unconstrained` with the given `value`.
    */
@@ -555,6 +562,16 @@ and Provable.asProver() blocks, which execute outside the proof.
     );
   }
 
+  /**
+   * Update an `Unconstrained` by a witness computation.
+   */
+  updateAsProver(compute: (value: T) => T) {
+    return Provable.asProver(() => {
+      let value = this.get();
+      this.set(compute(value));
+    });
+  }
+
   static provable: Provable<Unconstrained<any>> & {
     toInput: (x: Unconstrained<any>) => {
       fields?: Field[];