@aztec/pxe 0.0.1-commit.6c91f13 → 0.0.1-commit.7d4e6cd

package/src/block_synchronizer/block_synchronizer.ts

@@ -7,6 +7,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import type { PXEConfig } from '../config/index.js';
 import type { AnchorBlockStore } from '../storage/anchor_block_store/anchor_block_store.js';
 import type { NoteStore } from '../storage/note_store/note_store.js';
+import type { PrivateEventStore } from '../storage/private_event_store/private_event_store.js';
 
 /**
  * The BlockSynchronizer class orchestrates synchronization between PXE and Aztec node, maintaining an up-to-date
@@ -22,6 +23,7 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
     private node: AztecNode,
     private anchorBlockStore: AnchorBlockStore,
     private noteStore: NoteStore,
+    private privateEventStore: PrivateEventStore,
     private l2TipsStore: L2TipsKVStore,
     config: Partial<Pick<PXEConfig, 'l2BlockBatchSize'>> = {},
     loggerOrSuffix?: string | Logger,
@@ -48,13 +50,13 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
 
     switch (event.type) {
       case 'blocks-added': {
-        const lastBlock = event.blocks.at(-1)!.block;
+        const lastBlock = event.blocks.at(-1)!;
         this.log.verbose(`Updated pxe last block to ${lastBlock.number}`, {
           blockHash: lastBlock.hash(),
           archive: lastBlock.archive.root.toString(),
           header: lastBlock.header.toInspect(),
         });
-        await this.anchorBlockStore.setHeader(lastBlock.getBlockHeader());
+        await this.anchorBlockStore.setHeader(lastBlock.header);
         break;
       }
       case 'chain-pruned': {
@@ -62,6 +64,7 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
         // We first unnullify and then remove so that unnullified notes that were created after the block number end up deleted.
         const lastSynchedBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
         await this.noteStore.rollbackNotesAndNullifiers(event.block.number, lastSynchedBlockNumber);
+        await this.privateEventStore.rollbackEventsAfterBlock(event.block.number, lastSynchedBlockNumber);
         // Update the header to the last block.
         const newHeader = await this.node.getBlockHeader(event.block.number);
         if (!newHeader) {

package/src/config/package_info.ts

@@ -1,3 +1,3 @@
 export function getPackageInfo() {
-  return { version: '3.0.0', name: '@aztec/pxe' };
+  return { version: '4.0.0', name: '@aztec/pxe' };
 }

package/src/contract_function_simulator/contract_function_simulator.ts

@@ -123,6 +123,7 @@ export class ContractFunctionSimulator {
    * @param senderForTags - The address that is used as a tagging sender when emitting private logs. Returned from
    * the `privateGetSenderForTags` oracle.
    * @param scopes - The accounts whose notes we can access in this call. Currently optional and will default to all.
+   * @param jobId - The job ID for staged writes.
    * @returns The result of the execution.
    */
   public async run(
@@ -131,13 +132,14 @@ export class ContractFunctionSimulator {
     selector: FunctionSelector,
     msgSender = AztecAddress.fromField(Fr.MAX_FIELD_VALUE),
     anchorBlockHeader: BlockHeader,
-    senderForTags?: AztecAddress,
-    scopes?: AztecAddress[],
+    senderForTags: AztecAddress | undefined,
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<PrivateExecutionResult> {
     const simulatorSetupTimer = new Timer();
 
     await this.contractStore.syncPrivateState(contractAddress, selector, privateSyncCall =>
-      this.runUtility(privateSyncCall, [], anchorBlockHeader, scopes),
+      this.runUtility(privateSyncCall, [], anchorBlockHeader, scopes, jobId),
     );
 
     await verifyCurrentClassId(contractAddress, this.aztecNode, this.contractStore, anchorBlockHeader);
@@ -174,7 +176,7 @@ export class ContractFunctionSimulator {
       callContext,
       anchorBlockHeader,
       async call => {
-        await this.runUtility(call, [], anchorBlockHeader, scopes);
+        await this.runUtility(call, [], anchorBlockHeader, scopes, jobId);
       },
       request.authWitnesses,
       request.capsules,
@@ -192,6 +194,7 @@ export class ContractFunctionSimulator {
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      jobId,
       0, // totalPublicArgsCount
       startSideEffectCounter,
       undefined, // log
@@ -261,7 +264,8 @@ export class ContractFunctionSimulator {
     call: FunctionCall,
     authwits: AuthWitness[],
     anchorBlockHeader: BlockHeader,
-    scopes?: AztecAddress[],
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<Fr[]> {
     await verifyCurrentClassId(call.to, this.aztecNode, this.contractStore, anchorBlockHeader);
 
@@ -286,6 +290,7 @@ export class ContractFunctionSimulator {
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      jobId,
       undefined,
       scopes,
     );

package/src/contract_function_simulator/noir-structs/event_validation_request.ts

@@ -15,6 +15,7 @@ export class EventValidationRequest {
   constructor(
     public contractAddress: AztecAddress,
     public eventTypeId: EventSelector,
+    public randomness: Fr,
     public serializedEvent: Fr[],
     public eventCommitment: Fr,
     public txHash: TxHash,
@@ -27,6 +28,8 @@ export class EventValidationRequest {
     const contractAddress = AztecAddress.fromField(reader.readField());
     const eventTypeId = EventSelector.fromField(reader.readField());
 
+    const randomness = reader.readField();
+
     const eventStorage = reader.readFieldArray(MAX_EVENT_SERIALIZED_LEN);
     const eventLen = reader.readField().toNumber();
     const serializedEvent = eventStorage.slice(0, eventLen);
@@ -38,6 +41,7 @@ export class EventValidationRequest {
     return new EventValidationRequest(
       contractAddress,
       eventTypeId,
+      randomness,
       serializedEvent,
       eventCommitment,
       txHash,

package/src/contract_function_simulator/oracle/private_execution.ts

@@ -85,7 +85,6 @@ export async function executePrivateFunction(
 
   const rawReturnValues = await privateExecutionOracle.privateLoadFromExecutionCache(publicInputs.returnsHash);
 
-  const noteHashLeafIndexMap = privateExecutionOracle.getNoteHashLeafIndexMap();
   const newNotes = privateExecutionOracle.getNewNotes();
   const noteHashNullifierCounterMap = privateExecutionOracle.getNoteHashNullifierCounterMap();
   const offchainEffects = privateExecutionOracle.getOffchainEffects();
@@ -108,7 +107,6 @@ export async function executePrivateFunction(
     Buffer.from(artifact.verificationKey!, 'base64'),
     partialWitness,
     publicInputs,
-    noteHashLeafIndexMap,
     newNotes,
     noteHashNullifierCounterMap,
     rawReturnValues,

package/src/contract_function_simulator/oracle/private_execution_oracle.ts

@@ -14,7 +14,7 @@ import {
 } from '@aztec/stdlib/abi';
 import type { AuthWitness } from '@aztec/stdlib/auth-witness';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { computeUniqueNoteHash, siloNoteHash, siloNullifier } from '@aztec/stdlib/hash';
+import { siloNullifier } from '@aztec/stdlib/hash';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import { PrivateContextInputs } from '@aztec/stdlib/kernel';
 import { type ContractClassLog, DirectionalAppTaggingSecret, type PreTag } from '@aztec/stdlib/logs';
@@ -64,16 +64,6 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
    * Users can also use this to get a clearer idea of what's happened during a simulation.
    */
   private newNotes: NoteAndSlot[] = [];
-  /**
-   * Notes from previous transactions that are returned to the oracle call `getNotes` during this execution.
-   * The mapping maps from the unique siloed note hash to the index for notes created in private executions.
-   * It maps from siloed note hash to the index for notes created by public functions.
-   *
-   * They are not part of the ExecutionNoteCache and being forwarded to nested contexts via `extend()`
-   * because these notes are meant to be maintained on a per-call basis
-   * They should act as references for the read requests output by an app circuit via public inputs.
-   */
-  private noteHashLeafIndexMap: Map<bigint, bigint> = new Map();
   private noteHashNullifierCounterMap: Map<number, number> = new Map();
   private contractClassLogs: CountedContractClassLog[] = [];
   private offchainEffects: { data: Fr[] }[] = [];
@@ -104,6 +94,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
     senderAddressBookStore: SenderAddressBookStore,
     capsuleStore: CapsuleStore,
     privateEventStore: PrivateEventStore,
+    jobId: string,
     private totalPublicCalldataCount: number = 0,
     protected sideEffectCounter: number = 0,
     log = createLogger('simulator:client_execution_context'),
@@ -126,6 +117,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       senderAddressBookStore,
       capsuleStore,
       privateEventStore,
+      jobId,
       log,
       scopes,
     );
@@ -160,14 +152,6 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
     return toACVMWitness(0, fields);
   }
 
-  /**
-   * The KernelProver will use this to fully populate witnesses and provide hints to the kernel circuit
-   * regarding which note hash each settled read request corresponds to.
-   */
-  public getNoteHashLeafIndexMap() {
-    return this.noteHashLeafIndexMap;
-  }
-
   /**
    * Get the data for the newly created notes.
    */
@@ -402,23 +386,6 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
         .join(', ')}`,
     );
 
-    const noteHashesAndIndexes = await Promise.all(
-      notes.map(async n => {
-        if (n.index !== undefined) {
-          const siloedNoteHash = await siloNoteHash(n.contractAddress, n.noteHash);
-          const uniqueNoteHash = await computeUniqueNoteHash(n.noteNonce, siloedNoteHash);
-
-          return { hash: uniqueNoteHash, index: n.index };
-        }
-      }),
-    );
-
-    noteHashesAndIndexes
-      .filter(n => n !== undefined)
-      .forEach(n => {
-        this.noteHashLeafIndexMap.set(n!.hash.toBigInt(), n!.index);
-      });
-
     return notes;
   }
 
@@ -586,6 +553,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      this.jobId,
       this.totalPublicCalldataCount,
       sideEffectCounter,
       this.log,

package/src/contract_function_simulator/oracle/utility_execution_oracle.ts

@@ -64,6 +64,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     protected readonly senderAddressBookStore: SenderAddressBookStore,
     protected readonly capsuleStore: CapsuleStore,
     protected readonly privateEventStore: PrivateEventStore,
+    protected readonly jobId: string,
     protected log = createLogger('simulator:client_view_context'),
     protected readonly scopes?: AztecAddress[],
   ) {}
@@ -412,6 +413,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       eventService.deliverEvent(
         request.contractAddress,
         request.eventTypeId,
+        request.randomness,
         request.serializedEvent,
         request.eventCommitment,
         request.txHash,

package/src/entrypoints/client/bundle/utils.ts

@@ -1,3 +1,4 @@
+import { BBPrivateKernelProver } from '@aztec/bb-prover/client';
 import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
 import { randomBytes } from '@aztec/foundation/crypto/random';
 import { createLogger } from '@aztec/foundation/log';
@@ -49,7 +50,12 @@ export async function createPXE(
     ? loggers.prover
     : createLogger('pxe:bb:wasm:bundle' + (logSuffix ? `:${logSuffix}` : ''));
 
-  const prover = options.prover ?? new BBBundlePrivateKernelProver(simulator, proverLogger);
+  let prover;
+  if (options.proverOrOptions instanceof BBPrivateKernelProver) {
+    prover = options.proverOrOptions;
+  } else {
+    prover = new BBBundlePrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });
+  }
   const protocolContractsProvider = new BundledProtocolContractsProvider();
 
   const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));

package/src/entrypoints/client/lazy/utils.ts

@@ -1,3 +1,4 @@
+import { BBPrivateKernelProver } from '@aztec/bb-prover/client';
 import { BBLazyPrivateKernelProver } from '@aztec/bb-prover/client/lazy';
 import { randomBytes } from '@aztec/foundation/crypto/random';
 import { createLogger } from '@aztec/foundation/log';
@@ -48,8 +49,12 @@ export async function createPXE(
     ? loggers.prover
     : createLogger('pxe:bb:wasm:bundle' + (logSuffix ? `:${logSuffix}` : ''));
 
-  const prover = options.prover ?? new BBLazyPrivateKernelProver(simulator, proverLogger);
-
+  let prover;
+  if (options.proverOrOptions instanceof BBPrivateKernelProver) {
+    prover = options.proverOrOptions;
+  } else {
+    prover = new BBLazyPrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });
+  }
   const protocolContractsProvider = new LazyProtocolContractsProvider();
 
   const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));

package/src/entrypoints/pxe_creation_options.ts

@@ -1,3 +1,4 @@
+import type { BBPrivateKernelProverOptions } from '@aztec/bb-prover/client';
 import type { Logger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import type { CircuitSimulator } from '@aztec/simulator/client';
@@ -6,7 +7,7 @@ import type { PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
 export type PXECreationOptions = {
   loggers?: { store?: Logger; pxe?: Logger; prover?: Logger };
   useLogSuffix?: boolean | string;
-  prover?: PrivateKernelProver;
+  proverOrOptions?: PrivateKernelProver | BBPrivateKernelProverOptions;
   store?: AztecAsyncKVStore;
   simulator?: CircuitSimulator;
 };

package/src/entrypoints/server/utils.ts

@@ -1,13 +1,10 @@
+import { BBPrivateKernelProver } from '@aztec/bb-prover/client';
 import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
 import { randomBytes } from '@aztec/foundation/crypto/random';
-import { type Logger, createLogger } from '@aztec/foundation/log';
+import { createLogger } from '@aztec/foundation/log';
+import { createStore } from '@aztec/kv-store/lmdb-v2';
 import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import {
-  type CircuitSimulator,
-  MemoryCircuitRecorder,
-  SimulatorRecorderWrapper,
-  WASMSimulator,
-} from '@aztec/simulator/client';
+import { MemoryCircuitRecorder, SimulatorRecorderWrapper, WASMSimulator } from '@aztec/simulator/client';
 import { FileCircuitRecorder } from '@aztec/simulator/testing';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
@@ -46,9 +43,6 @@ export async function createPXE(
   };
 
   if (!options.store) {
-    // TODO once https://github.com/AztecProtocol/aztec-packages/issues/13656 is fixed, we can remove this and always
-    // import the lmdb-v2 version
-    const { createStore } = await import('@aztec/kv-store/lmdb-v2');
     const storeLogger = loggers.store
       ? loggers.store
       : createLogger('pxe:data:lmdb' + (logSuffix ? `:${logSuffix}` : ''));
@@ -58,7 +52,13 @@ export async function createPXE(
     ? loggers.prover
     : createLogger('pxe:bb:native' + (logSuffix ? `:${logSuffix}` : ''));
 
-  const prover = options.prover ?? createProver(simulator, proverLogger);
+  let prover;
+  if (options.proverOrOptions instanceof BBPrivateKernelProver) {
+    prover = options.proverOrOptions;
+  } else {
+    prover = new BBBundlePrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });
+  }
+
   const protocolContractsProvider = new BundledProtocolContractsProvider();
 
   const pxeLogger = loggers.pxe ? loggers.pxe : createLogger('pxe:service' + (logSuffix ? `:${logSuffix}` : ''));
@@ -73,7 +73,3 @@ export async function createPXE(
   );
   return pxe;
 }
-
-function createProver(simulator: CircuitSimulator, logger?: Logger) {
-  return new BBBundlePrivateKernelProver(simulator, logger);
-}

package/src/events/event_service.ts

@@ -19,6 +19,7 @@ export class EventService {
   public async deliverEvent(
     contractAddress: AztecAddress,
     selector: EventSelector,
+    randomness: Fr,
     content: Fr[],
     eventCommitment: Fr,
     txHash: TxHash,
@@ -63,6 +64,7 @@ export class EventService {
 
     return this.privateEventStore.storePrivateEventLog(
       selector,
+      randomness,
       content,
       Number(nullifierIndex.data), // Index of the event commitment in the nullifier tree
       {

package/src/job_coordinator/job_coordinator.ts

@@ -0,0 +1,149 @@
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+
+/**
+ * Interface that stores must implement to support staged writes.
+ */
+export interface StagedStore {
+  /** Unique name identifying this store (used for tracking staged stores from JobCoordinator) */
+  readonly storeName: string;
+
+  /**
+   * Commits staged data to main storage.
+   * Should be called within a transaction for atomicity.
+   *
+   * @param jobId - The job identifier
+   */
+  commit(jobId: string): Promise<void>;
+
+  /**
+   * Discards staged data without committing.
+   * Called on abort.
+   *
+   * @param jobId - The job identifier
+   */
+  discardStaged(jobId: string): Promise<void>;
+}
+
+/**
+ * JobCoordinator manages job lifecycle and provides crash resilience for PXE operations.
+ *
+ * It uses a staged writes pattern:
+ * 1. When a job begins, a unique job ID is created
+ * 2. During the job, all writes go to staging (keyed by job ID)
+ * 3. On commit, staging is promoted to main storage
+ * 4. On abort, staged data is discarded
+ *
+ * Note: PXE should only rely on a single JobCoordinator instance, so it can eventually
+ * orchestrate concurrent jobs. Right now it doesn't make a difference because we're
+ * using a job queue with concurrency=1.
+ */
+export class JobCoordinator {
+  private readonly log = createLogger('pxe:job_coordinator');
+
+  /** The underlying KV store */
+  kvStore: AztecAsyncKVStore;
+
+  #currentJobId: string | undefined;
+  #stores: Map<string, StagedStore> = new Map();
+
+  constructor(kvStore: AztecAsyncKVStore) {
+    this.kvStore = kvStore;
+  }
+
+  /**
+   * Registers a staged store.
+   * Must be called during initialization for all stores that need staging support.
+   */
+  registerStore(store: StagedStore): void {
+    if (this.#stores.has(store.storeName)) {
+      throw new Error(`Store "${store.storeName}" is already registered`);
+    }
+    this.#stores.set(store.storeName, store);
+    this.log.debug(`Registered staged store: ${store.storeName}`);
+  }
+
+  /**
+   * Registers multiple staged stores.
+   */
+  registerStores(stores: StagedStore[]): void {
+    for (const store of stores) {
+      this.registerStore(store);
+    }
+  }
+
+  /**
+   * Begins a new job and returns a job ID for staged writes.
+   *
+   * @returns Job ID to pass to store operations
+   */
+  beginJob(): string {
+    if (this.#currentJobId) {
+      throw new Error(
+        `Cannot begin job: job ${this.#currentJobId} is already in progress. ` +
+          `This should not happen - ensure jobs are properly committed or aborted.`,
+      );
+    }
+
+    const jobId = randomBytes(8).toString('hex');
+    this.#currentJobId = jobId;
+
+    this.log.debug(`Started job ${jobId}`);
+    return jobId;
+  }
+
+  /**
+   * Commits a job by promoting all staged data to main storage.
+   *
+   * @param jobId - The job ID returned from beginJob
+   */
+  async commitJob(jobId: string): Promise<void> {
+    if (!this.#currentJobId || this.#currentJobId !== jobId) {
+      throw new Error(
+        `Cannot commit job ${jobId}: no matching job in progress. ` + `Current job: ${this.#currentJobId ?? 'none'}`,
+      );
+    }
+
+    this.log.debug(`Committing job ${jobId}`);
+
+    // Commit all stores atomically in a single transaction.
+    // Each store's commit is a no-op if it has no staged data (but that's up to each store to handle).
+    await this.kvStore.transactionAsync(async () => {
+      for (const store of this.#stores.values()) {
+        await store.commit(jobId);
+      }
+    });
+
+    this.#currentJobId = undefined;
+    this.log.debug(`Job ${jobId} committed successfully`);
+  }
+
+  /**
+   * Aborts a job by discarding all staged data.
+   *
+   * @param jobId - The job ID returned from beginJob
+   */
+  async abortJob(jobId: string): Promise<void> {
+    if (!this.#currentJobId || this.#currentJobId !== jobId) {
+      // Job may have already been aborted or never started properly
+      this.log.warn(`Abort called for job ${jobId} but current job is ${this.#currentJobId ?? 'none'}`);
+    }
+
+    this.log.debug(`Aborting job ${jobId}`);
+
+    for (const store of this.#stores.values()) {
+      await store.discardStaged(jobId);
+    }
+
+    this.#currentJobId = undefined;
+    this.log.debug(`Job ${jobId} aborted`);
+  }
+
+  /**
+   * Checks if there's a job currently in progress.
+   */
+  hasJobInProgress(): boolean {
+    return this.#currentJobId !== undefined;
+  }
+}

package/src/private_kernel/hints/build_private_kernel_reset_private_inputs.ts

@@ -138,7 +138,7 @@ export class PrivateKernelResetPrivateInputsBuilder {
     }
   }
 
-  async build(oracle: PrivateKernelOracle, noteHashLeafIndexMap: Map<bigint, bigint>) {
+  async build(oracle: PrivateKernelOracle) {
     if (privateKernelResetDimensionNames.every(name => !this.requestedDimensions[name])) {
       throw new Error('Reset is not required.');
     }
@@ -190,7 +190,6 @@ export class PrivateKernelResetPrivateInputsBuilder {
           this.previousKernel.validationRequests.noteHashReadRequests,
           this.previousKernel.end.noteHashes,
           this.noteHashResetActions,
-          noteHashLeafIndexMap,
         ),
         await buildNullifierReadRequestHintsFromResetActions(
           { getNullifierMembershipWitness: getNullifierMembershipWitnessResolver(oracle) },

package/src/private_kernel/private_kernel_execution_prover.ts

@@ -28,7 +28,6 @@ import {
   type PrivateCallExecutionResult,
   type PrivateExecutionResult,
   TxRequest,
-  collectNoteHashLeafIndexMap,
   collectNoteHashNullifierCounterMap,
   getFinalMinRevertibleSideEffectCounter,
 } from '@aztec/stdlib/tx';
@@ -101,7 +100,6 @@ export class PrivateKernelExecutionProver {
 
     const executionSteps: PrivateExecutionStep[] = [];
 
-    const noteHashLeafIndexMap = collectNoteHashLeafIndexMap(executionResult);
     const noteHashNullifierCounterMap = collectNoteHashNullifierCounterMap(executionResult);
     const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(executionResult);
     const splitCounter = isPrivateOnlyTx ? 0 : minRevertibleSideEffectCounter;
@@ -116,7 +114,7 @@ export class PrivateKernelExecutionProver {
         );
         while (resetBuilder.needsReset()) {
           const witgenTimer = new Timer();
-          const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
+          const privateInputs = await resetBuilder.build(this.oracle);
           output = generateWitnesses
             ? await this.proofCreator.generateResetOutput(privateInputs)
             : await this.proofCreator.simulateReset(privateInputs);
@@ -224,7 +222,7 @@ export class PrivateKernelExecutionProver {
     );
     while (resetBuilder.needsReset()) {
       const witgenTimer = new Timer();
-      const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
+      const privateInputs = await resetBuilder.build(this.oracle);
       output = generateWitnesses
         ? await this.proofCreator.generateResetOutput(privateInputs)
         : await this.proofCreator.simulateReset(privateInputs);

package/src/private_kernel/private_kernel_oracle.ts

@@ -43,7 +43,8 @@ export interface PrivateKernelOracle {
 
   /**
    * Returns a membership witness with the sibling path and leaf index in our private function indexed merkle tree.
-   */ getNoteHashMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>>;
+   */
+  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined>;
 
   /**
    * Returns a membership witness with the sibling path and leaf index in our nullifier indexed merkle tree.

package/src/private_kernel/private_kernel_oracle_impl.ts

@@ -2,7 +2,6 @@ import { NOTE_HASH_TREE_HEIGHT, PUBLIC_DATA_TREE_HEIGHT, VK_TREE_HEIGHT } from '
 import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { GrumpkinScalar, Point } from '@aztec/foundation/curves/grumpkin';
 import { createLogger } from '@aztec/foundation/log';
-import type { Tuple } from '@aztec/foundation/serialize';
 import { MembershipWitness } from '@aztec/foundation/trees';
 import type { KeyStore } from '@aztec/key-store';
 import { getVKIndex, getVKSiblingPath } from '@aztec/noir-protocol-circuits-types/vk-tree';
@@ -70,13 +69,8 @@ export class PrivateKernelOracleImpl implements PrivateKernelOracle {
     return Promise.resolve(new MembershipWitness(VK_TREE_HEIGHT, BigInt(leafIndex), getVKSiblingPath(leafIndex)));
   }
 
-  async getNoteHashMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>> {
-    const path = await this.node.getNoteHashSiblingPath(this.blockNumber, leafIndex);
-    return new MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>(
-      path.pathSize,
-      leafIndex,
-      path.toFields() as Tuple<Fr, typeof NOTE_HASH_TREE_HEIGHT>,
-    );
+  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined> {
+    return this.node.getNoteHashMembershipWitness(this.blockNumber, noteHash);
   }
 
   getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitness | undefined> {