refactor: Streamline VMTracer construction in EDR stack trace tests

Author: Xanewok

packages/hardhat-core/src/internal/hardhat-network/provider/provider.ts

@@ -12,12 +12,9 @@ import type {
 import type {
   EdrContext,
   Provider as EdrProviderT,
-  ExecutionResult,
   RawTrace,
   Response,
   SubscriptionEvent,
-  TracingMessage,
-  TracingStep,
 } from "@nomicfoundation/edr";
 import { Common } from "@nomicfoundation/ethereumjs-common";
 import chalk from "chalk";
@@ -152,12 +149,6 @@ export function getNodeConfig(
   };
 }
 
-export interface RawTraceCallbacks {
-  onStep?: (messageTrace: TracingStep) => Promise<void>;
-  onBeforeMessage?: (messageTrace: TracingMessage) => Promise<void>;
-  onAfterMessage?: (messageTrace: ExecutionResult) => Promise<void>;
-}
-
 class EdrProviderEventAdapter extends EventEmitter {}
 
 type CallOverrideCallback = (
@@ -184,7 +175,7 @@ export class EdrProviderWrapper
     },
     private readonly _eventAdapter: EdrProviderEventAdapter,
     private readonly _vmTraceDecoder: VmTraceDecoder,
-    private readonly _rawTraceCallbacks: RawTraceCallbacks,
+    private readonly _vmTracer: VMTracer | undefined,
     // The common configuration for EthereumJS VM is not used by EDR, but tests expect it as part of the provider.
     private readonly _common: Common,
     tracingConfig?: TracingConfig
@@ -199,8 +190,8 @@ export class EdrProviderWrapper
   public static async create(
     config: HardhatNetworkProviderConfig,
     loggerConfig: LoggerConfig,
-    rawTraceCallbacks: RawTraceCallbacks,
-    tracingConfig?: TracingConfig
+    tracingConfig?: TracingConfig,
+    vmTracer?: VMTracer
   ): Promise<EdrProviderWrapper> {
     const { Provider } = requireNapiRsModule(
       "@nomicfoundation/edr"
@@ -325,7 +316,7 @@ export class EdrProviderWrapper
       minimalEthereumJsNode,
       eventAdapter,
       vmTraceDecoder,
-      rawTraceCallbacks,
+      vmTracer,
       common,
       tracingConfig
     );
@@ -371,9 +362,7 @@ export class EdrProviderWrapper
     const needsTraces =
       this._node._vm.evm.events.eventNames().length > 0 ||
       this._node._vm.events.eventNames().length > 0 ||
-      this._rawTraceCallbacks.onStep !== undefined ||
-      this._rawTraceCallbacks.onAfterMessage !== undefined ||
-      this._rawTraceCallbacks.onBeforeMessage !== undefined;
+      this._vmTracer !== undefined;
 
     if (needsTraces) {
       const rawTraces = responseObject.traces;
@@ -394,9 +383,8 @@ export class EdrProviderWrapper
                 edrTracingStepToMinimalInterpreterStep(traceItem)
               );
             }
-            if (this._rawTraceCallbacks.onStep !== undefined) {
-              await this._rawTraceCallbacks.onStep(traceItem);
-            }
+
+            this._vmTracer?.addStep(traceItem);
           }
           // afterMessage event
           else if ("executionResult" in traceItem) {
@@ -406,11 +394,8 @@ export class EdrProviderWrapper
                 edrTracingMessageResultToMinimalEVMResult(traceItem)
               );
             }
-            if (this._rawTraceCallbacks.onAfterMessage !== undefined) {
-              await this._rawTraceCallbacks.onAfterMessage(
-                traceItem.executionResult
-              );
-            }
+
+            this._vmTracer?.addAfterMessage(traceItem.executionResult);
           }
           // beforeMessage event
           else {
@@ -420,9 +405,8 @@ export class EdrProviderWrapper
                 edrTracingMessageToMinimalMessage(traceItem)
               );
             }
-            if (this._rawTraceCallbacks.onBeforeMessage !== undefined) {
-              await this._rawTraceCallbacks.onBeforeMessage(traceItem);
-            }
+
+            this._vmTracer?.addBeforeMessage(traceItem);
           }
         }
 
@@ -484,6 +468,11 @@ export class EdrProviderWrapper
     }
   }
 
