Skip to content

Commit

Permalink
feat(perf): cleanup and minor speedups in AVM Ephemeral trees
Browse files Browse the repository at this point in the history
  • Loading branch information
@dbanks12
dbanks12 committed Feb 10, 2025
1 parent 18151df commit 0719cf4
Show file tree
Hide file tree
Showing 5 changed files with 61 additions and 45 deletions.
2 changes: 1 addition & 1 deletion yarn-project/simulator/src/avm/avm_simulator.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ export class AvmSimulator {

const endTotalTime = performance.now();
const totalTime = endTotalTime - startTotalTime;
this.log.debug(`Total execution time: ${totalTime}ms`);
this.log.debug(`Core AVM simulation took ${totalTime}ms`);

// Return results for processing by calling context
return results;
Expand Down
28 changes: 22 additions & 6 deletions yarn-project/simulator/src/avm/avm_tree.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,21 @@ import { type IndexedTreeLeafPreimage, type TreeLeafPreimage } from '@aztec/foun
import { strict as assert } from 'assert';
import cloneDeep from 'lodash.clonedeep';

const MAX_TREE_DEPTH = 128;

/**
* Helper function to precompute zero hashes
*/
async function preComputeZeroHashes(): Promise<Fr[]> {
let currentHash = Fr.zero();
const zeroHashes: Fr[] = [];
for (let i = 0; i < MAX_TREE_DEPTH; i++) {
zeroHashes.push(currentHash);
currentHash = await poseidon2Hash([currentHash, currentHash]);
}
return zeroHashes;
}

/****************************************************/
/****** Structs Used by the AvmEphemeralForest ******/
/****************************************************/
Expand Down Expand Up @@ -554,6 +569,8 @@ export class EphemeralAvmTree {
private tree: Tree;
public frontier: Fr[];

private static precomputedZeroHashes: Fr[] | undefined;

private constructor(private root: Leaf, public leafCount: bigint, public depth: number, private zeroHashes: Fr[]) {
this.tree = root;
this.frontier = [];
Expand All @@ -565,13 +582,12 @@ export class EphemeralAvmTree {
treeDb: MerkleTreeReadOperations,
merkleId: MerkleTreeId,
): Promise<EphemeralAvmTree> {
let zeroHash = Fr.zero();
// Can probably cache this elsewhere
const zeroHashes = [];
for (let i = 0; i < depth; i++) {
zeroHashes.push(zeroHash);
zeroHash = await poseidon2Hash([zeroHash, zeroHash]);
let zeroHashes = EphemeralAvmTree.precomputedZeroHashes;
if (zeroHashes === undefined) {
zeroHashes = await preComputeZeroHashes();
EphemeralAvmTree.precomputedZeroHashes = zeroHashes;
}
const zeroHash = zeroHashes[depth];
const tree = new EphemeralAvmTree(Leaf(zeroHash), forkedLeafCount, depth, zeroHashes);
await tree.initializeFrontier(treeDb, merkleId);
return tree;
Expand Down
46 changes: 23 additions & 23 deletions yarn-project/simulator/src/avm/journal/journal.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,16 +144,16 @@ export class AvmPersistableStateManager {
* @param value - the value being written to the slot
*/
public async writeStorage(contractAddress: AztecAddress, slot: Fr, value: Fr, protocolWrite = false): Promise<void> {
this.log.debug(`Storage write (address=${contractAddress}, slot=${slot}): value=${value}`);
this.log.trace(`Storage write (address=${contractAddress}, slot=${slot}): value=${value}`);
const leafSlot = await computePublicDataTreeLeafSlot(contractAddress, slot);
this.log.debug(`leafSlot=${leafSlot}`);
this.log.trace(`leafSlot=${leafSlot}`);
// Cache storage writes for later reference/reads
this.publicStorage.write(contractAddress, slot, value);

if (this.doMerkleOperations) {
const result = await this.merkleTrees.writePublicStorage(leafSlot, value);
assert(result !== undefined, 'Public data tree insertion error. You might want to disable doMerkleOperations.');
this.log.debug(`Inserted public data tree leaf at leafSlot ${leafSlot}, value: ${value}`);
this.log.trace(`Inserted public data tree leaf at leafSlot ${leafSlot}, value: ${value}`);

const lowLeafInfo = result.lowWitness;
const lowLeafPreimage = result.lowWitness.preimage as PublicDataTreeLeafPreimage;
Expand Down Expand Up @@ -195,9 +195,9 @@ export class AvmPersistableStateManager {
*/
public async readStorage(contractAddress: AztecAddress, slot: Fr): Promise<Fr> {
const { value, cached } = await this.publicStorage.read(contractAddress, slot);
this.log.debug(`Storage read (address=${contractAddress}, slot=${slot}): value=${value}, cached=${cached}`);
this.log.trace(`Storage read (address=${contractAddress}, slot=${slot}): value=${value}, cached=${cached}`);
const leafSlot = await computePublicDataTreeLeafSlot(contractAddress, slot);
this.log.debug(`leafSlot=${leafSlot}`);
this.log.trace(`leafSlot=${leafSlot}`);

if (this.doMerkleOperations) {
// Get leaf if present, low leaf if absent
Expand All @@ -212,8 +212,8 @@ export class AvmPersistableStateManager {
const leafPath = await this.merkleTrees.getSiblingPath(MerkleTreeId.PUBLIC_DATA_TREE, leafIndex);
const leafPreimage = preimage as PublicDataTreeLeafPreimage;

this.log.debug(`leafPreimage.slot: ${leafPreimage.slot}, leafPreimage.value: ${leafPreimage.value}`);
this.log.debug(
this.log.trace(`leafPreimage.slot: ${leafPreimage.slot}, leafPreimage.value: ${leafPreimage.value}`);
this.log.trace(
`leafPreimage.nextSlot: ${leafPreimage.nextSlot}, leafPreimage.nextIndex: ${Number(leafPreimage.nextIndex)}`,
);

Expand All @@ -223,7 +223,7 @@ export class AvmPersistableStateManager {
`Value mismatch when performing public data read (got value: ${value}, value in ephemeral tree: ${leafPreimage.value})`,
);
} else {
this.log.debug(`Slot has never been written before!`);
this.log.trace(`Slot has never been written before!`);
// Sanity check that the leaf slot is skipped by low leaf when it doesn't exist
assert(
leafPreimage.slot.toBigInt() < leafSlot.toBigInt() &&
Expand All @@ -250,7 +250,7 @@ export class AvmPersistableStateManager {
*/
public async peekStorage(contractAddress: AztecAddress, slot: Fr): Promise<Fr> {
const { value, cached } = await this.publicStorage.read(contractAddress, slot);
this.log.debug(`Storage peek (address=${contractAddress}, slot=${slot}): value=${value}, cached=${cached}`);
this.log.trace(`Storage peek (address=${contractAddress}, slot=${slot}): value=${value}, cached=${cached}`);
return Promise.resolve(value);
}

Expand All @@ -266,7 +266,7 @@ export class AvmPersistableStateManager {
public async checkNoteHashExists(contractAddress: AztecAddress, noteHash: Fr, leafIndex: Fr): Promise<boolean> {
const gotLeafValue = (await this.worldStateDB.getCommitmentValue(leafIndex.toBigInt())) ?? Fr.ZERO;
const exists = gotLeafValue.equals(noteHash);
this.log.debug(
this.log.trace(
`noteHashes(${contractAddress})@${noteHash} ?? leafIndex: ${leafIndex} | gotLeafValue: ${gotLeafValue}, exists: ${exists}.`,
);
if (this.doMerkleOperations) {
Expand Down Expand Up @@ -306,7 +306,7 @@ export class AvmPersistableStateManager {
* @param noteHash - the siloed unique hash to write
*/
public async writeUniqueNoteHash(noteHash: Fr): Promise<void> {
this.log.debug(`noteHashes += @${noteHash}.`);
this.log.trace(`noteHashes += @${noteHash}.`);

if (this.doMerkleOperations) {
// Should write a helper for this
Expand All @@ -325,7 +325,7 @@ export class AvmPersistableStateManager {
* @returns exists - whether the nullifier exists in the nullifier set
*/
public async checkNullifierExists(contractAddress: AztecAddress, nullifier: Fr): Promise<boolean> {
this.log.debug(`Checking existence of nullifier (address=${contractAddress}, nullifier=${nullifier})`);
this.log.trace(`Checking existence of nullifier (address=${contractAddress}, nullifier=${nullifier})`);
const siloedNullifier = await siloNullifier(contractAddress, nullifier);
const [exists, leafOrLowLeafPreimage, leafOrLowLeafIndex, leafOrLowLeafPath] = await this.getNullifierMembership(
siloedNullifier,
Expand Down Expand Up @@ -367,7 +367,7 @@ export class AvmPersistableStateManager {
]
> {
const [exists, isPending, _] = await this.nullifiers.checkExists(siloedNullifier);
this.log.debug(`Checked siloed nullifier ${siloedNullifier} (exists=${exists}), pending=${isPending}`);
this.log.trace(`Checked siloed nullifier ${siloedNullifier} (exists=${exists}), pending=${isPending}`);

if (this.doMerkleOperations) {
// Get leaf if present, low leaf if absent
Expand All @@ -386,7 +386,7 @@ export class AvmPersistableStateManager {
);

if (exists) {
this.log.debug(`Siloed nullifier ${siloedNullifier} exists at leafIndex=${leafIndex}`);
this.log.trace(`Siloed nullifier ${siloedNullifier} exists at leafIndex=${leafIndex}`);
} else {
// Sanity check that the leaf value is skipped by low leaf when it doesn't exist
assert(
Expand All @@ -407,7 +407,7 @@ export class AvmPersistableStateManager {
* @param nullifier - the unsiloed nullifier to write
*/
public async writeNullifier(contractAddress: AztecAddress, nullifier: Fr) {
this.log.debug(`Inserting new nullifier (address=${nullifier}, nullifier=${contractAddress})`);
this.log.trace(`Inserting new nullifier (address=${nullifier}, nullifier=${contractAddress})`);
const siloedNullifier = await siloNullifier(contractAddress, nullifier);
await this.writeSiloedNullifier(siloedNullifier);
}
Expand All @@ -417,7 +417,7 @@ export class AvmPersistableStateManager {
* @param siloedNullifier - the siloed nullifier to write
*/
public async writeSiloedNullifier(siloedNullifier: Fr) {
this.log.debug(`Inserting siloed nullifier=${siloedNullifier}`);
this.log.trace(`Inserting siloed nullifier=${siloedNullifier}`);

if (this.doMerkleOperations) {
// Maybe overkill, but we should check if the nullifier is already present in the tree before attempting to insert
Expand Down Expand Up @@ -447,13 +447,13 @@ export class AvmPersistableStateManager {
// Cache pending nullifiers for later access
await this.nullifiers.append(siloedNullifier);
// We append the new nullifier
this.log.debug(
this.log.trace(
`Nullifier tree root before insertion ${await this.merkleTrees.treeMap
.get(MerkleTreeId.NULLIFIER_TREE)!
.getRoot()}`,
);
const appendResult = await this.merkleTrees.appendNullifier(siloedNullifier);
this.log.debug(
this.log.trace(
`Nullifier tree root after insertion ${await this.merkleTrees.treeMap
.get(MerkleTreeId.NULLIFIER_TREE)!
.getRoot()}`,
Expand Down Expand Up @@ -496,7 +496,7 @@ export class AvmPersistableStateManager {
): Promise<boolean> {
const valueAtIndex = (await this.worldStateDB.getL1ToL2LeafValue(msgLeafIndex.toBigInt())) ?? Fr.ZERO;
const exists = valueAtIndex.equals(msgHash);
this.log.debug(
this.log.trace(
`l1ToL2Messages(@${msgLeafIndex}) ?? exists: ${exists}, expected: ${msgHash}, found: ${valueAtIndex}.`,
);

Expand All @@ -522,7 +522,7 @@ export class AvmPersistableStateManager {
* @param content - Message content.
*/
public writeL2ToL1Message(contractAddress: AztecAddress, recipient: Fr, content: Fr) {
this.log.debug(`L2ToL1Messages(${contractAddress}) += (recipient: ${recipient}, content: ${content}).`);
this.log.trace(`L2ToL1Messages(${contractAddress}) += (recipient: ${recipient}, content: ${content}).`);
this.trace.traceNewL2ToL1Message(contractAddress, recipient, content);
}

Expand All @@ -532,7 +532,7 @@ export class AvmPersistableStateManager {
* @param log - log contents
*/
public writePublicLog(contractAddress: AztecAddress, log: Fr[]) {
this.log.debug(`PublicLog(${contractAddress}) += event with ${log.length} fields.`);
this.log.trace(`PublicLog(${contractAddress}) += event with ${log.length} fields.`);
this.trace.tracePublicLog(contractAddress, log);
}

Expand All @@ -542,7 +542,7 @@ export class AvmPersistableStateManager {
* @returns the contract instance or undefined if it does not exist.
*/
public async getContractInstance(contractAddress: AztecAddress): Promise<SerializableContractInstance | undefined> {
this.log.debug(`Getting contract instance for address ${contractAddress}`);
this.log.trace(`Getting contract instance for address ${contractAddress}`);
const instanceWithAddress = await this.worldStateDB.getContractInstance(contractAddress);
const exists = instanceWithAddress !== undefined;

Expand All @@ -568,7 +568,7 @@ export class AvmPersistableStateManager {

if (exists) {
const instance = new SerializableContractInstance(instanceWithAddress);
this.log.debug(
this.log.trace(
`Got contract instance (address=${contractAddress}): exists=${exists}, instance=${jsonStringify(instance)}`,
);
if (this.doMerkleOperations) {
Expand Down
4 changes: 2 additions & 2 deletions yarn-project/simulator/src/public/public_tx_context.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -162,15 +162,15 @@ export class PublicTxContext {
* NOTE: this does not "halt" the entire transaction execution.
*/
revert(phase: TxExecutionPhase, revertReason: SimulationError | undefined = undefined, culprit = '') {
this.log.debug(`${TxExecutionPhase[phase]} phase reverted! ${culprit} failed with reason: ${revertReason}`);
this.log.warn(`${TxExecutionPhase[phase]} phase reverted! ${culprit} failed with reason: ${revertReason}`);

if (revertReason && !this.revertReason) {
// don't override revertReason
// (if app logic and teardown both revert, we want app logic's reason)
this.revertReason = revertReason;
}
if (phase === TxExecutionPhase.SETUP) {
this.log.debug(`Setup phase reverted! The transaction will be thrown out.`);
this.log.warn(`Setup phase reverted! The transaction will be thrown out.`);
if (revertReason) {
throw revertReason;
} else {
Expand Down
26 changes: 13 additions & 13 deletions yarn-project/simulator/src/public/side_effect_trace.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,6 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
/** We need to track the set of class IDs used for bytecode retrieval to deduplicate and enforce limits. */
private gotBytecodeFromClassIds: UniqueClassIds = new UniqueClassIds(),
) {
this.log.debug(`Creating trace instance with startSideEffectCounter: ${startSideEffectCounter}`);
this.sideEffectCounter = startSideEffectCounter;
this.avmCircuitHints = AvmExecutionHints.empty();
}
Expand Down Expand Up @@ -215,7 +214,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
path: Fr[] = emptyPublicDataPath(),
) {
this.avmCircuitHints.publicDataReads.items.push(new AvmPublicDataReadTreeHint(leafPreimage, leafIndex, path));
this.log.debug(
this.log.trace(
`Tracing storage read (address=${contractAddress}, slot=${slot}): value=${value} (counter=${this.sideEffectCounter})`,
);
this.incrementSideEffectCounter();
Expand Down Expand Up @@ -265,7 +264,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
new AvmPublicDataWriteTreeHint(readHint, newLeafPreimage, insertionPath),
);

this.log.debug(
this.log.trace(
`Traced public data write (address=${contractAddress}, slot=${slot}): value=${value} (counter=${this.sideEffectCounter}, isProtocol:${protocolWrite})`,
);
this.incrementSideEffectCounter();
Expand All @@ -282,6 +281,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
// New Hinting
this.avmCircuitHints.noteHashReads.items.push(new AvmAppendTreeHint(leafIndex, noteHash, path));
// NOTE: counter does not increment for note hash checks (because it doesn't rely on pending note hashes)
this.log.trace(`Tracing note hash check (counter=${this.sideEffectCounter})`);
}

public traceNewNoteHash(noteHash: Fr, leafIndex: Fr = Fr.zero(), path: Fr[] = emptyNoteHashPath()) {
Expand All @@ -290,8 +290,8 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
}

this.noteHashes.push(new NoteHash(noteHash, this.sideEffectCounter));
this.log.debug(`NEW_NOTE_HASH cnt: ${this.sideEffectCounter}`);
this.avmCircuitHints.noteHashWrites.items.push(new AvmAppendTreeHint(leafIndex, noteHash, path));
this.log.trace(`Tracing new note hash (counter=${this.sideEffectCounter})`);
this.incrementSideEffectCounter();
}

Expand All @@ -305,7 +305,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
this.avmCircuitHints.nullifierReads.items.push(
new AvmNullifierReadTreeHint(lowLeafPreimage, lowLeafIndex, lowLeafPath),
);
this.log.debug(`NULLIFIER_EXISTS cnt: ${this.sideEffectCounter}`);
this.log.trace(`Tracing nullifier check (counter=${this.sideEffectCounter})`);
this.incrementSideEffectCounter();
}

Expand All @@ -324,7 +324,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {

const lowLeafReadHint = new AvmNullifierReadTreeHint(lowLeafPreimage, lowLeafIndex, lowLeafPath);
this.avmCircuitHints.nullifierWrites.items.push(new AvmNullifierWriteTreeHint(lowLeafReadHint, insertionPath));
this.log.debug(`NEW_NULLIFIER cnt: ${this.sideEffectCounter}`);
this.log.trace(`Tracing new nullifier (counter=${this.sideEffectCounter})`);
this.incrementSideEffectCounter();
}

Expand All @@ -337,6 +337,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
path: Fr[] = emptyL1ToL2MessagePath(),
) {
this.avmCircuitHints.l1ToL2MessageReads.items.push(new AvmAppendTreeHint(msgLeafIndex, msgHash, path));
this.log.trace(`Tracing l1 to l2 message check (counter=${this.sideEffectCounter})`);
}

public traceNewL2ToL1Message(contractAddress: AztecAddress, recipient: Fr, content: Fr) {
Expand All @@ -348,7 +349,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
this.l2ToL1Messages.push(
new L2ToL1Message(recipientAddress, content, this.sideEffectCounter).scope(contractAddress),
);
this.log.debug(`NEW_L2_TO_L1_MSG cnt: ${this.sideEffectCounter}`);
this.log.trace(`Tracing new l2 to l1 message (counter=${this.sideEffectCounter})`);
this.incrementSideEffectCounter();
}

Expand All @@ -362,7 +363,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
}
const publicLog = new PublicLog(contractAddress, padArrayEnd(log, Fr.ZERO, PUBLIC_LOG_DATA_SIZE_IN_FIELDS));
this.publicLogs.push(publicLog);
this.log.debug(`NEW_PUBLIC_LOG cnt: ${this.sideEffectCounter}`);
this.log.trace(`Tracing new public log (counter=${this.sideEffectCounter})`);
this.incrementSideEffectCounter();
}

Expand All @@ -387,7 +388,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
membershipHint,
),
);
this.log.debug(`CONTRACT_INSTANCE cnt: ${this.sideEffectCounter}`);
this.log.trace(`Tracing contract instance retrieval (counter=${this.sideEffectCounter})`);
this.incrementSideEffectCounter();
}

Expand Down Expand Up @@ -432,7 +433,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
// Since the bytecode hints are keyed by class ID, we need to hint the instance separately
// since there might be multiple instances hinted for the same class ID.
this.avmCircuitHints.contractInstances.items.push(instance);
this.log.debug(
this.log.trace(
`Tracing contract instance for bytecode retrieval: exists=${exists}, instance=${jsonStringify(contractInstance)}`,
);

Expand All @@ -447,7 +448,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
// To do so, the circuit needs to know the class ID in the
if (this.gotBytecodeFromClassIds.has(contractInstance.contractClassId.toString())) {
// this ensures there are no duplicates
this.log.debug(
this.log.trace(
`Contract class id ${contractInstance.contractClassId.toString()} already exists in previous hints`,
);
return;
Expand All @@ -470,7 +471,7 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
);
}

this.log.debug(`Tracing bytecode & contract class for bytecode retrieval: class=${jsonStringify(contractClass)}`);
this.log.trace(`Tracing bytecode & contract class for bytecode retrieval: class=${jsonStringify(contractClass)}`);
this.avmCircuitHints.contractBytecodeHints.set(
contractInstance.contractClassId.toString(),
new AvmContractBytecodeHints(bytecode, instance, contractClass),
Expand All @@ -492,7 +493,6 @@ export class SideEffectTrace implements PublicSideEffectTraceInterface {
/** Did the call revert? */
_reverted: boolean,
) {
this.log.debug(`Tracing enqueued call`);
// TODO(4805): check if some threshold is reached for max enqueued or nested calls (to unique contracts?)
this.enqueuedCalls.push(publicCallRequest);
this.avmCircuitHints.enqueuedCalls.items.push(new AvmEnqueuedCallHint(publicCallRequest.contractAddress, calldata));
Expand Down

0 comments on commit 0719cf4

Please sign in to comment.