diff --git a/src/index.ts b/src/index.ts index 296b3c95f..967223db1 100644 --- a/src/index.ts +++ b/src/index.ts @@ -52,7 +52,10 @@ export { Gadgets } from './lib/provable/gadgets/gadgets.js'; export { Types } from './bindings/mina-transaction/types.js'; export { MerkleList, MerkleListIterator } from './lib/provable/merkle-list.js'; -import { IndexedMerkleMap } from './lib/provable/merkle-tree-indexed.js'; +import { + IndexedMerkleMap, + IndexedMerkleMapBase, +} from './lib/provable/merkle-tree-indexed.js'; export { Option } from './lib/provable/option.js'; export * as Mina from './lib/mina/mina.js'; @@ -146,6 +149,7 @@ namespace Experimental { // indexed merkle map export let IndexedMerkleMap = Experimental_.IndexedMerkleMap; + export type IndexedMerkleMap = IndexedMerkleMapBase; // offchain state export let OffchainState = OffchainState_.OffchainState; diff --git a/src/lib/mina/actions/offchain-contract.unit-test.ts b/src/lib/mina/actions/offchain-contract.unit-test.ts index 882e86efc..0398a14ee 100644 --- a/src/lib/mina/actions/offchain-contract.unit-test.ts +++ b/src/lib/mina/actions/offchain-contract.unit-test.ts @@ -2,7 +2,6 @@ import { SmartContract, method, Mina, - State, state, PublicKey, UInt64, @@ -12,23 +11,22 @@ import assert from 'assert'; const proofsEnabled = true; -const { OffchainState, OffchainStateCommitments } = Experimental; +const { OffchainState } = Experimental; -const offchainState = OffchainState({ - accounts: OffchainState.Map(PublicKey, UInt64), - totalSupply: OffchainState.Field(UInt64), -}); +const offchainState = OffchainState( + { + accounts: OffchainState.Map(PublicKey, UInt64), + totalSupply: OffchainState.Field(UInt64), + }, + { logTotalCapacity: 10, maxActionsPerProof: 5 } +); class StateProof extends offchainState.Proof {} // example contract that interacts with offchain state class ExampleContract extends SmartContract { - // TODO could have sugar for this like - // @OffchainState.commitment offchainState = OffchainState.Commitment(); - @state(OffchainStateCommitments) offchainState = State( - OffchainStateCommitments.empty() - ); + @state(OffchainState.Commitments) offchainState = offchainState.commitments(); @method async createAccount(address: PublicKey, amountToMint: UInt64) { diff --git a/src/lib/mina/actions/offchain-state-rollup.ts b/src/lib/mina/actions/offchain-state-rollup.ts index f4cfa6d00..c64e0a789 100644 --- a/src/lib/mina/actions/offchain-state-rollup.ts +++ b/src/lib/mina/actions/offchain-state-rollup.ts @@ -1,9 +1,11 @@ import { Proof, ZkProgram } from '../../proof-system/zkprogram.js'; import { Bool, Field } from '../../provable/wrapped.js'; -import { Unconstrained } from '../../provable/types/unconstrained.js'; import { MerkleList, MerkleListIterator } from '../../provable/merkle-list.js'; import { Actions } from '../../../bindings/mina-transaction/transaction-leaves.js'; -import { MerkleTree, MerkleWitness } from '../../provable/merkle-tree.js'; +import { + IndexedMerkleMap, + IndexedMerkleMapBase, +} from '../../provable/merkle-tree-indexed.js'; import { Struct } from '../../provable/types/struct.js'; import { SelfProof } from '../../proof-system/zkprogram.js'; import { Provable } from '../../provable/provable.js'; @@ -15,7 +17,6 @@ import { MerkleLeaf, updateMerkleMap, } from './offchain-state-serialization.js'; -import { MerkleMap } from '../../provable/merkle-map.js'; import { getProofsEnabled } from '../mina.js'; export { OffchainStateRollup, OffchainStateCommitments }; @@ -43,8 +44,8 @@ class OffchainStateCommitments extends Struct({ // actionState: ActionIterator.provable, actionState: Field, }) { - static empty() { - let emptyMerkleRoot = new MerkleMap().getRoot(); + static emptyFromHeight(height: number) { + let emptyMerkleRoot = new (IndexedMerkleMap(height))().root; return new OffchainStateCommitments({ root: emptyMerkleRoot, actionState: Actions.emptyActionState(), @@ -52,9 +53,6 @@ class OffchainStateCommitments extends Struct({ } } -const TREE_HEIGHT = 256; -class MerkleMapWitness extends MerkleWitness(TREE_HEIGHT) {} - // TODO: it would be nice to abstract the logic for proving a chain of state transition proofs /** @@ -62,15 +60,15 @@ class MerkleMapWitness extends MerkleWitness(TREE_HEIGHT) {} */ function merkleUpdateBatch( { - maxActionsPerBatch, + maxActionsPerProof, maxActionsPerUpdate, }: { - maxActionsPerBatch: number; + maxActionsPerProof: number; maxActionsPerUpdate: number; }, stateA: OffchainStateCommitments, actions: ActionIterator, - tree: Unconstrained + tree: IndexedMerkleMapBase ): OffchainStateCommitments { // this would be unnecessary if the iterator could just be the public input actions.currentHash.assertEquals(stateA.actionState); @@ -78,7 +76,7 @@ function merkleUpdateBatch( // linearize actions into a flat MerkleList, so we don't process an insane amount of dummy actions let linearActions = LinearizedActionList.empty(); - for (let i = 0; i < maxActionsPerBatch; i++) { + for (let i = 0; i < maxActionsPerProof; i++) { let inner = actions.next().startIterating(); let isAtEnd = Bool(false); for (let i = 0; i < maxActionsPerUpdate; i++) { @@ -99,96 +97,66 @@ function merkleUpdateBatch( } actions.assertAtEnd(); - // update merkle root at once for the actions of each account update - let root = stateA.root; - let intermediateRoot = root; - - let intermediateUpdates: { key: Field; value: Field }[] = []; - let intermediateTree = Unconstrained.witness(() => tree.get().clone()); + // tree must match the public Merkle root; the method operates on the tree internally + // TODO: this would be simpler if the tree was the public input directly + stateA.root.assertEquals(tree.root); + let intermediateTree = tree.clone(); let isValidUpdate = Bool(true); - linearActions.forEach(maxActionsPerBatch, (element, isDummy) => { + linearActions.forEach(maxActionsPerProof, (element, isDummy) => { let { action, isCheckPoint } = element; let { key, value, usesPreviousValue, previousValue } = action; - // merkle witness - let witness = Provable.witness( - MerkleMapWitness, - () => - new MerkleMapWitness(intermediateTree.get().getWitness(key.toBigInt())) - ); + // make sure that if this is a dummy action, we use the canonical dummy (key, value) pair + key = Provable.if(isDummy, Field(0n), key); + value = Provable.if(isDummy, Field(0n), value); - // previous value at the key - let actualPreviousValue = Provable.witness(Field, () => - intermediateTree.get().getLeaf(key.toBigInt()) - ); - - // prove that the witness and `actualPreviousValue` is correct, by comparing the implied root and key - // note: this just works if the (key, value) is a (0,0) dummy, because the value at the 0 key will always be 0 - witness.calculateIndex().assertEquals(key, 'key mismatch'); - witness - .calculateRoot(actualPreviousValue) - .assertEquals(intermediateRoot, 'root mismatch'); + // set (key, value) in the intermediate tree + // note: this just works if (key, value) is a (0,0) dummy, because the value at the 0 key will always be 0 + let actualPreviousValue = intermediateTree.set(key, value); // if an expected previous value was provided, check whether it matches the actual previous value // otherwise, the entire update in invalidated - let matchesPreviousValue = actualPreviousValue.equals(previousValue); + let matchesPreviousValue = actualPreviousValue + .orElse(0n) + .equals(previousValue); let isValidAction = usesPreviousValue.implies(matchesPreviousValue); isValidUpdate = isValidUpdate.and(isValidAction); - // store new value in at the key - let newRoot = witness.calculateRoot(value); - - // update intermediate root if this wasn't a dummy action - intermediateRoot = Provable.if(isDummy, intermediateRoot, newRoot); + // at checkpoints, update the tree, if the entire update was valid + tree.overwriteIf(isCheckPoint.and(isValidUpdate), intermediateTree); - // at checkpoints, update the root, if the entire update was valid - root = Provable.if(isCheckPoint.and(isValidUpdate), intermediateRoot, root); // at checkpoints, reset intermediate values - let wasValidUpdate = isValidUpdate; isValidUpdate = Provable.if(isCheckPoint, Bool(true), isValidUpdate); - intermediateRoot = Provable.if(isCheckPoint, root, intermediateRoot); - - // update the tree, outside the circuit (this should all be part of a better merkle tree API) - Provable.asProver(() => { - // ignore dummy value - if (isDummy.toBoolean()) return; - - intermediateTree.get().setLeaf(key.toBigInt(), value.toConstant()); - intermediateUpdates.push({ key, value }); - - if (isCheckPoint.toBoolean()) { - // if the update was valid, apply the intermediate updates to the actual tree - if (wasValidUpdate.toBoolean()) { - intermediateUpdates.forEach(({ key, value }) => { - tree.get().setLeaf(key.toBigInt(), value.toConstant()); - }); - } - // otherwise, we have to roll back the intermediate tree (TODO: inefficient) - else { - intermediateTree.set(tree.get().clone()); - } - intermediateUpdates = []; - } - }); + intermediateTree.overwriteIf(isCheckPoint, tree); }); - return { root, actionState: actions.currentHash }; + return { root: tree.root, actionState: actions.currentHash }; } /** * This program represents a proof that we can go from OffchainStateCommitments A -> B */ function OffchainStateRollup({ - // 1 action uses about 7.5k constraints - // we can fit at most 7 * 7.5k = 52.5k constraints in one method next to proof verification - // => we use `maxActionsPerBatch = 6` to safely stay below the constraint limit - // the second parameter `maxActionsPerUpdate` only weakly affects # constraints, but has to be <= `maxActionsPerBatch` - // => so we set it to the same value - maxActionsPerBatch = 6, - maxActionsPerUpdate = 6, + /** + * the constraints used in one batch proof with a height-31 tree are: + * + * 1967*A + 87*A*U + 2 + * + * where A = maxActionsPerProof and U = maxActionsPerUpdate. + * + * To determine defaults, we set U=4 which should cover most use cases while ensuring + * that the main loop which is independent of U dominates. + * + * Targeting ~50k constraints, to leave room for recursive verification, yields A=22. + */ + maxActionsPerProof = 22, + maxActionsPerUpdate = 4, + logTotalCapacity = 30, } = {}) { + class IndexedMerkleMapN extends IndexedMerkleMap(logTotalCapacity + 1) {} + let offchainStateRollup = ZkProgram({ name: 'merkle-map-rollup', publicInput: OffchainStateCommitments, @@ -199,15 +167,15 @@ function OffchainStateRollup({ */ firstBatch: { // [actions, tree] - privateInputs: [ActionIterator.provable, Unconstrained.provable], + privateInputs: [ActionIterator.provable, IndexedMerkleMapN.provable], async method( stateA: OffchainStateCommitments, actions: ActionIterator, - tree: Unconstrained + tree: IndexedMerkleMapN ): Promise { return merkleUpdateBatch( - { maxActionsPerBatch, maxActionsPerUpdate }, + { maxActionsPerProof, maxActionsPerUpdate }, stateA, actions, tree @@ -221,14 +189,14 @@ function OffchainStateRollup({ // [actions, tree, proof] privateInputs: [ ActionIterator.provable, - Unconstrained.provable, + IndexedMerkleMapN.provable, SelfProof, ], async method( stateA: OffchainStateCommitments, actions: ActionIterator, - tree: Unconstrained, + tree: IndexedMerkleMapN, recursiveProof: Proof< OffchainStateCommitments, OffchainStateCommitments @@ -247,7 +215,7 @@ function OffchainStateRollup({ let stateB = recursiveProof.publicOutput; return merkleUpdateBatch( - { maxActionsPerBatch, maxActionsPerUpdate }, + { maxActionsPerProof, maxActionsPerUpdate }, stateB, actions, tree @@ -272,8 +240,11 @@ function OffchainStateRollup({ return result; }, - async prove(tree: MerkleTree, actions: MerkleList>) { - assert(tree.height === TREE_HEIGHT, 'Tree height must match'); + async prove( + tree: IndexedMerkleMapN, + actions: MerkleList> + ) { + assert(tree.height === logTotalCapacity + 1, 'Tree height must match'); if (getProofsEnabled()) await this.compile(); // clone the tree so we don't modify the input tree = tree.clone(); @@ -281,7 +252,7 @@ function OffchainStateRollup({ // input state let iterator = actions.startIterating(); let inputState = new OffchainStateCommitments({ - root: tree.getRoot(), + root: tree.root, actionState: iterator.currentHash, }); @@ -304,7 +275,7 @@ function OffchainStateRollup({ updateMerkleMap(actionsList, tree); let finalState = new OffchainStateCommitments({ - root: tree.getRoot(), + root: tree.root, actionState: iterator.hash, }); let proof = await RollupProof.dummy(inputState, finalState, 2, 15); @@ -312,12 +283,13 @@ function OffchainStateRollup({ } // base proof - let slice = sliceActions(iterator, maxActionsPerBatch); - let proof = await offchainStateRollup.firstBatch( - inputState, - slice, - Unconstrained.from(tree) - ); + let slice = sliceActions(iterator, maxActionsPerProof); + let proof = await offchainStateRollup.firstBatch(inputState, slice, tree); + + // update tree root/length again, they aren't mutated :( + // TODO: this shows why the full tree should be the public output + tree.root = proof.publicOutput.root; + tree.length = Field(tree.data.get().sortedLeaves.length); // recursive proofs let nProofs = 1; @@ -325,13 +297,17 @@ function OffchainStateRollup({ if (iterator.isAtEnd().toBoolean()) break; nProofs++; - let slice = sliceActions(iterator, maxActionsPerBatch); + let slice = sliceActions(iterator, maxActionsPerProof); proof = await offchainStateRollup.nextBatch( inputState, slice, - Unconstrained.from(tree), + tree, proof ); + + // update tree root/length again, they aren't mutated :( + tree.root = proof.publicOutput.root; + tree.length = Field(tree.data.get().sortedLeaves.length); } return { proof, tree, nProofs }; diff --git a/src/lib/mina/actions/offchain-state-serialization.ts b/src/lib/mina/actions/offchain-state-serialization.ts index 61d7d3d03..2bc1df8dd 100644 --- a/src/lib/mina/actions/offchain-state-serialization.ts +++ b/src/lib/mina/actions/offchain-state-serialization.ts @@ -24,8 +24,11 @@ import * as Mina from '../mina.js'; import { PublicKey } from '../../provable/crypto/signature.js'; import { Provable } from '../../provable/provable.js'; import { Actions } from '../account-update.js'; -import { MerkleTree } from '../../provable/merkle-tree.js'; import { Option } from '../../provable/option.js'; +import { + IndexedMerkleMap, + IndexedMerkleMapBase, +} from '../../provable/merkle-tree-indexed.js'; export { toKeyHash, @@ -256,9 +259,13 @@ async function fetchMerkleLeaves( * We also deserialize a keyHash -> value map from the leaves. */ async function fetchMerkleMap( + height: number, contract: { address: PublicKey; tokenId: Field }, endActionState?: Field -): Promise<{ merkleMap: MerkleTree; valueMap: Map }> { +): Promise<{ + merkleMap: IndexedMerkleMapBase; + valueMap: Map; +}> { let result = await Mina.fetchActions( contract.address, { endActionState }, @@ -272,7 +279,7 @@ async function fetchMerkleMap( .reverse() ); - let merkleMap = new MerkleTree(256); + let merkleMap = new (IndexedMerkleMap(height))(); let valueMap = new Map(); updateMerkleMap(leaves, merkleMap, valueMap); @@ -282,44 +289,43 @@ async function fetchMerkleMap( function updateMerkleMap( updates: MerkleLeaf[][], - tree: MerkleTree, + tree: IndexedMerkleMapBase, valueMap?: Map ) { let intermediateTree = tree.clone(); for (let leaves of updates) { let isValidUpdate = true; - let updates: { key: bigint; value: bigint; fullValue: Field[] }[] = []; + let updates: { key: bigint; fullValue: Field[] }[] = []; for (let leaf of leaves) { let { key, value, usesPreviousValue, previousValue, prefix } = MerkleLeaf.toValue(leaf); // the update is invalid if there is an unsatisfied precondition + let previous = intermediateTree.getOption(key).orElse(0n); let isValidAction = - !usesPreviousValue || - intermediateTree.getLeaf(key).toBigInt() === previousValue; + !usesPreviousValue || previous.toBigInt() === previousValue; if (!isValidAction) { isValidUpdate = false; - break; } // update the intermediate tree, save updates for final tree - intermediateTree.setLeaf(key, Field(value)); - updates.push({ key, value, fullValue: prefix.get() }); + intermediateTree.set(key, value); + updates.push({ key, fullValue: prefix.get() }); } if (isValidUpdate) { // if the update was valid, we can commit the updates - for (let { key, value, fullValue } of updates) { - tree.setLeaf(key, Field(value)); + tree.overwrite(intermediateTree); + for (let { key, fullValue } of updates) { if (valueMap) valueMap.set(key, fullValue); } } else { // if the update was invalid, we have to roll back the intermediate tree - intermediateTree = tree.clone(); + intermediateTree.overwrite(tree); } } } diff --git a/src/lib/mina/actions/offchain-state.ts b/src/lib/mina/actions/offchain-state.ts index c0582b4e3..18a506441 100644 --- a/src/lib/mina/actions/offchain-state.ts +++ b/src/lib/mina/actions/offchain-state.ts @@ -22,7 +22,7 @@ import { Actions } from '../account-update.js'; import { Provable } from '../../provable/provable.js'; import { Poseidon } from '../../provable/crypto/poseidon.js'; import { smartContractContext } from '../smart-contract-context.js'; -import { MerkleTree, MerkleWitness } from '../../provable/merkle-tree.js'; +import { IndexedMerkleMap } from '../../provable/merkle-tree-indexed.js'; export { OffchainState, OffchainStateCommitments }; @@ -92,14 +92,17 @@ type OffchainState = { settle( proof: Proof ): Promise; + + /** + * Commitments to the offchain state, to use in your onchain state. + */ + commitments(): State; }; type OffchainStateContract = SmartContract & { offchainState: State; }; -const MerkleWitness256 = MerkleWitness(256); - /** * Offchain state for a `SmartContract`. * @@ -131,17 +134,55 @@ const MerkleWitness256 = MerkleWitness(256); */ function OffchainState< const Config extends { [key: string]: OffchainStateKind } ->(config: Config): OffchainState { +>( + config: Config, + options?: { + /** + * The base-2 logarithm of the total capacity of the offchain state. + * + * Example: if you want to have 1 million individual state fields and map entries available, + * set this to 20, because 2^20 ~= 1M. + * + * The default is 30, which allows for ~1 billion entries. + * + * Passing in lower numbers will reduce the number of constraints required to prove offchain state updates, + * which we will make proof creation slightly faster. + * Instead, you could also use a smaller total capacity to increase the `maxActionsPerProof`, so that fewer proofs are required, + * which will reduce the proof time even more, but only in the case of many actions. + */ + logTotalCapacity?: number; + /** + * The maximum number of offchain state actions that can be included in a single account update. + * + * In other words, you must not call `.update()` or `.overwrite()` more than this number of times in any of your smart contract methods. + * + * The default is 4. + * + * Note: When increasing this, consider decreasing `maxActionsPerProof` or `logTotalCapacity` in order to not exceed the circuit size limit. + */ + maxActionsPerUpdate?: number; + maxActionsPerProof?: number; + } +): OffchainState { + // read options + let { + logTotalCapacity = 30, + maxActionsPerUpdate = 4, + maxActionsPerProof, + } = options ?? {}; + const height = logTotalCapacity + 1; + class IndexedMerkleMapN extends IndexedMerkleMap(height) {} + // setup internal state of this "class" let internal = { _contract: undefined as OffchainStateContract | undefined, - _merkleMap: undefined as MerkleTree | undefined, + _merkleMap: undefined as IndexedMerkleMapN | undefined, _valueMap: undefined as Map | undefined, get contract() { assert( internal._contract !== undefined, - 'Must call `setContractAccount()` first' + 'Must call `setContractInstance()` first' ); return internal._contract; }, @@ -159,6 +200,7 @@ function OffchainState< } let actionState = await onchainActionState(); let { merkleMap, valueMap } = await fetchMerkleMap( + height, internal.contract, actionState ); @@ -167,7 +209,11 @@ function OffchainState< return { merkleMap, valueMap }; }; - let rollup = OffchainStateRollup(); + let rollup = OffchainStateRollup({ + logTotalCapacity, + maxActionsPerProof, + maxActionsPerUpdate, + }); function contract() { let ctx = smartContractContext.get(); @@ -189,7 +235,17 @@ function OffchainState< // get onchain merkle root let stateRoot = contract().offchainState.getAndRequireEquals().root; - // witness the actual value + // witness the merkle map & anchor against the onchain root + let map = await Provable.witnessAsync( + IndexedMerkleMapN.provable, + async () => (await merkleMaps()).merkleMap + ); + map.root.assertEquals(stateRoot, 'root mismatch'); + + // get the value hash + let valueHash = map.getOption(key); + + // witness the full value const optionType = Option(valueType); let value = await Provable.witnessAsync(optionType, async () => { let { valueMap } = await merkleMaps(); @@ -201,23 +257,12 @@ function OffchainState< return optionType.from(value); }); - // witness a merkle witness - let witness = await Provable.witnessAsync(MerkleWitness256, async () => { - let { merkleMap } = await merkleMaps(); - return new MerkleWitness256(merkleMap.getWitness(key.toBigInt())); - }); - - // anchor the value against the onchain root and passed in key - // we also allow the value to be missing, in which case the map must contain the 0 element - let valueHash = Provable.if( - value.isSome, - Poseidon.hashPacked(valueType, value.value), - Field(0) + // assert that the value hash matches the value, or both are none + let hashMatches = Poseidon.hashPacked(valueType, value.value).equals( + valueHash.value ); - let actualKey = witness.calculateIndex(); - let actualRoot = witness.calculateRoot(valueHash); - key.assertEquals(actualKey, 'key mismatch'); - stateRoot.assertEquals(actualRoot, 'root mismatch'); + let bothNone = value.isSome.or(valueHash.isSome).not(); + assert(hashMatches.or(bothNone), 'value hash mismatch'); return value; } @@ -340,7 +385,7 @@ function OffchainState< // - take new tree from `result` // - update value map in `prove()`, or separately based on `actions` let { merkleMap: newMerkleMap, valueMap: newValueMap } = - await fetchMerkleMap(internal.contract); + await fetchMerkleMap(height, internal.contract); internal._merkleMap = newMerkleMap; internal._valueMap = newValueMap; @@ -374,11 +419,16 @@ function OffchainState< : map(i, kind.keyType, kind.valueType), ]) ) as any, + + commitments() { + return State(OffchainStateCommitments.emptyFromHeight(height)); + }, }; } OffchainState.Map = OffchainMap; OffchainState.Field = OffchainField; +OffchainState.Commitments = OffchainStateCommitments; // type helpers diff --git a/src/lib/provable/merkle-tree-indexed.ts b/src/lib/provable/merkle-tree-indexed.ts index 484ef3393..262e788df 100644 --- a/src/lib/provable/merkle-tree-indexed.ts +++ b/src/lib/provable/merkle-tree-indexed.ts @@ -11,7 +11,7 @@ import { conditionalSwap } from './merkle-tree.js'; import { provableFromClass } from './types/provable-derivers.js'; // external API -export { IndexedMerkleMap }; +export { IndexedMerkleMap, IndexedMerkleMapBase }; // internal API export { Leaf }; @@ -156,6 +156,30 @@ class IndexedMerkleMapBase { return cloned; } + /** + * Overwrite the entire Merkle map with another one. + * + * This method is provable. + */ + overwrite(other: IndexedMerkleMapBase) { + this.overwriteIf(true, other); + } + + /** + * Overwrite the entire Merkle map with another one, if the condition is true. + * + * This method is provable. + */ + overwriteIf(condition: Bool | boolean, other: IndexedMerkleMapBase) { + condition = Bool(condition); + + this.root = Provable.if(condition, other.root, this.root); + this.length = Provable.if(condition, other.length, this.length); + this.data.updateAsProver(() => + Bool(condition).toBoolean() ? other.clone().data.get() : this.data.get() + ); + } + /** * Insert a new leaf `(key, value)`. *