chore: cleanup/speedup an AVM test

AztecProtocol · Feb 9, 2025 · 8c2bb0b · 8c2bb0b
1 parent 4a8136c
commit 8c2bb0b
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 7 deletions.
diff --git a/yarn-project/bb-prover/src/avm_proving_tests/avm_proving_tester.ts b/yarn-project/bb-prover/src/avm_proving_tests/avm_proving_tester.ts
@@ -1,6 +1,7 @@
 import { type MerkleTreeWriteOperations } from '@aztec/circuit-types';
 import { type AvmCircuitInputs, AztecAddress, VerificationKeyData } from '@aztec/circuits.js';
 import { PublicTxSimulationTester, type TestEnqueuedCall } from '@aztec/simulator/public/fixtures';
+import { WorldStateDB } from '@aztec/simulator/server';
 import { NativeWorldStateService } from '@aztec/world-state';
 
 import fs from 'node:fs/promises';
@@ -25,21 +26,24 @@ export class AvmProvingTester extends PublicTxSimulationTester {
   constructor(
     private bbWorkingDirectory: string,
     private checkCircuitOnly: boolean,
+    worldStateDB: WorldStateDB,
     contractDataSource: SimpleContractDataSource,
     merkleTrees: MerkleTreeWriteOperations,
     skipContractDeployments: boolean,
   ) {
-    super(contractDataSource, merkleTrees, skipContractDeployments);
+    super(worldStateDB, contractDataSource, merkleTrees, skipContractDeployments);
   }
 
   static override async create(checkCircuitOnly: boolean = false, skipContractDeployments: boolean = false) {
     const bbWorkingDirectory = await fs.mkdtemp(path.join(tmpdir(), 'bb-'));
 
     const contractDataSource = new SimpleContractDataSource();
     const merkleTrees = await (await NativeWorldStateService.tmp()).fork();
+    const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
     return new AvmProvingTester(
       bbWorkingDirectory,
       checkCircuitOnly,
+      worldStateDB,
       contractDataSource,
       merkleTrees,
       skipContractDeployments,
@@ -110,19 +114,27 @@ export class AvmProvingTester extends PublicTxSimulationTester {
 export class AvmProvingTesterV2 extends PublicTxSimulationTester {
   constructor(
     private bbWorkingDirectory: string,
+    worldStateDB: WorldStateDB,
     contractDataSource: SimpleContractDataSource,
     merkleTrees: MerkleTreeWriteOperations,
     skipContractDeployments: boolean,
   ) {
-    super(contractDataSource, merkleTrees, skipContractDeployments);
+    super(worldStateDB, contractDataSource, merkleTrees, skipContractDeployments);
   }
 
   static override async create(skipContractDeployments: boolean = false) {
     const bbWorkingDirectory = await fs.mkdtemp(path.join(tmpdir(), 'bb-'));
 
     const contractDataSource = new SimpleContractDataSource();
     const merkleTrees = await (await NativeWorldStateService.tmp()).fork();
-    return new AvmProvingTesterV2(bbWorkingDirectory, contractDataSource, merkleTrees, skipContractDeployments);
+    const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
+    return new AvmProvingTesterV2(
+      bbWorkingDirectory,
+      worldStateDB,
+      contractDataSource,
+      merkleTrees,
+      skipContractDeployments,
+    );
   }
 
   async proveV2(avmCircuitInputs: AvmCircuitInputs): Promise<BBResult> {

diff --git a/yarn-project/simulator/src/public/fixtures/public_tx_simulation_tester.ts b/yarn-project/simulator/src/public/fixtures/public_tx_simulation_tester.ts
@@ -1,4 +1,4 @@
-import { MerkleTreeId, PublicExecutionRequest, type Tx } from '@aztec/circuit-types';
+import { MerkleTreeId, type MerkleTreeWriteOperations, PublicExecutionRequest, type Tx } from '@aztec/circuit-types';
 import {
   type AvmCircuitPublicInputs,
   CallContext,
@@ -43,10 +43,20 @@ export type TestEnqueuedCall = {
 export class PublicTxSimulationTester extends BaseAvmSimulationTester {
   private txCount = 0;
 
+  constructor(
+    private worldStateDB: WorldStateDB,
+    contractDataSource: SimpleContractDataSource,
+    merkleTrees: MerkleTreeWriteOperations,
+    skipContractDeployments: boolean,
+  ) {
+    super(contractDataSource, merkleTrees, skipContractDeployments);
+  }
+
   public static async create(skipContractDeployments = false): Promise<PublicTxSimulationTester> {
     const contractDataSource = new SimpleContractDataSource();
     const merkleTrees = await (await NativeWorldStateService.tmp()).fork();
-    return new PublicTxSimulationTester(contractDataSource, merkleTrees, skipContractDeployments);
+    const worldStateDB = new WorldStateDB(merkleTrees, contractDataSource);
+    return new PublicTxSimulationTester(worldStateDB, contractDataSource, merkleTrees, skipContractDeployments);
   }
 
   public async simulateTx(
@@ -60,8 +70,7 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
     globals.timestamp = TIMESTAMP;
     globals.gasFees = DEFAULT_GAS_FEES;
 
-    const worldStateDB = new WorldStateDB(this.merkleTrees, this.contractDataSource);
-    const simulator = new PublicTxSimulator(this.merkleTrees, worldStateDB, globals, /*doMerkleOperations=*/ true);
+    const simulator = new PublicTxSimulator(this.merkleTrees, this.worldStateDB, globals, /*doMerkleOperations=*/ true);
 
     const setupExecutionRequests: PublicExecutionRequest[] = [];
     for (let i = 0; i < setupCalls.length; i++) {
@@ -118,7 +127,10 @@ export class PublicTxSimulationTester extends BaseAvmSimulationTester {
       feePayer,
     );
 
+    const startTime = performance.now();
     const avmResult = await simulator.simulate(tx);
+    const endTime = performance.now();
+    this.logger.debug(`Public transaction simulation took ${endTime - startTime}ms`);
 
     if (avmResult.revertCode.isOK()) {
       await this.commitTxStateUpdates(avmResult.avmProvingRequest.inputs.publicInputs);