@aztec/simulator 0.55.1 → 0.57.0

package/src/public/enqueued_calls_processor.ts

@@ -0,0 +1,388 @@
+import {
+  type MerkleTreeOperations,
+  type NestedProcessReturnValues,
+  type ProcessedTx,
+  ProvingRequestType,
+  type PublicExecutionRequest,
+  type PublicKernelMergeRequest,
+  PublicKernelPhase,
+  type PublicProvingRequest,
+  type SimulationError,
+  type Tx,
+} from '@aztec/circuit-types';
+import {
+  EnqueuedCallData,
+  Fr,
+  Gas,
+  type GlobalVariables,
+  type Header,
+  type KernelCircuitPublicInputs,
+  NESTED_RECURSIVE_PROOF_LENGTH,
+  type PublicCallRequest,
+  PublicKernelCircuitPrivateInputs,
+  type PublicKernelCircuitPublicInputs,
+  PublicKernelData,
+  type VMCircuitPublicInputs,
+  makeEmptyProof,
+  makeEmptyRecursiveProof,
+} from '@aztec/circuits.js';
+import { type DebugLogger, createDebugLogger } from '@aztec/foundation/log';
+import { Timer } from '@aztec/foundation/timer';
+import { ProtocolCircuitVks, getVKIndex, getVKSiblingPath } from '@aztec/noir-protocol-circuits-types';
+
+import { inspect } from 'util';
+
+import { EnqueuedCallSimulator } from './enqueued_call_simulator.js';
+import { type PublicExecutor } from './executor.js';
+import { type WorldStateDB } from './public_db_sources.js';
+import { type PublicKernelCircuitSimulator } from './public_kernel_circuit_simulator.js';
+import { PublicKernelTailSimulator } from './public_kernel_tail_simulator.js';
+
+const PhaseIsRevertible: Record<PublicKernelPhase, boolean> = {
+  [PublicKernelPhase.SETUP]: false,
+  [PublicKernelPhase.APP_LOGIC]: true,
+  [PublicKernelPhase.TEARDOWN]: true,
+};
+
+type PublicPhaseResult = {
+  /** The output of the public kernel circuit simulation for this phase */
+  publicKernelOutput: PublicKernelCircuitPublicInputs;
+  /** The collection of public proving requests */
+  provingRequests: PublicProvingRequest[];
+  /** Return values of simulating complete callstack */
+  returnValues: NestedProcessReturnValues[];
+  /** Gas used during the execution this phase */
+  gasUsed: Gas;
+  /** Time spent for the execution this phase */
+  durationMs: number;
+  /** Revert reason, if any */
+  revertReason?: SimulationError;
+};
+
+export type ProcessedPhase = {
+  phase: PublicKernelPhase;
+  durationMs: number;
+  revertReason?: SimulationError;
+};
+
+export type TxPublicCallsResult = {
+  /** The collection of public proving requests */
+  provingRequests: PublicProvingRequest[];
+  /** The output of the public kernel tail circuit simulation for this tx */
+  tailKernelOutput: KernelCircuitPublicInputs;
+  /** Return values of simulating complete callstack */
+  returnValues: NestedProcessReturnValues[];
+  /** Gas used during the execution this tx */
+  gasUsed: ProcessedTx['gasUsed'];
+  /** Revert reason, if any */
+  revertReason?: SimulationError;
+  processedPhases: ProcessedPhase[];
+};
+
+export class EnqueuedCallsProcessor {
+  private log: DebugLogger;
+
+  constructor(
+    private publicKernelSimulator: PublicKernelCircuitSimulator,
+    private globalVariables: GlobalVariables,
+    private worldStateDB: WorldStateDB,
+    private enqueuedCallSimulator: EnqueuedCallSimulator,
+    private publicKernelTailSimulator: PublicKernelTailSimulator,
+  ) {
+    this.log = createDebugLogger(`aztec:sequencer`);
+  }
+
+  static create(
+    db: MerkleTreeOperations,
+    publicExecutor: PublicExecutor,
+    publicKernelSimulator: PublicKernelCircuitSimulator,
+    globalVariables: GlobalVariables,
+    historicalHeader: Header,
+    worldStateDB: WorldStateDB,
+  ) {
+    const enqueuedCallSimulator = new EnqueuedCallSimulator(
+      db,
+      publicExecutor,
+      publicKernelSimulator,
+      globalVariables,
+      historicalHeader,
+    );
+
+    const publicKernelTailSimulator = PublicKernelTailSimulator.create(db, publicKernelSimulator);
+
+    return new EnqueuedCallsProcessor(
+      publicKernelSimulator,
+      globalVariables,
+      worldStateDB,
+      enqueuedCallSimulator,
+      publicKernelTailSimulator,
+    );
+  }
+
+  static getExecutionRequestsByPhase(tx: Tx, phase: PublicKernelPhase): PublicExecutionRequest[] {
+    switch (phase) {
+      case PublicKernelPhase.SETUP:
+        return tx.getNonRevertiblePublicExecutionRequests();
+      case PublicKernelPhase.APP_LOGIC:
+        return tx.getRevertiblePublicExecutionRequests();
+      case PublicKernelPhase.TEARDOWN: {
+        const request = tx.getPublicTeardownExecutionRequest();
+        return request ? [request] : [];
+      }
+      default:
+        throw new Error(`Unknown phase: ${phase}`);
+    }
+  }
+
+  static getCallRequestsByPhase(tx: Tx, phase: PublicKernelPhase): PublicCallRequest[] {
+    switch (phase) {
+      case PublicKernelPhase.SETUP:
+        return tx.data.getNonRevertiblePublicCallRequests();
+      case PublicKernelPhase.APP_LOGIC:
+        return tx.data.getRevertiblePublicCallRequests();
+      case PublicKernelPhase.TEARDOWN: {
+        const request = tx.data.getTeardownPublicCallRequest();
+        return request ? [request] : [];
+      }
+      default:
+        throw new Error(`Unknown phase: ${phase}`);
+    }
+  }
+
+  async process(tx: Tx): Promise<TxPublicCallsResult> {
+    this.log.verbose(`Processing tx ${tx.getTxHash()}`);
+
+    const phases: PublicKernelPhase[] = [
+      PublicKernelPhase.SETUP,
+      PublicKernelPhase.APP_LOGIC,
+      PublicKernelPhase.TEARDOWN,
+    ];
+    const processedPhases: ProcessedPhase[] = [];
+    const provingRequests: PublicProvingRequest[] = [];
+    const gasUsed: ProcessedTx['gasUsed'] = {};
+    let publicKernelOutput = tx.data.toPublicKernelCircuitPublicInputs();
+    let isFromPrivate = true;
+    let returnValues: NestedProcessReturnValues[] = [];
+    let revertReason: SimulationError | undefined;
+    for (let i = 0; i < phases.length; i++) {
+      const phase = phases[i];
+      const callRequests = EnqueuedCallsProcessor.getCallRequestsByPhase(tx, phase);
+      if (callRequests.length) {
+        const executionRequests = EnqueuedCallsProcessor.getExecutionRequestsByPhase(tx, phase);
+        const result = await this.processPhase(
+          tx,
+          callRequests,
+          executionRequests,
+          publicKernelOutput,
+          phase,
+          isFromPrivate,
+        ).catch(async err => {
+          await this.worldStateDB.rollbackToCommit();
+          throw err;
+        });
+
+        publicKernelOutput = result.publicKernelOutput;
+        isFromPrivate = false;
+
+        provingRequests.push(...result.provingRequests);
+        if (phase === PublicKernelPhase.APP_LOGIC) {
+          returnValues = result.returnValues;
+        }
+
+        gasUsed[phase] = result.gasUsed;
+
+        processedPhases.push({
+          phase,
+          durationMs: result.durationMs,
+          revertReason: result.revertReason,
+        });
+
+        revertReason ??= result.revertReason;
+      }
+    }
+
+    const { output: tailKernelOutput, provingRequest } = await this.publicKernelTailSimulator
+      .simulate(publicKernelOutput)
+      .catch(
+        // the abstract phase manager throws if simulation gives error in non-revertible phase
+        async err => {
+          await this.worldStateDB.rollbackToCommit();
+          throw err;
+        },
+      );
+    provingRequests.push(provingRequest);
+
+    return {
+      provingRequests: provingRequests,
+      tailKernelOutput,
+      returnValues,
+      gasUsed,
+      processedPhases,
+      revertReason,
+    };
+  }
+
+  private async processPhase(
+    tx: Tx,
+    callRequests: PublicCallRequest[],
+    executionRequests: PublicExecutionRequest[],
+    previousPublicKernelOutput: PublicKernelCircuitPublicInputs,
+    phase: PublicKernelPhase,
+    isFromPrivate: boolean,
+  ): Promise<PublicPhaseResult> {
+    this.log.debug(`Beginning processing in phase ${PublicKernelPhase[phase]} for tx ${tx.getTxHash()}`);
+
+    const phaseTimer = new Timer();
+    const provingRequests: PublicProvingRequest[] = [];
+    const returnValues: NestedProcessReturnValues[] = [];
+    let publicKernelOutput = previousPublicKernelOutput;
+    let gasUsed = Gas.empty();
+    let revertReason: SimulationError | undefined;
+    for (let i = callRequests.length - 1; i >= 0 && !revertReason; i--) {
+      const callRequest = callRequests[i];
+      const executionRequest = executionRequests[i];
+
+      // add new contracts to the contracts db so that their functions may be found and called
+      // TODO(#4073): This is catching only private deployments, when we add public ones, we'll
+      // have to capture contracts emitted in that phase as well.
+      // TODO(@spalladino): Should we allow emitting contracts in the fee preparation phase?
+      // TODO(#6464): Should we allow emitting contracts in the private setup phase?
+      // if so, this should only add contracts that were deployed during private app logic.
+      await this.worldStateDB.addNewContracts(tx);
+
+      const availableGas = this.getAvailableGas(tx, publicKernelOutput, phase);
+      const transactionFee = this.getTransactionFee(tx, publicKernelOutput, phase);
+
+      const enqueuedCallResult = await this.enqueuedCallSimulator.simulate(
+        callRequest,
+        executionRequest,
+        tx,
+        publicKernelOutput,
+        availableGas,
+        transactionFee,
+        phase,
+      );
+
+      if (enqueuedCallResult.revertReason && !PhaseIsRevertible[phase]) {
+        this.log.debug(
+          `Simulation error on ${executionRequest.contractAddress.toString()}:${
+            executionRequest.callContext.functionSelector
+          } with reason: ${enqueuedCallResult.revertReason}`,
+        );
+        throw enqueuedCallResult.revertReason;
+      }
+
+      provingRequests.push(...enqueuedCallResult.provingRequests);
+      returnValues.push(enqueuedCallResult.returnValues);
+      gasUsed = gasUsed.add(enqueuedCallResult.gasUsed);
+      revertReason ??= enqueuedCallResult.revertReason;
+
+      if (revertReason) {
+        // TODO(#6464): Should we allow emitting contracts in the private setup phase?
+        // if so, this is removing contracts deployed in private setup
+        await this.worldStateDB.removeNewContracts(tx);
+        await this.worldStateDB.rollbackToCheckpoint();
+        tx.filterRevertedLogs(publicKernelOutput);
+      } else {
+        // TODO(#6470): we should be adding contracts deployed in those logs to the publicContractsDB
+        tx.unencryptedLogs.addFunctionLogs([enqueuedCallResult.newUnencryptedLogs]);
+      }
+
+      const { output, provingRequest } = await this.runMergeKernelCircuit(
+        publicKernelOutput,
+        enqueuedCallResult.kernelOutput,
+        isFromPrivate,
+      );
+      publicKernelOutput = output;
+      isFromPrivate = false;
+      provingRequests.push(provingRequest);
+    }
+
+    if (phase === PublicKernelPhase.SETUP) {
+      await this.worldStateDB.checkpoint();
+    }
+
+    return {
+      publicKernelOutput,
+      provingRequests,
+      durationMs: phaseTimer.ms(),
+      gasUsed,
+      returnValues: revertReason ? [] : returnValues,
+      revertReason,
+    };
+  }
+
+  private getAvailableGas(
+    tx: Tx,
+    previousPublicKernelOutput: PublicKernelCircuitPublicInputs,
+    phase: PublicKernelPhase,
+  ) {
+    if (phase === PublicKernelPhase.TEARDOWN) {
+      return tx.data.constants.txContext.gasSettings.getTeardownLimits();
+    } else {
+      return tx.data.constants.txContext.gasSettings
+        .getLimits() // No need to subtract teardown limits since they are already included in end.gasUsed
+        .sub(previousPublicKernelOutput.end.gasUsed)
+        .sub(previousPublicKernelOutput.endNonRevertibleData.gasUsed);
+    }
+  }
+
+  private getTransactionFee(
+    tx: Tx,
+    previousPublicKernelOutput: PublicKernelCircuitPublicInputs,
+    phase: PublicKernelPhase,
+  ): Fr {
+    if (phase !== PublicKernelPhase.TEARDOWN) {
+      return Fr.ZERO;
+    } else {
+      const gasSettings = tx.data.constants.txContext.gasSettings;
+      const gasFees = this.globalVariables.gasFees;
+      // No need to add teardown limits since they are already included in end.gasUsed
+      const gasUsed = previousPublicKernelOutput.end.gasUsed.add(
+        previousPublicKernelOutput.endNonRevertibleData.gasUsed,
+      );
+      const txFee = gasSettings.inclusionFee.add(gasUsed.computeFee(gasFees));
+      this.log.debug(`Computed tx fee`, { txFee, gasUsed: inspect(gasUsed), gasFees: inspect(gasFees) });
+      return txFee;
+    }
+  }
+
+  private async runMergeKernelCircuit(
+    previousOutput: PublicKernelCircuitPublicInputs,
+    enqueuedCallData: VMCircuitPublicInputs,
+    isFromPrivate: boolean,
+  ): Promise<{ output: PublicKernelCircuitPublicInputs; provingRequest: PublicKernelMergeRequest }> {
+    const previousKernel = this.getPreviousKernelData(previousOutput, isFromPrivate);
+
+    // The proof is not used in simulation.
+    const vmProof = makeEmptyProof();
+    const callData = new EnqueuedCallData(enqueuedCallData, vmProof);
+
+    const inputs = new PublicKernelCircuitPrivateInputs(previousKernel, callData);
+
+    const output = await this.publicKernelSimulator.publicKernelCircuitMerge(inputs);
+
+    const provingRequest: PublicKernelMergeRequest = {
+      type: ProvingRequestType.PUBLIC_KERNEL_MERGE,
+      inputs,
+    };
+
+    return { output, provingRequest };
+  }
+
+  private getPreviousKernelData(
+    previousOutput: PublicKernelCircuitPublicInputs,
+    isFromPrivate: boolean,
+  ): PublicKernelData {
+    // The proof is not used in simulation.
+    const proof = makeEmptyRecursiveProof(NESTED_RECURSIVE_PROOF_LENGTH);
+
+    const vk = isFromPrivate
+      ? ProtocolCircuitVks.PrivateKernelTailToPublicArtifact
+      : ProtocolCircuitVks.PublicKernelMergeArtifact;
+    const vkIndex = getVKIndex(vk);
+    const siblingPath = getVKSiblingPath(vkIndex);
+
+    return new PublicKernelData(previousOutput, proof, vk, vkIndex, siblingPath);
+  }
+}

