diff --git a/l1-contracts/lib/forge-std b/l1-contracts/lib/forge-std index 0e7097750918..75f1746c949a 160000 --- a/l1-contracts/lib/forge-std +++ b/l1-contracts/lib/forge-std @@ -1 +1 @@ -Subproject commit 0e7097750918380d84dd3cfdef595bee74dabb70 +Subproject commit 75f1746c949ae56377611e5f2128aa93d381431f diff --git a/l1-contracts/lib/openzeppelin-contracts b/l1-contracts/lib/openzeppelin-contracts index 448efeea6640..7222a31d5486 160000 --- a/l1-contracts/lib/openzeppelin-contracts +++ b/l1-contracts/lib/openzeppelin-contracts @@ -1 +1 @@ -Subproject commit 448efeea6640bbbc09373f03fbc9c88e280147ba +Subproject commit 7222a31d548695998a475c9661fa159ef45a0e88 diff --git a/yarn-project/prover-client/src/mocks/test_context.ts b/yarn-project/prover-client/src/mocks/test_context.ts index 59eab78431c9..ebecd07801aa 100644 --- a/yarn-project/prover-client/src/mocks/test_context.ts +++ b/yarn-project/prover-client/src/mocks/test_context.ts @@ -82,6 +82,7 @@ export class TestContext { publicDb = await ws.getLatest(); proverDb = await ws.getLatest(); } + worldStateDB.getMerkleInterface.mockReturnValue(publicDb); const publicTxSimulator = new PublicTxSimulator(publicDb, worldStateDB, telemetry, globalVariables); const processor = new PublicProcessor( diff --git a/yarn-project/simulator/src/avm/avm_simulator.test.ts b/yarn-project/simulator/src/avm/avm_simulator.test.ts index 6697b235a65b..f92e013a58a2 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, @@ -1128,39 +1130,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 as jest.Mocked).getMerkleInterface.mockReturnValue(merkleTrees); + ephemeralForest = await AvmEphemeralForest.create(worldStateDB.getMerkleInterface()); + + persistableState = initPersistableStateManager({ + worldStateDB, + trace, + doMerkleOperations: true, + merkleTrees: ephemeralForest, + }); }); const createContext = (calldata: Fr[] = []) => { @@ -1173,13 +1169,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); @@ -1214,12 +1214,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); @@ -1243,17 +1248,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 as jest.Mocked).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); @@ -1278,12 +1285,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); @@ -1301,11 +1312,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 3a8d594e4b57..5711c9d9dc68 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/index.ts b/yarn-project/simulator/src/avm/index.ts index 611d697850cd..4beb53410cbf 100644 --- a/yarn-project/simulator/src/avm/index.ts +++ b/yarn-project/simulator/src/avm/index.ts @@ -1,2 +1,3 @@ export * from './avm_simulator.js'; export * from './journal/index.js'; +export * from './avm_tree.js'; diff --git a/yarn-project/simulator/src/avm/journal/journal.test.ts b/yarn-project/simulator/src/avm/journal/journal.test.ts index 68c009a8afef..976f5f748ea2 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 dd62d155e73e..0396c60830cd 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,7 +19,8 @@ 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 { NullifierManager } from './nullifiers.js'; +import { AvmEphemeralForest } from '../avm_tree.js'; +import { NullifierCollisionError, 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; - /** Make sure a forked state is never merged twice. */ private alreadyMergedIntoParent = false; @@ -53,33 +49,43 @@ export class AvmPersistableStateManager { /** 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[], doMerkleOperations: boolean = false, - merkleTrees?: MerkleTreeWriteOperations, ) { 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, - merkleTrees, + 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, ); } @@ -93,6 +99,7 @@ export class AvmPersistableStateManager { this.publicStorage.fork(), this.nullifiers.fork(), this.doMerkleOperations, + this.merkleTrees.fork(), ); } @@ -143,27 +150,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, @@ -197,13 +195,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)}`, ); @@ -212,15 +214,16 @@ 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( `Tracing storage leaf preimage slot=${slot}, leafSlot=${leafSlot}, value=${value}, nextKey=${leafPreimage.nextSlot}, nextIndex=${leafPreimage.nextIndex}`, ); // On non-existence, AVM circuit will need to recognize that leafPreimage.slot != leafSlot, - // prove that this is a low leaf that skips leafSlot, and then prove memebership of the leaf. + // prove that this is a low leaf that skips leafSlot, and then prove membership of the leaf. this.trace.tracePublicStorageRead(contractAddress, slot, value, leafPreimage, new Fr(leafIndex), leafPath); } else { this.trace.tracePublicStorageRead(contractAddress, slot, value); @@ -263,7 +266,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); } @@ -274,19 +277,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); } @@ -306,15 +305,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}.`, @@ -354,40 +353,55 @@ export class AvmPersistableStateManager { */ public async writeNullifier(contractAddress: AztecAddress, nullifier: Fr) { this.log.debug(`nullifiers(${contractAddress}) += ${nullifier}.`); - // Cache pending nullifiers for later access - await this.nullifiers.append(contractAddress, nullifier); 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 membership path + const path = await this.merkleTrees.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, index); + // This just becomes a nullifier read hint + this.trace.traceNullifierCheck( + contractAddress, + nullifier, + /*exists=*/ update, + preimage as NullifierLeafPreimage, + new Fr(index), + path, + ); + throw new NullifierCollisionError( + `Nullifier ${nullifier} at contract ${contractAddress} already exists in parent cache or host.`, + ); + } else { + // Cache pending nullifiers for later access + await this.nullifiers.append(contractAddress, nullifier); + // 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 { + // Cache pending nullifiers for later access + await this.nullifiers.append(contractAddress, nullifier); this.trace.traceNewNullifier(contractAddress, nullifier); } } @@ -412,7 +426,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); @@ -538,30 +556,3 @@ export class AvmPersistableStateManager { this.trace.traceEnqueuedCall(publicCallRequest, calldata, 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/journal/public_storage.test.ts b/yarn-project/simulator/src/avm/journal/public_storage.test.ts index 442761b67861..5f26cdce4312 100644 --- a/yarn-project/simulator/src/avm/journal/public_storage.test.ts +++ b/yarn-project/simulator/src/avm/journal/public_storage.test.ts @@ -20,6 +20,7 @@ describe('avm public storage', () => { const contractAddress = AztecAddress.fromNumber(1); const slot = new Fr(2); // never written! + publicDb.storageRead.mockResolvedValue(Fr.ZERO); const { exists, value: gotValue, cached } = await publicStorage.read(contractAddress, slot); // doesn't exist, value is zero expect(exists).toEqual(false); diff --git a/yarn-project/simulator/src/avm/journal/public_storage.ts b/yarn-project/simulator/src/avm/journal/public_storage.ts index 794cb7450485..75fa7bc76cf8 100644 --- a/yarn-project/simulator/src/avm/journal/public_storage.ts +++ b/yarn-project/simulator/src/avm/journal/public_storage.ts @@ -77,17 +77,18 @@ export class PublicStorage { let value = this.readHereOrParent(contractAddress, slot); // Finally try the host's Aztec state (a trip to the database) if (!value) { - value = await this.hostPublicStorage.storageRead(contractAddress, slot); + // This functions returns Fr.ZERO if it has never been written to before + // we explicity coalesce to Fr.ZERO in case we have some implementations that cause this to return undefined + value = (await this.hostPublicStorage.storageRead(contractAddress, slot)) ?? Fr.ZERO; // TODO(dbanks12): if value retrieved from host storage, we can cache it here // any future reads to the same slot can read from cache instead of more expensive // DB access } else { cached = true; } - // if value is undefined, that means this slot has never been written to! - const exists = value !== undefined; - const valueOrZero = exists ? value : Fr.ZERO; - return Promise.resolve({ value: valueOrZero, exists, cached }); + // if value is Fr.ZERO here, it that means this slot has never been written to! + const exists = value !== Fr.ZERO; + return Promise.resolve({ value, exists, cached }); } /** diff --git a/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts b/yarn-project/simulator/src/avm/opcodes/accrued_substate.ts index b93bfa40de76..0a8bd01f1409 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/public_tx_context.ts b/yarn-project/simulator/src/public/public_tx_context.ts index 65c287fc769d..fd4c860a35c2 100644 --- a/yarn-project/simulator/src/public/public_tx_context.ts +++ b/yarn-project/simulator/src/public/public_tx_context.ts @@ -111,7 +111,7 @@ export class PublicTxContext { 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( worldStateDB, trace, nonRevertibleNullifiersFromPrivate, diff --git a/yarn-project/simulator/src/public/public_tx_simulator.test.ts b/yarn-project/simulator/src/public/public_tx_simulator.test.ts index 813c074cda4f..a530981a7b64 100644 --- a/yarn-project/simulator/src/public/public_tx_simulator.test.ts +++ b/yarn-project/simulator/src/public/public_tx_simulator.test.ts @@ -1,10 +1,4 @@ -import { - type MerkleTreeWriteOperations, - SimulationError, - type TreeInfo, - TxExecutionPhase, - mockTx, -} from '@aztec/circuit-types'; +import { type MerkleTreeWriteOperations, SimulationError, TxExecutionPhase, mockTx } from '@aztec/circuit-types'; import { AppendOnlyTreeSnapshot, AztecAddress, @@ -16,7 +10,6 @@ import { Header, PUBLIC_DATA_TREE_HEIGHT, PartialStateReference, - PublicDataTreeLeafPreimage, PublicDataWrite, RevertCode, StateReference, @@ -28,6 +21,7 @@ import { type AztecKVStore } from '@aztec/kv-store'; import { openTmpStore } from '@aztec/kv-store/utils'; import { type AppendOnlyTree, Poseidon, StandardTree, newTree } from '@aztec/merkle-tree'; import { NoopTelemetryClient } from '@aztec/telemetry-client/noop'; +import { MerkleTrees } from '@aztec/world-state'; import { jest } from '@jest/globals'; import { type MockProxy, mock } from 'jest-mock-extended'; @@ -50,10 +44,9 @@ describe('public_tx_simulator', () => { // gasUsed for each enqueued call. const enqueuedCallGasUsed = new Gas(12, 34); - let db: MockProxy; + let db: MerkleTreeWriteOperations; let worldStateDB: MockProxy; - let root: Buffer; let publicDataTree: AppendOnlyTree; let treeStore: AztecKVStore; @@ -148,9 +141,10 @@ describe('public_tx_simulator', () => { }; beforeEach(async () => { - db = mock(); + const tmp = openTmpStore(); + const telemetryClient = new NoopTelemetryClient(); + db = await (await MerkleTrees.new(tmp, telemetryClient)).fork(); worldStateDB = mock(); - root = Buffer.alloc(32, 5); treeStore = openTmpStore(); @@ -175,11 +169,7 @@ describe('public_tx_simulator', () => { // Clone the whole state because somewhere down the line (AbstractPhaseManager) the public data root is modified in the referenced header directly :/ header.state = StateReference.fromBuffer(stateReference.toBuffer()); - db.getTreeInfo.mockResolvedValue({ root } as TreeInfo); - db.getStateReference.mockResolvedValue(stateReference); - db.getSiblingPath.mockResolvedValue(publicDataTree.getSiblingPath(0n, false)); - db.getPreviousValueIndex.mockResolvedValue({ index: 0n, alreadyPresent: true }); - db.getLeafPreimage.mockResolvedValue(new PublicDataTreeLeafPreimage(new Fr(0), new Fr(0), new Fr(0), 0n)); + worldStateDB.getMerkleInterface.mockReturnValue(db); simulator = new PublicTxSimulator( db,