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/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[];

src/lib/hash.ts

@@ -66,6 +66,17 @@ const Poseidon = {
     return MlFieldArray.from(newState) as [Field, Field, Field];
   },
 
+  hashWithPrefix(prefix: string, input: Field[]) {
+    let init = Poseidon.update(Poseidon.initialState(), [
+      prefixToField(prefix),
+    ]);
+    return Poseidon.update(init, input)[0];
+  },
+
+  initialState(): [Field, Field, Field] {
+    return [Field(0), Field(0), Field(0)];
+  },
+
   hashToGroup(input: Field[]) {
     if (isConstant(input)) {
       let result = PoseidonBigint.hashToGroup(toBigints(input));
@@ -118,10 +129,6 @@ const Poseidon = {
     return Poseidon.hash(packed);
   },
 
-  initialState(): [Field, Field, Field] {
-    return [Field(0), Field(0), Field(0)];
-  },
-
   Sponge,
 };