Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(avm): integrate ephemeral trees #9917

Draft
wants to merge 1 commit into
base: ir/11-09-feat_avm_more_efficient_low_leaf_search
Choose a base branch
from
Draft

feat(avm): integrate ephemeral trees #9917

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions yarn-project/bb-prover/src/avm_proving.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -71,7 +72,11 @@ const proveAndVerifyAvmTestContract = async (
globals.timestamp = TIMESTAMP;

const worldStateDB = mock<WorldStateDB>();
//
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');
Expand Down Expand Up @@ -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,
Expand Down
10 changes: 7 additions & 3 deletions yarn-project/ivc-integration/src/avm_integration.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -161,6 +162,11 @@ const proveAvmTestContract = async (
assertionErrString?: string,
): Promise<BBSuccess> => {
const worldStateDB = mock<WorldStateDB>();
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];
Expand Down Expand Up @@ -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,
Expand Down
129 changes: 71 additions & 58 deletions yarn-project/simulator/src/avm/avm_simulator.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
import { MerkleTreeId, type MerkleTreeWriteOperations } from '@aztec/circuit-types';
import {
type ContractDataSource,
GasFees,
GlobalVariables,
PublicDataTreeLeaf,
PublicDataTreeLeafPreimage,
type PublicFunction,
PublicKeys,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -153,6 +152,11 @@ describe('AVM simulator: transpiled Noir contracts', () => {
),
}).withAddress(contractInstance.address);
const worldStateDB = mock<WorldStateDB>();
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
Expand All @@ -165,9 +169,7 @@ describe('AVM simulator: transpiled Noir contracts', () => {
mockStorageRead(worldStateDB, storageValue);

const trace = mock<PublicSideEffectTraceInterface>();
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,
Expand Down Expand Up @@ -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<PublicSideEffectTraceInterface>();

const telemetry = new NoopTelemetryClient();
merkleTrees = await (await MerkleTrees.new(openTmpStore(), telemetry)).fork();
worldStateDB = new WorldStateDB(merkleTrees, mock<ContractDataSource>() as ContractDataSource);

persistableState = initPersistableStateManager({ worldStateDB, trace, doMerkleOperations: true, merkleTrees });
worldStateDB = mock<WorldStateDB>();
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[] = []) => {
Expand All @@ -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, PublicDataTreeLeafPreimage>(
MerkleTreeId.PUBLIC_DATA_TREE,
leafSlot0,
);

const lowLeafPath = await ephemeralForest.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, lowLeafIndex);

const [lowLeafIndex, lowLeafPreimage, lowLeafPath, leafAlreadyPresent] =
await getLeafOrLowLeaf<PublicDataTreeLeafPreimage>(
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);
Expand Down Expand Up @@ -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<PublicDataTreeLeafPreimage>(
MerkleTreeId.PUBLIC_DATA_TREE,
leafSlot0.toBigInt(),
merkleTrees,
);
const {
preimage: lowLeafPreimage,
index: lowLeafIndex,
update: leafAlreadyPresent,
} = await ephemeralForest.getLeafOrLowLeafInfo<MerkleTreeId.PUBLIC_DATA_TREE, PublicDataTreeLeafPreimage>(
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);
Expand All @@ -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<PublicDataTreeLeafPreimage>(

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);
Expand All @@ -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<PublicDataTreeLeafPreimage>(
MerkleTreeId.PUBLIC_DATA_TREE,
leafSlot0.toBigInt(),
merkleTrees,
);
const {
preimage: lowLeafPreimage,
index: lowLeafIndex,
update: leafAlreadyPresent,
} = await ephemeralForest.getLeafOrLowLeafInfo<MerkleTreeId.PUBLIC_DATA_TREE, PublicDataTreeLeafPreimage>(
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);
Expand All @@ -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<PublicDataTreeLeafPreimage>(
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);
Expand Down
7 changes: 4 additions & 3 deletions yarn-project/simulator/src/avm/fixtures/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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';
Expand Down Expand Up @@ -43,7 +44,7 @@ export function initPersistableStateManager(overrides?: {
publicStorage?: PublicStorage;
nullifiers?: NullifierManager;
doMerkleOperations?: boolean;
merkleTrees?: MerkleTreeWriteOperations;
merkleTrees?: AvmEphemeralForest;
}): AvmPersistableStateManager {
const worldStateDB = overrides?.worldStateDB || mock<WorldStateDB>();
return new AvmPersistableStateManager(
Expand All @@ -52,7 +53,7 @@ export function initPersistableStateManager(overrides?: {
overrides?.publicStorage || new PublicStorage(worldStateDB),
overrides?.nullifiers || new NullifierManager(worldStateDB),
overrides?.doMerkleOperations || false,
overrides?.merkleTrees || mock<MerkleTreeWriteOperations>(),
overrides?.merkleTrees || mock<AvmEphemeralForest>(),
);
}

Expand Down
4 changes: 2 additions & 2 deletions yarn-project/simulator/src/avm/journal/journal.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
Expand Down
Loading
Loading