+  /** Used for internal stack traces integration tests. Only defined there. */
+  public vmTracer(): VMTracer | undefined {
+    return this._vmTracer;
+  }
+
   // temporarily added to make smock work with HH+EDR
   private _setCallOverrideCallback(callback: CallOverrideCallback) {
     this._callOverrideCallback = callback;
@@ -588,11 +577,11 @@ export class EdrProviderWrapper
     const trace = rawTrace.trace();
     for (const traceItem of trace) {
       if ("pc" in traceItem) {
-        await vmTracer.addStep(traceItem);
+        vmTracer.addStep(traceItem);
       } else if ("executionResult" in traceItem) {
-        await vmTracer.addAfterMessage(traceItem.executionResult);
+        vmTracer.addAfterMessage(traceItem.executionResult);
       } else {
-        await vmTracer.addBeforeMessage(traceItem);
+        vmTracer.addBeforeMessage(traceItem);
       }
     }
 
@@ -634,7 +623,6 @@ export async function createHardhatNetworkProvider(
   return EdrProviderWrapper.create(
     hardhatNetworkProviderConfig,
     loggerConfig,
-    {},
     await makeTracingConfig(artifacts)
   );
 }

packages/hardhat-core/src/internal/hardhat-network/stack-traces/vm-tracer.ts

@@ -54,7 +54,7 @@ export class VMTracer {
     return this._throwErrors || this._lastError === undefined;
   }
 
-  public async addBeforeMessage(message: TracingMessage) {
+  public addBeforeMessage(message: TracingMessage) {
     if (!this._shouldKeepTracing()) {
       return;
     }
@@ -151,7 +151,7 @@ export class VMTracer {
     }
   }
 
-  public async addStep(step: TracingStep) {
+  public addStep(step: TracingStep) {
     if (!this._shouldKeepTracing()) {
       return;
     }
@@ -177,7 +177,7 @@ export class VMTracer {
     }
   }
 
-  public async addAfterMessage(result: ExecutionResult, haltOverride?: Exit) {
+  public addAfterMessage(result: ExecutionResult, haltOverride?: Exit) {
     if (!this._shouldKeepTracing()) {
       return;
     }

packages/hardhat-core/test/internal/hardhat-network/stack-traces/execution.ts

@@ -1,3 +1,5 @@
+import { assert } from "chai";
+
 import {
   bigIntToHex,
   bytesToHex,
@@ -29,7 +31,7 @@ const senderAddress = bytesToHex(privateToAddress(toBuffer(senderPrivateKey)));
 export async function instantiateProvider(
   loggerConfig: LoggerConfig,
   tracingConfig: TracingConfig
-): Promise<[EdrProviderWrapper, VMTracer]> {
+): Promise<EdrProviderWrapper> {
   const config = {
     hardfork: "shanghai",
     chainId: 1,
@@ -55,26 +57,14 @@ export async function instantiateProvider(
     enableTransientStorage: false,
   };
 
-  const vmTracer = new VMTracer(false);
-
   const provider = await EdrProviderWrapper.create(
     config,
     loggerConfig,
-    {
-      onStep: async (step) => {
-        await vmTracer.addStep(step);
-      },
-      onAfterMessage: async (message) => {
-        await vmTracer.addAfterMessage(message);
-      },
-      onBeforeMessage: async (message) => {
-        await vmTracer.addBeforeMessage(message);
-      },
-    },
-    tracingConfig
+    tracingConfig,
+    new VMTracer(false)
   );
 
-  return [provider, vmTracer];
+  return provider;
 }
 
 export function encodeConstructorParams(
@@ -116,9 +106,13 @@ export interface TxData {
 
 export async function traceTransaction(
   provider: EdrProviderWrapper,
-  vmTracer: VMTracer,
   txData: TxData
 ): Promise<MessageTrace> {
+  const vmTracer = provider.vmTracer();
+  if (vmTracer === undefined) {
+    assert.fail("VMTracer should always be initialized in stack-traces");
+  }
+
   try {
     await provider.request({
       method: "eth_sendTransaction",

packages/hardhat-core/test/internal/hardhat-network/stack-traces/test.ts

@@ -27,7 +27,6 @@ import {
 } from "../../../../src/internal/hardhat-network/stack-traces/solidity-stack-trace";
 import { SolidityTracer } from "../../../../src/internal/hardhat-network/stack-traces/solidityTracer";
 import { VmTraceDecoder } from "../../../../src/internal/hardhat-network/stack-traces/vm-trace-decoder";
-import { VMTracer } from "../../../../src/internal/hardhat-network/stack-traces/vm-tracer";
 import {
   BuildInfo,
   CompilerInput,
@@ -498,7 +497,7 @@ async function runTest(
 
   const logger = new FakeModulesLogger();
   const solidityTracer = new SolidityTracer();
-  const [provider, vmTracer] = await instantiateProvider(
+  const provider = await instantiateProvider(
     {
       enabled: false,
       printLineFn: logger.printLineFn(),
@@ -517,7 +516,6 @@ async function runTest(
         txIndex,
         tx,
         provider,
-        vmTracer,
         compilerOutput,
         txIndexToContract
       );
@@ -541,7 +539,6 @@ async function runTest(
         txIndex,
         tx,
         provider,
-        vmTracer,
         compilerOutput,
         contract!
       );
@@ -650,7 +647,6 @@ async function runDeploymentTransactionTest(
   txIndex: number,
   tx: DeploymentTransaction,
   provider: EdrProviderWrapper,
-  vmTracer: VMTracer,
   compilerOutput: CompilerOutput,
   txIndexToContract: Map<number, DeployedContract>
 ): Promise<CreateMessageTrace> {
@@ -682,7 +678,7 @@ async function runDeploymentTransactionTest(
 
   const data = Buffer.concat([deploymentBytecode, params]);
 
-  const trace = await traceTransaction(provider, vmTracer, {
+  const trace = await traceTransaction(provider, {
     value: tx.value !== undefined ? BigInt(tx.value) : undefined,
     data,
     gas: tx.gas !== undefined ? BigInt(tx.gas) : undefined,
@@ -695,7 +691,6 @@ async function runCallTransactionTest(
   txIndex: number,
   tx: CallTransaction,
   provider: EdrProviderWrapper,
-  vmTracer: VMTracer,
   compilerOutput: CompilerOutput,
   contract: DeployedContract
 ): Promise<CallMessageTrace> {
@@ -716,7 +711,7 @@ async function runCallTransactionTest(
     data = Buffer.from([]);
   }
 
-  const trace = await traceTransaction(provider, vmTracer, {
+  const trace = await traceTransaction(provider, {
     to: contract.address,
     value: tx.value !== undefined ? BigInt(tx.value) : undefined,
     data,