package/src/public/execution.ts

@@ -1,4 +1,5 @@
 import {
+  NestedProcessReturnValues,
   type PublicExecutionRequest,
   type SimulationError,
   type UnencryptedFunctionL2Logs,
@@ -13,8 +14,8 @@ import {
   type LogHash,
   type NoteHash,
   type Nullifier,
-  PublicCallRequest,
   PublicCallStackItemCompressed,
+  PublicInnerCallRequest,
   type ReadRequest,
   RevertCode,
   type TreeLeafReadRequest,
@@ -86,10 +87,8 @@ export interface PublicExecutionResult {
    */
   allUnencryptedLogs: UnencryptedFunctionL2Logs;
 
-  // TODO(dbanks12): add contract instance read requests
-
   /** The requests to call public functions made by this call. */
-  publicCallRequests: PublicCallRequest[];
+  publicCallRequests: PublicInnerCallRequest[];
   /** The results of nested calls. */
   nestedExecutions: this[];
 
@@ -101,14 +100,19 @@ export interface PublicExecutionResult {
 }
 
 /**
- * Returns if the input is a public execution result and not just a public execution.
- * @param input - Public execution or public execution result.
- * @returns Whether the input is a public execution result and not just a public execution.
+ * Recursively accummulate the return values of a call result and its nested executions,
+ * so they can be retrieved in order.
+ * @param executionResult
+ * @returns
  */
-export function isPublicExecutionResult(
-  input: PublicExecutionRequest | PublicExecutionResult,
-): input is PublicExecutionResult {
-  return 'executionRequest' in input && input.executionRequest !== undefined;
+export function accumulatePublicReturnValues(executionResult: PublicExecutionResult): NestedProcessReturnValues {
+  const acc = new NestedProcessReturnValues(executionResult.returnValues);
+  acc.nested = executionResult.nestedExecutions.map(nestedExecution => accumulatePublicReturnValues(nestedExecution));
+  return acc;
+}
+
+export function collectExecutionResults(result: PublicExecutionResult): PublicExecutionResult[] {
+  return [result, ...result.nestedExecutions.map(collectExecutionResults)].flat();
 }
 
 /**
@@ -146,5 +150,5 @@ export function resultToPublicCallRequest(result: PublicExecutionResult) {
     Gas.from(result.startGasLeft),
     Gas.from(result.endGasLeft),
   );
-  return new PublicCallRequest(item, result.startSideEffectCounter.toNumber());
+  return new PublicInnerCallRequest(item, result.startSideEffectCounter.toNumber());
 }

package/src/public/executor.ts

@@ -1,6 +1,6 @@
 import { type PublicExecutionRequest } from '@aztec/circuit-types';
 import { type AvmSimulationStats } from '@aztec/circuit-types/stats';
-import { Fr, Gas, type GlobalVariables, type Header, type Nullifier, type TxContext } from '@aztec/circuits.js';
+import { Fr, Gas, type GlobalVariables, type Nullifier, type TxContext } from '@aztec/circuits.js';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
 import { type TelemetryClient } from '@aztec/telemetry-client';
@@ -9,11 +9,11 @@ import { AvmContext } from '../avm/avm_context.js';
 import { AvmExecutionEnvironment } from '../avm/avm_execution_environment.js';
 import { AvmMachineState } from '../avm/avm_machine_state.js';
 import { AvmSimulator } from '../avm/avm_simulator.js';
-import { HostStorage } from '../avm/journal/host_storage.js';
 import { AvmPersistableStateManager } from '../avm/journal/index.js';
-import { type CommitmentsDB, type PublicContractsDB, type PublicStateDB } from './db_interfaces.js';
+import { getPublicFunctionDebugName } from '../common/debug_fn_name.js';
 import { type PublicExecutionResult } from './execution.js';
 import { ExecutorMetrics } from './executor_metrics.js';
+import { type WorldStateDB } from './public_db_sources.js';
 import { PublicSideEffectTrace } from './side_effect_trace.js';
 
 /**
@@ -22,13 +22,7 @@ import { PublicSideEffectTrace } from './side_effect_trace.js';
 export class PublicExecutor {
   metrics: ExecutorMetrics;
 
-  constructor(
-    private readonly publicStorageDB: PublicStateDB,
-    private readonly contractsDb: PublicContractsDB,
-    private readonly commitmentsDb: CommitmentsDB,
-    private readonly header: Header,
-    client: TelemetryClient,
-  ) {
+  constructor(private readonly worldStateDB: WorldStateDB, client: TelemetryClient) {
     this.metrics = new ExecutorMetrics(client, 'PublicExecutor');
   }
 
@@ -56,25 +50,19 @@ export class PublicExecutor {
   ): Promise<PublicExecutionResult> {
     const address = executionRequest.contractAddress;
     const selector = executionRequest.callContext.functionSelector;
-    const fnName = (await this.contractsDb.getDebugFunctionName(address, selector)) ?? `${address}:${selector}`;
+    const fnName = await getPublicFunctionDebugName(this.worldStateDB, address, selector, executionRequest.args);
 
-    PublicExecutor.log.verbose(`[AVM] Executing public external function ${fnName}.`);
+    PublicExecutor.log.verbose(`[AVM] Executing public external function ${fnName}@${address}.`);
     const timer = new Timer();
 
-    const hostStorage = new HostStorage(this.publicStorageDB, this.contractsDb, this.commitmentsDb);
     const trace = new PublicSideEffectTrace(startSideEffectCounter);
     const avmPersistableState = AvmPersistableStateManager.newWithPendingSiloedNullifiers(
-      hostStorage,
+      this.worldStateDB,
       trace,
       pendingSiloedNullifiers.map(n => n.value),
     );
 
-    const avmExecutionEnv = createAvmExecutionEnvironment(
-      executionRequest,
-      this.header,
-      globalVariables,
-      transactionFee,
-    );
+    const avmExecutionEnv = createAvmExecutionEnvironment(executionRequest, globalVariables, transactionFee);
 
     const avmMachineState = new AvmMachineState(availableGas);
     const avmContext = new AvmContext(avmPersistableState, avmExecutionEnv, avmMachineState);
@@ -128,7 +116,6 @@ export class PublicExecutor {
  */
 function createAvmExecutionEnvironment(
   executionRequest: PublicExecutionRequest,
-  header: Header,
   globalVariables: GlobalVariables,
   transactionFee: Fr,
 ): AvmExecutionEnvironment {
@@ -139,7 +126,6 @@ function createAvmExecutionEnvironment(
     executionRequest.callContext.functionSelector,
     /*contractCallDepth=*/ Fr.zero(),
     transactionFee,
-    header,
     globalVariables,
     executionRequest.callContext.isStaticCall,
     executionRequest.callContext.isDelegateCall,

package/src/public/hints_builder.ts

@@ -13,6 +13,7 @@ import {
   NOTE_HASH_TREE_HEIGHT,
   NULLIFIER_TREE_HEIGHT,
   type Nullifier,
+  type NullifierLeafPreimage,
   PUBLIC_DATA_TREE_HEIGHT,
   type PublicDataRead,
   type PublicDataTreeLeafPreimage,
@@ -27,6 +28,7 @@ import {
 } from '@aztec/circuits.js';
 import { makeTuple } from '@aztec/foundation/array';
 import { type Tuple } from '@aztec/foundation/serialize';
+import { type IndexedTreeLeafPreimage } from '@aztec/foundation/trees';
 import { type MerkleTreeOperations } from '@aztec/world-state';
 
 export class HintsBuilder {
@@ -47,15 +49,13 @@ export class HintsBuilder {
     nullifierReadRequests: Tuple<ScopedReadRequest, typeof MAX_NULLIFIER_READ_REQUESTS_PER_TX>,
     pendingNullifiers: Tuple<Nullifier, typeof MAX_NULLIFIERS_PER_TX>,
   ) {
-    return (
-      await buildSiloedNullifierReadRequestHints(
-        this,
-        nullifierReadRequests,
-        pendingNullifiers,
-        MAX_NULLIFIER_READ_REQUESTS_PER_TX,
-        MAX_NULLIFIER_READ_REQUESTS_PER_TX,
-      )
-    ).hints;
+    return await buildSiloedNullifierReadRequestHints(
+      this,
+      nullifierReadRequests,
+      pendingNullifiers,
+      MAX_NULLIFIER_READ_REQUESTS_PER_TX,
+      MAX_NULLIFIER_READ_REQUESTS_PER_TX,
+    );
   }
 
   getNullifierNonExistentReadRequestHints(
@@ -112,7 +112,8 @@ export class HintsBuilder {
       throw new Error(`Nullifier ${nullifier.toBigInt()} already exists in the tree.`);
     }
 
-    return this.getMembershipWitnessWithPreimage<typeof NULLIFIER_TREE_HEIGHT>(
+    // Should find a way to stop casting IndexedTreeLeafPreimage as NullifierLeafPreimage
+    return this.getMembershipWitnessWithPreimage<typeof NULLIFIER_TREE_HEIGHT, NullifierLeafPreimage>(
       MerkleTreeId.NULLIFIER_TREE,
       NULLIFIER_TREE_HEIGHT,
       index,
@@ -125,23 +126,22 @@ export class HintsBuilder {
       throw new Error(`Cannot find the previous value index for public data ${leafSlot}.`);
     }
 
-    const { membershipWitness, leafPreimage } = await this.getMembershipWitnessWithPreimage<
-      typeof PUBLIC_DATA_TREE_HEIGHT
-    >(MerkleTreeId.PUBLIC_DATA_TREE, PUBLIC_DATA_TREE_HEIGHT, res.index);
-
     // Should find a way to stop casting IndexedTreeLeafPreimage as PublicDataTreeLeafPreimage everywhere.
-    return { membershipWitness, leafPreimage: leafPreimage as PublicDataTreeLeafPreimage };
+    return this.getMembershipWitnessWithPreimage<typeof PUBLIC_DATA_TREE_HEIGHT, PublicDataTreeLeafPreimage>(
+      MerkleTreeId.PUBLIC_DATA_TREE,
+      PUBLIC_DATA_TREE_HEIGHT,
+      res.index,
+    );
   }
 
-  private async getMembershipWitnessWithPreimage<TREE_HEIGHT extends number>(
-    treeId: IndexedTreeId,
-    treeHeight: TREE_HEIGHT,
-    index: bigint,
-  ) {
+  private async getMembershipWitnessWithPreimage<
+    TREE_HEIGHT extends number,
+    LEAF_PREIMAGE extends IndexedTreeLeafPreimage = IndexedTreeLeafPreimage,
+  >(treeId: IndexedTreeId, treeHeight: TREE_HEIGHT, index: bigint) {
     const siblingPath = await this.db.getSiblingPath<TREE_HEIGHT>(treeId, index);
     const membershipWitness = new MembershipWitness(treeHeight, index, siblingPath.toTuple());
 
-    const leafPreimage = await this.db.getLeafPreimage(treeId, index);
+    const leafPreimage = (await this.db.getLeafPreimage(treeId, index)) as LEAF_PREIMAGE;
     if (!leafPreimage) {
       throw new Error(`Cannot find the leaf preimage for tree ${treeId} at index ${index}.`);
     }

package/src/public/index.ts

@@ -1,6 +1,6 @@
-export * from './abstract_phase_manager.js';
 export * from './db_interfaces.js';
-export { isPublicExecutionResult, type PublicExecutionResult } from './execution.js';
+export * from './enqueued_calls_processor.js';
+export { type PublicExecutionResult } from './execution.js';
 export { PublicExecutor } from './executor.js';
 export * from './fee_payment.js';
 export { HintsBuilder } from './hints_builder.js';
@@ -8,3 +8,4 @@ export * from './public_db_sources.js';
 export * from './public_kernel.js';
 export * from './public_kernel_circuit_simulator.js';
 export { PublicProcessor, PublicProcessorFactory } from './public_processor.js';
+export { PublicSideEffectTrace } from './side_effect_trace.js';