diff --git a/yarn-project/bb-prover/src/avm_proving.test.ts b/yarn-project/bb-prover/src/avm_proving.test.ts index e3d7607fc82..b830a2614c3 100644 --- a/yarn-project/bb-prover/src/avm_proving.test.ts +++ b/yarn-project/bb-prover/src/avm_proving.test.ts @@ -30,6 +30,7 @@ import fs from 'node:fs/promises'; import { tmpdir } from 'node:os'; import path from 'path'; +import { AvmEphemeralForest } from '../../simulator/src/avm/avm_tree.js'; import { type BBSuccess, BB_RESULT, generateAvmProof, verifyAvmProof } from './bb/execute.js'; import { getPublicInputs } from './test/test_avm.js'; import { extractAvmVkData } from './verification_key/verification_key_data.js'; @@ -71,7 +72,11 @@ const proveAndVerifyAvmTestContract = async ( globals.timestamp = TIMESTAMP; const worldStateDB = mock(); - // + const tmp = openTmpStore(); + const telemetryClient = new NoopTelemetryClient(); + const merkleTree = await (await MerkleTrees.new(tmp, telemetryClient)).fork(); + worldStateDB.getMerkleInterface.mockReturnValue(merkleTree); + // Top level contract call const bytecode = getAvmTestContractBytecode('public_dispatch'); const fnSelector = getAvmTestContractFunctionSelector('public_dispatch'); @@ -106,9 +111,7 @@ const proveAndVerifyAvmTestContract = async ( worldStateDB.storageRead.mockResolvedValue(Promise.resolve(storageValue)); const trace = new PublicSideEffectTrace(startSideEffectCounter); - const telemetry = new NoopTelemetryClient(); - const merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork(); - worldStateDB.getMerkleInterface.mockReturnValue(merkleTrees); + const merkleTrees = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); const persistableState = initPersistableStateManager({ worldStateDB, trace, merkleTrees, doMerkleOperations: true }); const environment = initExecutionEnvironment({ functionSelector, diff --git a/yarn-project/ivc-integration/src/avm_integration.test.ts b/yarn-project/ivc-integration/src/avm_integration.test.ts index 967fd92bc1c..f7593549b3c 100644 --- a/yarn-project/ivc-integration/src/avm_integration.test.ts +++ b/yarn-project/ivc-integration/src/avm_integration.test.ts @@ -49,6 +49,7 @@ import os from 'os'; import path from 'path'; import { fileURLToPath } from 'url'; +import { AvmEphemeralForest } from '../../simulator/src/avm/avm_tree.js'; import { MockPublicKernelCircuit, witnessGenMockPublicKernelCircuit } from './index.js'; // Auto-generated types from noir are not in camel case. @@ -161,6 +162,11 @@ const proveAvmTestContract = async ( assertionErrString?: string, ): Promise => { const worldStateDB = mock(); + const tmp = openTmpStore(); + const telemetryClient = new NoopTelemetryClient(); + const merkleTree = await (await MerkleTrees.new(tmp, telemetryClient)).fork(); + worldStateDB.getMerkleInterface.mockReturnValue(merkleTree); + const startSideEffectCounter = 0; const functionSelector = getAvmTestContractFunctionSelector(functionName); calldata = [functionSelector.toField(), ...calldata]; @@ -200,9 +206,7 @@ const proveAvmTestContract = async ( worldStateDB.storageRead.mockResolvedValue(storageValue); const trace = new PublicSideEffectTrace(startSideEffectCounter); - const telemetry = new NoopTelemetryClient(); - const merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork(); - worldStateDB.getMerkleInterface.mockReturnValue(merkleTrees); + const merkleTrees = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); const persistableState = initPersistableStateManager({ worldStateDB, trace, merkleTrees, doMerkleOperations: true }); const environment = initExecutionEnvironment({ functionSelector, diff --git a/yarn-project/simulator/src/avm/avm_simulator.test.ts b/yarn-project/simulator/src/avm/avm_simulator.test.ts index f05c18b3e01..e4b2ab285d6 100644 --- a/yarn-project/simulator/src/avm/avm_simulator.test.ts +++ b/yarn-project/simulator/src/avm/avm_simulator.test.ts @@ -1,9 +1,7 @@ import { MerkleTreeId, type MerkleTreeWriteOperations } from '@aztec/circuit-types'; import { - type ContractDataSource, GasFees, GlobalVariables, - PublicDataTreeLeaf, PublicDataTreeLeafPreimage, type PublicFunction, PublicKeys, @@ -25,12 +23,13 @@ import { randomInt } from 'crypto'; import { mock } from 'jest-mock-extended'; import { PublicEnqueuedCallSideEffectTrace } from '../public/enqueued_call_side_effect_trace.js'; -import { WorldStateDB } from '../public/public_db_sources.js'; +import { type WorldStateDB } from '../public/public_db_sources.js'; import { type PublicSideEffectTraceInterface } from '../public/side_effect_trace_interface.js'; import { type AvmContext } from './avm_context.js'; import { type AvmExecutionEnvironment } from './avm_execution_environment.js'; import { type MemoryValue, TypeTag, type Uint8, type Uint64 } from './avm_memory_types.js'; import { AvmSimulator } from './avm_simulator.js'; +import { AvmEphemeralForest } from './avm_tree.js'; import { isAvmBytecode, markBytecodeAsAvm } from './bytecode_utils.js'; import { getAvmTestContractArtifact, @@ -46,7 +45,7 @@ import { randomMemoryUint64s, resolveAvmTestContractAssertionMessage, } from './fixtures/index.js'; -import { type AvmPersistableStateManager, getLeafOrLowLeaf } from './journal/journal.js'; +import { type AvmPersistableStateManager } from './journal/journal.js'; import { Add, CalldataCopy, @@ -153,6 +152,11 @@ describe('AVM simulator: transpiled Noir contracts', () => { ), }).withAddress(contractInstance.address); const worldStateDB = mock(); + const tmp = openTmpStore(); + const telemetryClient = new NoopTelemetryClient(); + const merkleTree = await (await MerkleTrees.new(tmp, telemetryClient)).fork(); + worldStateDB.getMerkleInterface.mockReturnValue(merkleTree); + worldStateDB.getContractInstance .mockResolvedValueOnce(contractInstance) .mockResolvedValueOnce(instanceGet) // test gets deployer @@ -165,9 +169,7 @@ describe('AVM simulator: transpiled Noir contracts', () => { mockStorageRead(worldStateDB, storageValue); const trace = mock(); - const telemetry = new NoopTelemetryClient(); - const merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork(); - worldStateDB.getMerkleInterface.mockReturnValue(merkleTrees); + const merkleTrees = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); const persistableState = initPersistableStateManager({ worldStateDB, trace, merkleTrees }); const environment = initExecutionEnvironment({ functionSelector, @@ -1129,39 +1131,33 @@ describe('AVM simulator: transpiled Noir contracts', () => { const sender = AztecAddress.fromNumber(42); const value0 = new Fr(420); - const value1 = new Fr(69); const slotNumber0 = 1; // must update Noir contract if changing this - const slotNumber1 = 2; // must update Noir contract if changing this const slot0 = new Fr(slotNumber0); - const slot1 = new Fr(slotNumber1); const leafSlot0 = computePublicDataTreeLeafSlot(address, slot0); - const leafSlot1 = computePublicDataTreeLeafSlot(address, slot1); - const publicDataTreeLeaf0 = new PublicDataTreeLeaf(leafSlot0, value0); - const _publicDataTreeLeaf1 = new PublicDataTreeLeaf(leafSlot1, value1); - - const INTERNAL_PUBLIC_DATA_SUBTREE_HEIGHT = 0; - - const listSlotNumber0 = 2; // must update Noir contract if changing this - const listSlotNumber1 = listSlotNumber0 + 1; - const listSlot0 = new Fr(listSlotNumber0); - const listSlot1 = new Fr(listSlotNumber1); - const _leafListSlot0 = computePublicDataTreeLeafSlot(address, listSlot0); - const _leafListSlot1 = computePublicDataTreeLeafSlot(address, listSlot1); let worldStateDB: WorldStateDB; let merkleTrees: MerkleTreeWriteOperations; let trace: PublicSideEffectTraceInterface; let persistableState: AvmPersistableStateManager; + let ephemeralForest: AvmEphemeralForest; beforeEach(async () => { trace = mock(); - const telemetry = new NoopTelemetryClient(); - merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork(); - worldStateDB = new WorldStateDB(merkleTrees, mock() as ContractDataSource); - - persistableState = initPersistableStateManager({ worldStateDB, trace, doMerkleOperations: true, merkleTrees }); + worldStateDB = mock(); + const tmp = openTmpStore(); + const telemetryClient = new NoopTelemetryClient(); + merkleTrees = await (await MerkleTrees.new(tmp, telemetryClient)).fork(); + worldStateDB.getMerkleInterface.mockReturnValue(merkleTrees); + ephemeralForest = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); + + persistableState = initPersistableStateManager({ + worldStateDB, + trace, + doMerkleOperations: true, + merkleTrees: ephemeralForest, + }); }); const createContext = (calldata: Fr[] = []) => { @@ -1174,13 +1170,17 @@ describe('AVM simulator: transpiled Noir contracts', () => { describe('Public storage accesses', () => { it('Should set value in storage (single)', async () => { const calldata = [value0]; + const { + preimage: lowLeafPreimage, + index: lowLeafIndex, + update: leafAlreadyPresent, + } = await ephemeralForest.getLeafOrLowLeafInfo( + MerkleTreeId.PUBLIC_DATA_TREE, + leafSlot0, + ); + + const lowLeafPath = await ephemeralForest.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, lowLeafIndex); - const [lowLeafIndex, lowLeafPreimage, lowLeafPath, leafAlreadyPresent] = - await getLeafOrLowLeaf( - MerkleTreeId.PUBLIC_DATA_TREE, - leafSlot0.toBigInt(), - merkleTrees, - ); // leafSlot0 should NOT be present in the tree! expect(leafAlreadyPresent).toEqual(false); expect(lowLeafPreimage.slot).not.toEqual(leafSlot0); @@ -1215,12 +1215,17 @@ describe('AVM simulator: transpiled Noir contracts', () => { it('Should read value in storage (single) - never written', async () => { const context = createContext(); - const [lowLeafIndex, lowLeafPreimage, lowLeafPath, leafAlreadyPresent] = - await getLeafOrLowLeaf( - MerkleTreeId.PUBLIC_DATA_TREE, - leafSlot0.toBigInt(), - merkleTrees, - ); + const { + preimage: lowLeafPreimage, + index: lowLeafIndex, + update: leafAlreadyPresent, + } = await ephemeralForest.getLeafOrLowLeafInfo( + MerkleTreeId.PUBLIC_DATA_TREE, + leafSlot0, + ); + + const lowLeafPath = await ephemeralForest.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, lowLeafIndex); + // leafSlot0 should NOT be present in the tree! expect(leafAlreadyPresent).toEqual(false); expect(lowLeafPreimage.slot).not.toEqual(leafSlot0); @@ -1244,17 +1249,19 @@ describe('AVM simulator: transpiled Noir contracts', () => { it('Should read value in storage (single) - written before, leaf exists', async () => { const context = createContext(); - - await merkleTrees.batchInsert( - MerkleTreeId.PUBLIC_DATA_TREE, - [publicDataTreeLeaf0.toBuffer()], - INTERNAL_PUBLIC_DATA_SUBTREE_HEIGHT, + worldStateDB.storageRead.mockImplementationOnce((_contractAddress: AztecAddress, _slot: Fr) => + Promise.resolve(value0), ); - const [leafIndex, leafPreimage, leafPath] = await getLeafOrLowLeaf( + + await ephemeralForest.writePublicStorage(leafSlot0, value0); + + const { preimage: leafPreimage, index: leafIndex } = await ephemeralForest.getLeafOrLowLeafInfo< MerkleTreeId.PUBLIC_DATA_TREE, - leafSlot0.toBigInt(), - merkleTrees, - ); + PublicDataTreeLeafPreimage + >(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot0); + + const leafPath = await ephemeralForest.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, leafIndex); + // leafSlot0 should be present in the tree! expect(leafPreimage.slot).toEqual(leafSlot0); expect(leafPreimage.value).toEqual(value0); @@ -1279,12 +1286,16 @@ describe('AVM simulator: transpiled Noir contracts', () => { it('Should set and read value in storage (single)', async () => { const calldata = [value0]; - const [lowLeafIndex, lowLeafPreimage, lowLeafPath, leafAlreadyPresent] = - await getLeafOrLowLeaf( - MerkleTreeId.PUBLIC_DATA_TREE, - leafSlot0.toBigInt(), - merkleTrees, - ); + const { + preimage: lowLeafPreimage, + index: lowLeafIndex, + update: leafAlreadyPresent, + } = await ephemeralForest.getLeafOrLowLeafInfo( + MerkleTreeId.PUBLIC_DATA_TREE, + leafSlot0, + ); + const lowLeafPath = await ephemeralForest.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, lowLeafIndex); + // leafSlot0 should NOT be present in the tree! expect(leafAlreadyPresent).toEqual(false); expect(lowLeafPreimage.slot).not.toEqual(leafSlot0); @@ -1302,11 +1313,13 @@ describe('AVM simulator: transpiled Noir contracts', () => { expect(results.reverted).toBe(false); expect(results.output).toEqual([value0]); - const [leafIndex, leafPreimage, leafPath] = await getLeafOrLowLeaf( + const { preimage: leafPreimage, index: leafIndex } = await ephemeralForest.getLeafOrLowLeafInfo< MerkleTreeId.PUBLIC_DATA_TREE, - leafSlot0.toBigInt(), - merkleTrees, - ); + PublicDataTreeLeafPreimage + >(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot0); + + const leafPath = await ephemeralForest.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, leafIndex); + // leafSlot0 should now be present in the tree! expect(leafPreimage.slot).toEqual(leafSlot0); expect(leafPreimage.value).toEqual(value0); diff --git a/yarn-project/simulator/src/avm/fixtures/index.ts b/yarn-project/simulator/src/avm/fixtures/index.ts index 3a8d594e4b5..5711c9d9dc6 100644 --- a/yarn-project/simulator/src/avm/fixtures/index.ts +++ b/yarn-project/simulator/src/avm/fixtures/index.ts @@ -1,4 +1,4 @@ -import { type MerkleTreeWriteOperations, isNoirCallStackUnresolved } from '@aztec/circuit-types'; +import { isNoirCallStackUnresolved } from '@aztec/circuit-types'; import { GasFees, GlobalVariables, MAX_L2_GAS_PER_ENQUEUED_CALL } from '@aztec/circuits.js'; import { type FunctionArtifact, FunctionSelector } from '@aztec/foundation/abi'; import { AztecAddress } from '@aztec/foundation/aztec-address'; @@ -16,6 +16,7 @@ import { AvmContext } from '../avm_context.js'; import { AvmExecutionEnvironment } from '../avm_execution_environment.js'; import { AvmMachineState } from '../avm_machine_state.js'; import { Field, Uint8, Uint32, Uint64 } from '../avm_memory_types.js'; +import { type AvmEphemeralForest } from '../avm_tree.js'; import { type AvmRevertReason } from '../errors.js'; import { AvmPersistableStateManager } from '../journal/journal.js'; import { NullifierManager } from '../journal/nullifiers.js'; @@ -43,7 +44,7 @@ export function initPersistableStateManager(overrides?: { publicStorage?: PublicStorage; nullifiers?: NullifierManager; doMerkleOperations?: boolean; - merkleTrees?: MerkleTreeWriteOperations; + merkleTrees?: AvmEphemeralForest; }): AvmPersistableStateManager { const worldStateDB = overrides?.worldStateDB || mock(); return new AvmPersistableStateManager( @@ -52,7 +53,7 @@ export function initPersistableStateManager(overrides?: { overrides?.publicStorage || new PublicStorage(worldStateDB), overrides?.nullifiers || new NullifierManager(worldStateDB), overrides?.doMerkleOperations || false, - overrides?.merkleTrees || mock(), + overrides?.merkleTrees || mock(), ); } diff --git a/yarn-project/simulator/src/avm/journal/journal.test.ts b/yarn-project/simulator/src/avm/journal/journal.test.ts index 68c009a8afe..976f5f748ea 100644 --- a/yarn-project/simulator/src/avm/journal/journal.test.ts +++ b/yarn-project/simulator/src/avm/journal/journal.test.ts @@ -75,8 +75,8 @@ describe('journal', () => { expect(trace.traceNoteHashCheck).toHaveBeenCalledWith(address, utxo, leafIndex, exists); }); - it('writeNoteHash works', async () => { - await persistableState.writeNoteHash(address, utxo); + it('writeNoteHash works', () => { + persistableState.writeNoteHash(address, utxo); expect(trace.traceNewNoteHash).toHaveBeenCalledTimes(1); expect(trace.traceNewNoteHash).toHaveBeenCalledWith(expect.objectContaining(address), /*noteHash=*/ utxo); }); diff --git a/yarn-project/simulator/src/avm/journal/journal.ts b/yarn-project/simulator/src/avm/journal/journal.ts index bbaa5633160..945ef18d792 100644 --- a/yarn-project/simulator/src/avm/journal/journal.ts +++ b/yarn-project/simulator/src/avm/journal/journal.ts @@ -1,18 +1,16 @@ -import { type IndexedTreeId, MerkleTreeId, type MerkleTreeWriteOperations } from '@aztec/circuit-types'; +import { MerkleTreeId } from '@aztec/circuit-types'; import { type AztecAddress, type Gas, type NullifierLeafPreimage, type PublicCallRequest, - PublicDataTreeLeaf, - PublicDataTreeLeafPreimage, + type PublicDataTreeLeafPreimage, SerializableContractInstance, computePublicBytecodeCommitment, } from '@aztec/circuits.js'; import { computePublicDataTreeLeafSlot, siloNoteHash, siloNullifier } from '@aztec/circuits.js/hash'; import { Fr } from '@aztec/foundation/fields'; import { createDebugLogger } from '@aztec/foundation/log'; -import { type IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import assert from 'assert'; @@ -21,6 +19,7 @@ import { type WorldStateDB } from '../../public/public_db_sources.js'; import { type PublicSideEffectTraceInterface } from '../../public/side_effect_trace_interface.js'; import { type AvmContractCallResult } from '../avm_contract_call_result.js'; import { type AvmExecutionEnvironment } from '../avm_execution_environment.js'; +import { AvmEphemeralForest } from '../avm_tree.js'; import { NullifierManager } from './nullifiers.js'; import { PublicStorage } from './public_storage.js'; @@ -36,9 +35,6 @@ import { PublicStorage } from './public_storage.js'; export class AvmPersistableStateManager { private readonly log = createDebugLogger('aztec:avm_simulator:state_manager'); - /** Interface to perform merkle tree operations */ - public merkleTrees: MerkleTreeWriteOperations; - constructor( /** Reference to node storage */ private readonly worldStateDB: WorldStateDB, @@ -46,33 +42,46 @@ export class AvmPersistableStateManager { // TODO(5818): make private once no longer accessed in executor public readonly trace: PublicSideEffectTraceInterface, /** Public storage, including cached writes */ - private readonly publicStorage: PublicStorage = new PublicStorage(worldStateDB), + public readonly publicStorage: PublicStorage = new PublicStorage(worldStateDB), /** Nullifier set, including cached/recently-emitted nullifiers */ private readonly nullifiers: NullifierManager = new NullifierManager(worldStateDB), private readonly doMerkleOperations: boolean = false, - merkleTrees?: MerkleTreeWriteOperations, - ) { - if (merkleTrees) { - this.merkleTrees = merkleTrees; - } else { - this.merkleTrees = worldStateDB.getMerkleInterface(); - } - } + /** Ephmeral forest for merkle tree operations */ + public readonly merkleTrees: AvmEphemeralForest, + ) {} /** * Create a new state manager with some preloaded pending siloed nullifiers */ - public static newWithPendingSiloedNullifiers( + public static async newWithPendingSiloedNullifiers( worldStateDB: WorldStateDB, trace: PublicSideEffectTraceInterface, pendingSiloedNullifiers: Fr[], ) { const parentNullifiers = NullifierManager.newWithPendingSiloedNullifiers(worldStateDB, pendingSiloedNullifiers); + const ephemeralForest = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); return new AvmPersistableStateManager( worldStateDB, trace, /*publicStorage=*/ new PublicStorage(worldStateDB), /*nullifiers=*/ parentNullifiers.fork(), + /*doMerkleOperations=*/ true, + ephemeralForest, + ); + } + + /** + * Create a new state manager + */ + public static async create(worldStateDB: WorldStateDB, trace: PublicSideEffectTraceInterface) { + const ephemeralForest = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); + return new AvmPersistableStateManager( + worldStateDB, + trace, + /*publicStorage=*/ new PublicStorage(worldStateDB), + /*nullifiers=*/ new NullifierManager(worldStateDB), + /*doMerkleOperations=*/ true, + ephemeralForest, ); } @@ -86,6 +95,7 @@ export class AvmPersistableStateManager { this.publicStorage.fork(), this.nullifiers.fork(), this.doMerkleOperations, + this.merkleTrees.fork(), ); } @@ -102,27 +112,18 @@ export class AvmPersistableStateManager { this.publicStorage.write(contractAddress, slot, value); const leafSlot = computePublicDataTreeLeafSlot(contractAddress, slot); if (this.doMerkleOperations) { - const result = await this.merkleTrees.batchInsert( - MerkleTreeId.PUBLIC_DATA_TREE, - [new PublicDataTreeLeaf(leafSlot, value).toBuffer()], - 0, - ); + const result = await this.merkleTrees.writePublicStorage(leafSlot, value); assert(result !== undefined, 'Public data tree insertion error. You might want to disable skipMerkleOperations.'); this.log.debug(`Inserted public data tree leaf at leafSlot ${leafSlot}, value: ${value}`); - const lowLeafInfo = result.lowLeavesWitnessData![0]; - const lowLeafPreimage = lowLeafInfo.leafPreimage as PublicDataTreeLeafPreimage; + const lowLeafInfo = result.lowWitness; + const lowLeafPreimage = result.lowWitness.preimage as PublicDataTreeLeafPreimage; const lowLeafIndex = lowLeafInfo.index; - const lowLeafPath = lowLeafInfo.siblingPath.toFields(); + const lowLeafPath = lowLeafInfo.siblingPath; + + const insertionPath = result.insertionPath; + const newLeafPreimage = result.newOrElementToUpdate.element as PublicDataTreeLeafPreimage; - const insertionPath = result.newSubtreeSiblingPath.toFields(); - const newLeafPreimage = new PublicDataTreeLeafPreimage( - leafSlot, - value, - lowLeafPreimage.nextSlot, - lowLeafPreimage.nextIndex, - ); - // FIXME: Why do we need to hint both preimages for public data writes, but not for nullifier insertions? this.trace.tracePublicStorageWrite( contractAddress, slot, @@ -156,13 +157,17 @@ export class AvmPersistableStateManager { if (this.doMerkleOperations) { // Get leaf if present, low leaf if absent // If leaf is present, hint/trace it. Otherwise, hint/trace the low leaf. - const [leafIndex, leafPreimage, leafPath, _alreadyPresent] = await getLeafOrLowLeaf( - MerkleTreeId.PUBLIC_DATA_TREE, - leafSlot.toBigInt(), - this.merkleTrees, - ); - // FIXME: cannot have this assertion until "caching" is done via ephemeral merkle writes - //assert(alreadyPresent == exists, 'WorldStateDB contains public data leaf, but merkle tree does not.... This is a bug!'); + const { + preimage, + index: leafIndex, + update, + } = await this.merkleTrees.getLeafOrLowLeafInfo(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot); + // The index and preimage here is either the low leaf or the leaf itself (depending on the value of update flag) + // In either case, we just want the sibling path to this leaf - it's up to the avm to distinguish if it's a low leaf or not + const leafPath = await this.merkleTrees.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, leafIndex); + const leafPreimage = preimage as PublicDataTreeLeafPreimage; + + assert(update == exists, 'WorldStateDB contains public data leaf, but merkle tree does not.... This is a bug!'); this.log.debug( `leafPreimage.nextSlot: ${leafPreimage.nextSlot}, leafPreimage.nextIndex: ${Number(leafPreimage.nextIndex)}`, ); @@ -171,8 +176,9 @@ export class AvmPersistableStateManager { if (!exists) { // Sanity check that the leaf slot is skipped by low leaf when it doesn't exist assert( - leafSlot.toBigInt() > leafPreimage.slot.toBigInt() && leafSlot.toBigInt() < leafPreimage.nextSlot.toBigInt(), - 'Public data tree low leaf should skip the target leaf slot when the target leaf does not exist.', + leafPreimage.slot.toBigInt() < leafSlot.toBigInt() && + (leafPreimage.nextIndex === 0n || leafPreimage.nextSlot.toBigInt() > leafSlot.toBigInt()), + 'Public data tree low leaf should skip the target leaf slot when the target leaf does not exist or is the max value.', ); } this.log.debug( @@ -222,7 +228,7 @@ export class AvmPersistableStateManager { // TODO(8287): We still return exists here, but we need to transmit both the requested noteHash and the gotLeafValue // such that the VM can constrain the equality and decide on exists based on that. const path = await this.merkleTrees.getSiblingPath(MerkleTreeId.NOTE_HASH_TREE, leafIndex.toBigInt()); - this.trace.traceNoteHashCheck(contractAddress, gotLeafValue, leafIndex, exists, path.toFields()); + this.trace.traceNoteHashCheck(contractAddress, gotLeafValue, leafIndex, exists, path); } else { this.trace.traceNoteHashCheck(contractAddress, gotLeafValue, leafIndex, exists); } @@ -233,19 +239,15 @@ export class AvmPersistableStateManager { * Write a note hash, trace the write. * @param noteHash - the unsiloed note hash to write */ - public async writeNoteHash(contractAddress: AztecAddress, noteHash: Fr): Promise { + public writeNoteHash(contractAddress: AztecAddress, noteHash: Fr): void { this.log.debug(`noteHashes(${contractAddress}) += @${noteHash}.`); if (this.doMerkleOperations) { - // TODO: We should track this globally here in the state manager - const info = await this.merkleTrees.getTreeInfo(MerkleTreeId.NOTE_HASH_TREE); - const leafIndex = new Fr(info.size + 1n); - - const path = await this.merkleTrees.getSiblingPath(MerkleTreeId.NOTE_HASH_TREE, leafIndex.toBigInt()); + // Should write a helper for this + const leafIndex = new Fr(this.merkleTrees.treeMap.get(MerkleTreeId.NOTE_HASH_TREE)!.leafCount); const siloedNoteHash = siloNoteHash(contractAddress, noteHash); - - await this.merkleTrees.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, [siloedNoteHash]); - this.trace.traceNewNoteHash(contractAddress, noteHash, leafIndex, path.toFields()); + const insertionPath = this.merkleTrees.appendNoteHash(siloedNoteHash); + this.trace.traceNewNoteHash(contractAddress, noteHash, leafIndex, insertionPath); } else { this.trace.traceNewNoteHash(contractAddress, noteHash); } @@ -265,15 +267,15 @@ export class AvmPersistableStateManager { if (this.doMerkleOperations) { // Get leaf if present, low leaf if absent // If leaf is present, hint/trace it. Otherwise, hint/trace the low leaf. - const [leafIndex, leafPreimage, leafPath, alreadyPresent] = await getLeafOrLowLeaf( - MerkleTreeId.NULLIFIER_TREE, - siloedNullifier.toBigInt(), - this.merkleTrees, - ); - assert( - alreadyPresent == exists, - 'WorldStateDB contains nullifier leaf, but merkle tree does not.... This is a bug!', - ); + const { + preimage, + index: leafIndex, + update, + } = await this.merkleTrees.getLeafOrLowLeafInfo(MerkleTreeId.NULLIFIER_TREE, siloedNullifier); + const leafPreimage = preimage as NullifierLeafPreimage; + const leafPath = await this.merkleTrees.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, leafIndex); + + assert(update == exists, 'WorldStateDB contains nullifier leaf, but merkle tree does not.... This is a bug!'); this.log.debug( `nullifiers(${contractAddress})@${nullifier} ?? leafIndex: ${leafIndex}, exists: ${exists}, pending: ${isPending}.`, @@ -319,33 +321,41 @@ export class AvmPersistableStateManager { const siloedNullifier = siloNullifier(contractAddress, nullifier); if (this.doMerkleOperations) { + // Maybe overkill, but we should check if the nullifier is already present in the tree before attempting to insert + // It might be better to catch the error from the insert operation // Trace all nullifier creations, even duplicate insertions that fail - const alreadyPresent = await this.merkleTrees.getPreviousValueIndex( + const { preimage, index, update } = await this.merkleTrees.getLeafOrLowLeafInfo( MerkleTreeId.NULLIFIER_TREE, - siloedNullifier.toBigInt(), + siloedNullifier, ); - if (alreadyPresent) { - this.log.verbose(`Nullifier already present in tree: ${nullifier} at index ${alreadyPresent.index}.`); + if (update) { + this.log.verbose(`Nullifier already present in tree: ${nullifier} at index ${index}.`); + // If the nullifier is already present, we should not insert it again + // instead we provide the direct memebership path + const path = await this.merkleTrees.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, index); + this.trace.traceNewNullifier( + contractAddress, + nullifier, + preimage as NullifierLeafPreimage, + new Fr(index), + path, + ); + } else { + // We append the new nullifier + const appendResult = await this.merkleTrees.appendNullifier(siloedNullifier); + const lowLeafPreimage = appendResult.lowWitness.preimage as NullifierLeafPreimage; + const lowLeafIndex = appendResult.lowWitness.index; + const lowLeafPath = appendResult.lowWitness.siblingPath; + const insertionPath = appendResult.insertionPath; + this.trace.traceNewNullifier( + contractAddress, + nullifier, + lowLeafPreimage, + new Fr(lowLeafIndex), + lowLeafPath, + insertionPath, + ); } - const insertionResult = await this.merkleTrees.batchInsert( - MerkleTreeId.NULLIFIER_TREE, - [siloedNullifier.toBuffer()], - 0, - ); - const lowLeafInfo = insertionResult.lowLeavesWitnessData![0]; - const lowLeafPreimage = lowLeafInfo.leafPreimage as NullifierLeafPreimage; - const lowLeafIndex = lowLeafInfo.index; - const lowLeafPath = lowLeafInfo.siblingPath.toFields(); - const insertionPath = insertionResult.newSubtreeSiblingPath.toFields(); - - this.trace.traceNewNullifier( - contractAddress, - nullifier, - lowLeafPreimage, - new Fr(lowLeafIndex), - lowLeafPath, - insertionPath, - ); } else { this.trace.traceNewNullifier(contractAddress, nullifier); } @@ -371,7 +381,11 @@ export class AvmPersistableStateManager { if (this.doMerkleOperations) { // TODO(8287): We still return exists here, but we need to transmit both the requested msgHash and the value // such that the VM can constrain the equality and decide on exists based on that. - const path = await this.merkleTrees.getSiblingPath(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, msgLeafIndex.toBigInt()); + // We should defintely add a helper here + const path = await this.merkleTrees.treeDb.getSiblingPath( + MerkleTreeId.L1_TO_L2_MESSAGE_TREE, + msgLeafIndex.toBigInt(), + ); this.trace.traceL1ToL2MessageCheck(contractAddress, valueAtIndex, msgLeafIndex, exists, path.toFields()); } else { this.trace.traceL1ToL2MessageCheck(contractAddress, valueAtIndex, msgLeafIndex, exists); @@ -551,30 +565,3 @@ export class AvmPersistableStateManager { this.trace.traceExecutionPhase(forkedState.trace, publicCallRequests, calldatas, reverted); } } - -/** - * Get leaf if present, low leaf if absent - */ -export async function getLeafOrLowLeaf( - treeId: IndexedTreeId, - key: bigint, - merkleTrees: MerkleTreeWriteOperations, -) { - // "key" is siloed slot (leafSlot) or siloed nullifier - const previousValueIndex = await merkleTrees.getPreviousValueIndex(treeId, key); - assert( - previousValueIndex !== undefined, - `${MerkleTreeId[treeId]} low leaf index should always be found (even if target leaf does not exist)`, - ); - const { index: leafIndex, alreadyPresent } = previousValueIndex; - - const leafPreimage = await merkleTrees.getLeafPreimage(treeId, leafIndex); - assert( - leafPreimage !== undefined, - `${MerkleTreeId[treeId]} low leaf preimage should never be undefined (even if target leaf does not exist)`, - ); - - const leafPath = await merkleTrees.getSiblingPath(treeId, leafIndex); - - return [leafIndex, leafPreimage as TreePreimageType, leafPath.toFields(), alreadyPresent] as const; -} diff --git a/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts b/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts index b93bfa40de7..0a8bd01f140 100644 --- a/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts +++ b/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts @@ -70,7 +70,7 @@ export class EmitNoteHash extends Instruction { } const noteHash = memory.get(noteHashOffset).toFr(); - await context.persistableState.writeNoteHash(context.environment.address, noteHash); + context.persistableState.writeNoteHash(context.environment.address, noteHash); memory.assert({ reads: 1, addressing }); } diff --git a/yarn-project/simulator/src/public/enqueued_calls_processor.ts b/yarn-project/simulator/src/public/enqueued_calls_processor.ts index f1610d2ae04..bc37c038806 100644 --- a/yarn-project/simulator/src/public/enqueued_calls_processor.ts +++ b/yarn-project/simulator/src/public/enqueued_calls_processor.ts @@ -207,7 +207,7 @@ export class EnqueuedCallsProcessor { const trace = new DualSideEffectTrace(innerCallTrace, enqueuedCallTrace); // Transaction level state manager that will be forked for revertible phases. - const txStateManager = AvmPersistableStateManager.newWithPendingSiloedNullifiers( + const txStateManager = await AvmPersistableStateManager.newWithPendingSiloedNullifiers( this.worldStateDB, trace, nonRevertibleNullifiersFromPrivate, diff --git a/yarn-project/simulator/src/public/executor.ts b/yarn-project/simulator/src/public/executor.ts index 3a90b863bdb..79d0dfd5713 100644 --- a/yarn-project/simulator/src/public/executor.ts +++ b/yarn-project/simulator/src/public/executor.ts @@ -116,7 +116,7 @@ export class PublicExecutor { const innerCallTrace = new PublicSideEffectTrace(startSideEffectCounter); const enqueuedCallTrace = new PublicEnqueuedCallSideEffectTrace(startSideEffectCounter); const trace = new DualSideEffectTrace(innerCallTrace, enqueuedCallTrace); - const stateManager = new AvmPersistableStateManager(this.worldStateDB, trace); + const stateManager = await AvmPersistableStateManager.create(this.worldStateDB, trace); return (await this.simulate( stateManager, executionRequest,