@aztec/pxe 0.0.1-commit.c7c42ec → 0.0.1-commit.f295ac2

package/dest/tree_membership/tree_membership_service.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"tree_membership_service.d.ts","sourceRoot":"","sources":["../../src/tree_membership/tree_membership_service.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,wBAAwB,EAAE,MAAM,kBAAkB,CAAC;AACjE,OAAO,KAAK,EAAE,EAAE,EAAE,MAAM,gCAAgC,CAAC;AACzD,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,yBAAyB,CAAC;AAC3D,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,6BAA6B,CAAC;AAChE,OAAO,KAAK,EAAE,cAAc,EAAE,MAAM,qBAAqB,CAAC;AAC1D,OAAO,KAAK,EAAE,SAAS,EAAE,MAAM,iCAAiC,CAAC;AAEjE,OAAO,EAAE,YAAY,EAAE,0BAA0B,EAAE,iBAAiB,EAAE,MAAM,qBAAqB,CAAC;AAElG,OAAO,KAAK,EAAE,gBAAgB,EAAE,MAAM,qDAAqD,CAAC;AAE5F,qBAAa,qBAAqB;;IAE9B,OAAO,CAAC,QAAQ,CAAC,SAAS;IAC1B,OAAO,CAAC,QAAQ,CAAC,gBAAgB;IAFnC,YACmB,SAAS,EAAE,SAAS,EACpB,gBAAgB,EAAE,gBAAgB,EACjD;IAEJ;;;OAGG;IACI,iBAAiB,CAAC,SAAS,EAAE,EAAE,+BAErC;IAED;;;;;;OAMG;IACU,oBAAoB,CAAC,WAAW,EAAE,cAAc,EAAE,MAAM,EAAE,YAAY,EAAE,SAAS,EAAE,EAAE,GAAG,OAAO,CAAC,EAAE,EAAE,CAAC,CAMjH;IAED;;;;;;;;OAQG;IACU,gCAAgC,CAC3C,WAAW,EAAE,cAAc,EAC3B,SAAS,EAAE,EAAE,GACZ,OAAO,CAAC,0BAA0B,GAAG,SAAS,CAAC,CAMjD;IAED;;;;OAIG;IACU,oBAAoB,CAAC,WAAW,EAAE,cAAc,EAAE,QAAQ,EAAE,EAAE,GAAG,OAAO,CAAC,iBAAiB,GAAG,SAAS,CAAC,CAMnH;IAED;;;;;;OAMG;IACI,0BAA0B,CAC/B,eAAe,EAAE,YAAY,EAC7B,WAAW,EAAE,EAAE,EACf,MAAM,EAAE,EAAE,GACT,OAAO,CAAC,CAAC,MAAM,EAAE,WAAW,CAAC,OAAO,wBAAwB,CAAC,CAAC,CAAC,CAEjE;CAyBF"}
+{"version":3,"file":"tree_membership_service.d.ts","sourceRoot":"","sources":["../../src/tree_membership/tree_membership_service.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,wBAAwB,EAAE,MAAM,kBAAkB,CAAC;AACjE,OAAO,KAAK,EAAE,EAAE,EAAE,MAAM,gCAAgC,CAAC;AACzD,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,yBAAyB,CAAC;AAC3D,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,6BAA6B,CAAC;AAChE,OAAO,EAAE,WAAW,EAAE,MAAM,qBAAqB,CAAC;AAClD,OAAO,KAAK,EAAE,SAAS,EAAE,MAAM,iCAAiC,CAAC;AAEjE,OAAO,EAAE,YAAY,EAAE,0BAA0B,EAAE,iBAAiB,EAAE,MAAM,qBAAqB,CAAC;AAMlG,qBAAa,qBAAqB;;IACpB,OAAO,CAAC,QAAQ,CAAC,SAAS;IAAtC,YAA6B,SAAS,EAAE,SAAS,EAAI;IAErD;;;OAGG;IACU,iBAAiB,CAAC,SAAS,EAAE,EAAE,GAAG,OAAO,CAAC,MAAM,GAAG,SAAS,CAAC,CAGzE;IAED;;;;;;OAMG;IACU,oBAAoB,CAAC,SAAS,EAAE,WAAW,EAAE,MAAM,EAAE,YAAY,EAAE,SAAS,EAAE,EAAE,GAAG,OAAO,CAAC,EAAE,EAAE,CAAC,CAQ5G;IAED;;;;;;;;OAQG;IACI,gCAAgC,CACrC,SAAS,EAAE,WAAW,EACtB,SAAS,EAAE,EAAE,GACZ,OAAO,CAAC,0BAA0B,GAAG,SAAS,CAAC,CAEjD;IAED;;;;OAIG;IACI,oBAAoB,CAAC,SAAS,EAAE,WAAW,EAAE,QAAQ,EAAE,EAAE,GAAG,OAAO,CAAC,iBAAiB,GAAG,SAAS,CAAC,CAExG;IAED;;;;;;OAMG;IACI,0BAA0B,CAC/B,eAAe,EAAE,YAAY,EAC7B,WAAW,EAAE,EAAE,EACf,MAAM,EAAE,EAAE,GACT,OAAO,CAAC,CAAC,MAAM,EAAE,WAAW,CAAC,OAAO,wBAAwB,CAAC,CAAC,CAAC,CAEjE;CAgBF"}

package/dest/tree_membership/tree_membership_service.js

@@ -1,56 +1,53 @@
 import { getNonNullifiedL1ToL2MessageWitness } from '@aztec/stdlib/messaging';
 import { MerkleTreeId } from '@aztec/stdlib/trees';
+// TODO: REmove this class as it no longer seems valuable. The only somewhat real functionality in this class is in
+// the `getMembershipWitness` method that is used by `utilityGetMembershipWitness` oracle by then in Aztec.nr we
+// have helper functions around that oracle for archive and note hash tree and the generic method itself is not used
+// anywhere else - make sense to just have simple specific handlers for the archive and note hash trees and drop this.
 export class TreeMembershipService {
     aztecNode;
-    anchorBlockStore;
-    constructor(aztecNode, anchorBlockStore){
+    constructor(aztecNode){
         this.aztecNode = aztecNode;
-        this.anchorBlockStore = anchorBlockStore;
     }
     /**
    * Gets the index of a nullifier in the nullifier tree.
    * @returns - The index of the nullifier. Undefined if it does not exist in the tree.
-   */ getNullifierIndex(nullifier) {
-        return this.#findLeafIndex('latest', MerkleTreeId.NULLIFIER_TREE, nullifier);
+   */ async getNullifierIndex(nullifier) {
+        const [leafIndex] = await this.aztecNode.findLeavesIndexes('latest', MerkleTreeId.NULLIFIER_TREE, [
+            nullifier
+        ]);
+        return leafIndex?.data;
     }
     /**
    * Fetches the index and sibling path of a leaf at a given block from a given tree.
-   * @param blockNumber - The block number at which to get the membership witness.
+   * @param blockHash - The block hash at which to get the membership witness.
    * @param treeId - Id of the tree to get the sibling path from.
    * @param leafValue - The leaf value
    * @returns The index and sibling path concatenated [index, sibling_path]
-   */ async getMembershipWitness(blockNumber, treeId, leafValue) {
-        const witness = await this.#tryGetMembershipWitness(blockNumber, treeId, leafValue);
+   */ async getMembershipWitness(blockHash, treeId, leafValue) {
+        const witness = await this.#tryGetMembershipWitness(blockHash, treeId, leafValue);
         if (!witness) {
-            throw new Error(`Leaf value ${leafValue} not found in tree ${MerkleTreeId[treeId]} at block ${blockNumber}`);
+            throw new Error(`Leaf value ${leafValue} not found in tree ${MerkleTreeId[treeId]} at block hash ${blockHash.toString()}`);
         }
         return witness;
     }
     /**
    * Returns a low nullifier membership witness for a given nullifier at a given block.
-   * @param blockNumber - The block number at which to get the index.
+   * @param blockHash - The block hash at which to get the index.
    * @param nullifier - Nullifier we try to find the low nullifier witness for.
    * @returns The low nullifier membership witness (if found).
    * @remarks Low nullifier witness can be used to perform a nullifier non-inclusion proof by leveraging the "linked
    * list structure" of leaves and proving that a lower nullifier is pointing to a bigger next value than the nullifier
    * we are trying to prove non-inclusion for.
-   */ async getLowNullifierMembershipWitness(blockNumber, nullifier) {
-        const anchorBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
-        if (blockNumber !== 'latest' && blockNumber > anchorBlockNumber) {
-            throw new Error(`Block number ${blockNumber} is higher than current block ${anchorBlockNumber}`);
-        }
-        return this.aztecNode.getLowNullifierMembershipWitness(blockNumber, nullifier);
+   */ getLowNullifierMembershipWitness(blockHash, nullifier) {
+        return this.aztecNode.getLowNullifierMembershipWitness(blockHash, nullifier);
     }
     /**
    * Returns a witness for a given slot of the public data tree at a given block.
-   * @param blockNumber - The block number at which to get the witness.
+   * @param blockHash - The block hash at which to get the witness.
    * @param leafSlot - The slot of the public data in the public data tree.
-   */ async getPublicDataWitness(blockNumber, leafSlot) {
-        const anchorBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
-        if (blockNumber !== 'latest' && blockNumber > anchorBlockNumber) {
-            throw new Error(`Block number ${blockNumber} is higher than current block ${anchorBlockNumber}`);
-        }
-        return await this.aztecNode.getPublicDataWitness(blockNumber, leafSlot);
+   */ getPublicDataWitness(blockHash, leafSlot) {
+        return this.aztecNode.getPublicDataWitness(blockHash, leafSlot);
     }
     /**
    * Looks for the L1 to L2 membership witness of a message at the Aztec node, given its hash.
@@ -61,24 +58,18 @@ export class TreeMembershipService {
    */ getL1ToL2MembershipWitness(contractAddress, messageHash, secret) {
         return getNonNullifiedL1ToL2MessageWitness(this.aztecNode, contractAddress, messageHash, secret);
     }
-    async #tryGetMembershipWitness(blockNumber, treeId, value) {
+    async #tryGetMembershipWitness(blockHash, treeId, value) {
         switch(treeId){
             case MerkleTreeId.NULLIFIER_TREE:
-                return (await this.aztecNode.getNullifierMembershipWitness(blockNumber, value))?.withoutPreimage().toFields();
+                return (await this.aztecNode.getNullifierMembershipWitness(blockHash, value))?.withoutPreimage().toFields();
             case MerkleTreeId.NOTE_HASH_TREE:
-                return (await this.aztecNode.getNoteHashMembershipWitness(blockNumber, value))?.toFields();
+                return (await this.aztecNode.getNoteHashMembershipWitness(blockHash, value))?.toFields();
             case MerkleTreeId.PUBLIC_DATA_TREE:
-                return (await this.aztecNode.getPublicDataWitness(blockNumber, value))?.withoutPreimage().toFields();
+                return (await this.aztecNode.getPublicDataWitness(blockHash, value))?.withoutPreimage().toFields();
             case MerkleTreeId.ARCHIVE:
-                return (await this.aztecNode.getArchiveMembershipWitness(blockNumber, value))?.toFields();
+                return (await this.aztecNode.getArchiveMembershipWitness(blockHash, value))?.toFields();
             default:
                 throw new Error('Not implemented');
         }
     }
-    async #findLeafIndex(blockNumber, treeId, leafValue) {
-        const [leafIndex] = await this.aztecNode.findLeavesIndexes(blockNumber, treeId, [
-            leafValue
-        ]);
-        return leafIndex?.data;
-    }
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aztec/pxe",
-  "version": "0.0.1-commit.c7c42ec",
+  "version": "0.0.1-commit.f295ac2",
   "type": "module",
   "exports": {
     "./server": "./dest/entrypoints/server/index.js",
@@ -11,7 +11,7 @@
   },
   "bin": "./dest/bin/index.js",
   "scripts": {
-    "build": "yarn clean && yarn generate && tsgo -b",
+    "build": "yarn clean && yarn generate && ../scripts/tsc.sh",
     "build:dev": "../scripts/tsc.sh --watch",
     "clean": "rm -rf ./dest .tsbuildinfo ./src/config/package_info.ts",
     "test": "NODE_NO_WARNINGS=1 node --experimental-vm-modules ../node_modules/.bin/jest --passWithNoTests --maxWorkers=${JEST_MAX_WORKERS:-8}",
@@ -61,19 +61,19 @@
     ]
   },
   "dependencies": {
-    "@aztec/bb-prover": "0.0.1-commit.c7c42ec",
-    "@aztec/bb.js": "0.0.1-commit.c7c42ec",
-    "@aztec/builder": "0.0.1-commit.c7c42ec",
-    "@aztec/constants": "0.0.1-commit.c7c42ec",
-    "@aztec/ethereum": "0.0.1-commit.c7c42ec",
-    "@aztec/foundation": "0.0.1-commit.c7c42ec",
-    "@aztec/key-store": "0.0.1-commit.c7c42ec",
-    "@aztec/kv-store": "0.0.1-commit.c7c42ec",
-    "@aztec/noir-protocol-circuits-types": "0.0.1-commit.c7c42ec",
-    "@aztec/noir-types": "0.0.1-commit.c7c42ec",
-    "@aztec/protocol-contracts": "0.0.1-commit.c7c42ec",
-    "@aztec/simulator": "0.0.1-commit.c7c42ec",
-    "@aztec/stdlib": "0.0.1-commit.c7c42ec",
+    "@aztec/bb-prover": "0.0.1-commit.f295ac2",
+    "@aztec/bb.js": "0.0.1-commit.f295ac2",
+    "@aztec/builder": "0.0.1-commit.f295ac2",
+    "@aztec/constants": "0.0.1-commit.f295ac2",
+    "@aztec/ethereum": "0.0.1-commit.f295ac2",
+    "@aztec/foundation": "0.0.1-commit.f295ac2",
+    "@aztec/key-store": "0.0.1-commit.f295ac2",
+    "@aztec/kv-store": "0.0.1-commit.f295ac2",
+    "@aztec/noir-protocol-circuits-types": "0.0.1-commit.f295ac2",
+    "@aztec/noir-types": "0.0.1-commit.f295ac2",
+    "@aztec/protocol-contracts": "0.0.1-commit.f295ac2",
+    "@aztec/simulator": "0.0.1-commit.f295ac2",
+    "@aztec/stdlib": "0.0.1-commit.f295ac2",
     "koa": "^2.16.1",
     "koa-router": "^13.1.1",
     "lodash.omit": "^4.5.0",
@@ -82,14 +82,14 @@
     "viem": "npm:@aztec/viem@2.38.2"
   },
   "devDependencies": {
-    "@aztec/merkle-tree": "0.0.1-commit.c7c42ec",
-    "@aztec/noir-test-contracts.js": "0.0.1-commit.c7c42ec",
+    "@aztec/merkle-tree": "0.0.1-commit.f295ac2",
+    "@aztec/noir-test-contracts.js": "0.0.1-commit.f295ac2",
     "@jest/globals": "^30.0.0",
     "@types/jest": "^30.0.0",
     "@types/lodash.omit": "^4.5.7",
     "@types/lodash.times": "^4.3.9",
     "@types/node": "^22.15.17",
-    "@typescript/native-preview": "7.0.0-dev.20251126.1",
+    "@typescript/native-preview": "7.0.0-dev.20260113.1",
     "jest": "^30.0.0",
     "jest-mock-extended": "^4.0.0",
     "lodash.times": "^4.3.2",

package/src/block_synchronizer/block_synchronizer.ts

@@ -1,12 +1,19 @@
 import { BlockNumber } from '@aztec/foundation/branded-types';
 import { type Logger, createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import type { L2TipsKVStore } from '@aztec/kv-store/stores';
-import { L2BlockStream, type L2BlockStreamEvent, type L2BlockStreamEventHandler } from '@aztec/stdlib/block';
+import {
+  L2BlockHash,
+  L2BlockStream,
+  type L2BlockStreamEvent,
+  type L2BlockStreamEventHandler,
+} from '@aztec/stdlib/block';
 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
@@ -20,8 +27,10 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
 
   constructor(
     private node: AztecNode,
+    private store: AztecAsyncKVStore,
     private anchorBlockStore: AnchorBlockStore,
     private noteStore: NoteStore,
+    private privateEventStore: PrivateEventStore,
     private l2TipsStore: L2TipsKVStore,
     config: Partial<Pick<PXEConfig, 'l2BlockBatchSize'>> = {},
     loggerOrSuffix?: string | Logger,
@@ -48,27 +57,36 @@ 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': {
         this.log.warn(`Pruning data after block ${event.block.number} due to reorg`);
-        // 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);
-        // Update the header to the last block.
-        const newHeader = await this.node.getBlockHeader(event.block.number);
-        if (!newHeader) {
-          this.log.error(`Block header not found for block number ${event.block.number} during chain prune`);
-        } else {
-          await this.anchorBlockStore.setHeader(newHeader);
+
+        const oldAnchorBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
+        // Note that the following is not necessarily the anchor block that will be used in the transaction - if
+        // the chain has already moved past the reorg, we'll also see blocks-added events that will push the anchor
+        // forward.
+        const newAnchorBlockHeader = await this.node.getBlockHeader(L2BlockHash.fromString(event.block.hash));
+
+        if (!newAnchorBlockHeader) {
+          throw new Error(
+            `Block header for block number ${event.block.number} and hash ${event.block.hash} not found during chain prune. This likely indicates a bug in the node, as we receive block stream events and fetch block headers from the same node.`,
+          );
         }
+
+        // Operations are wrapped in a single transaction to ensure atomicity.
+        await this.store.transactionAsync(async () => {
+          await this.noteStore.rollback(event.block.number, oldAnchorBlockNumber);
+          await this.privateEventStore.rollback(event.block.number, oldAnchorBlockNumber);
+          await this.anchorBlockStore.setHeader(newAnchorBlockHeader);
+        });
         break;
       }
     }

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/benchmarked_node.ts

@@ -0,0 +1,103 @@
+import { Timer } from '@aztec/foundation/timer';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { NodeStats, RoundTripStats } from '@aztec/stdlib/tx';
+
+/*
+ * Proxy generator for an AztecNode that tracks the time taken for each RPC call and the number of round trips (actual
+ * blocking waits for node responses).
+ *
+ * A round trip is counted when we transition from 0 to 1 in-flight calls, and ends when all concurrent calls complete.
+ * This means parallel calls in Promise.all count as a single round trip.
+ *
+ * Note that batching of RPC calls in `safe_json_rpc_client.ts` could affect the round trip counts but in places we
+ * currently use this information we do not even use HTTP as we have direct access to the Aztec Node instance in TS
+ * (i.e. not running against external node) so this is not a problem for now.
+ *
+ * If you want to use this against external node and the info gets skewed by batching you can set the `maxBatchSize`
+ * value in `safe_json_rpc_client.ts` to 1 (the main motivation for batching was to get around parallel http requests
+ * limits in web browsers which is not a problem when debugging in node.js).
+ */
+export type BenchmarkedNode = AztecNode & { getStats(): NodeStats };
+
+export class BenchmarkedNodeFactory {
+  static create(node: AztecNode): BenchmarkedNode {
+    // Per-method call stats
+    const perMethod: Partial<Record<keyof AztecNode, { times: number[] }>> = {};
+
+    // Round trip tracking
+    let inFlightCount = 0;
+    let currentRoundTripTimer: Timer | null = null;
+    let currentRoundTripMethods: string[] = [];
+    const roundTrips: RoundTripStats = {
+      roundTrips: 0,
+      totalBlockingTime: 0,
+      roundTripDurations: [],
+      roundTripMethods: [],
+    };
+
+    return new Proxy(node, {
+      get(target, prop: keyof BenchmarkedNode) {
+        if (prop === 'getStats') {
+          return (): NodeStats => {
+            return { perMethod, roundTrips };
+          };
+        } else {
+          return function (...args: any[]) {
+            // Track per-method stats
+            if (!perMethod[prop]) {
+              perMethod[prop] = { times: [] };
+            }
+
+            // Start of a new round trip batch?
+            if (inFlightCount === 0) {
+              roundTrips.roundTrips++;
+              currentRoundTripTimer = new Timer();
+              currentRoundTripMethods = [];
+            }
+            inFlightCount++;
+            currentRoundTripMethods.push(prop);
+
+            const callTimer = new Timer();
+            const result = (target[prop] as any).apply(target, args);
+
+            // Handle completion - called when the call finishes (after Promise resolves)
+            const handleCompletion = () => {
+              const callTime = callTimer.ms();
+              perMethod[prop]!.times.push(callTime);
+
+              inFlightCount--;
+
+              // End of round trip batch - all concurrent calls completed
+              if (inFlightCount === 0 && currentRoundTripTimer) {
+                const roundTripTime = currentRoundTripTimer.ms();
+                roundTrips.totalBlockingTime += roundTripTime;
+                roundTrips.roundTripDurations.push(roundTripTime);
+                roundTrips.roundTripMethods.push(currentRoundTripMethods);
+                currentRoundTripTimer = null;
+                currentRoundTripMethods = [];
+              }
+            };
+
+            // If the result is a Promise, chain the completion handler
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  handleCompletion();
+                  return value;
+                },
+                (error: any) => {
+                  handleCompletion();
+                  throw error;
+                },
+              );
+            } else {
+              // Synchronous method - handle completion immediately
+              handleCompletion();
+              return result;
+            }
+          };
+        }
+      },
+    }) as BenchmarkedNode;
+  }
+}

package/src/contract_function_simulator/contract_function_simulator.ts

@@ -80,6 +80,7 @@ import type { PrivateEventStore } from '../storage/private_event_store/private_e
 import type { RecipientTaggingStore } from '../storage/tagging_store/recipient_tagging_store.js';
 import type { SenderAddressBookStore } from '../storage/tagging_store/sender_address_book_store.js';
 import type { SenderTaggingStore } from '../storage/tagging_store/sender_tagging_store.js';
+import type { BenchmarkedNode } from './benchmarked_node.js';
 import { ExecutionNoteCache } from './execution_note_cache.js';
 import { ExecutionTaggingIndexCache } from './execution_tagging_index_cache.js';
 import { HashedValuesCache } from './hashed_values_cache.js';
@@ -87,7 +88,6 @@ import { Oracle } from './oracle/oracle.js';
 import { executePrivateFunction, verifyCurrentClassId } from './oracle/private_execution.js';
 import { PrivateExecutionOracle } from './oracle/private_execution_oracle.js';
 import { UtilityExecutionOracle } from './oracle/utility_execution_oracle.js';
-import type { ProxiedNode } from './proxied_node.js';
 
 /**
  * The contract function simulator.
@@ -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,11 +132,16 @@ 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, jobId),
+    );
+
     await verifyCurrentClassId(contractAddress, this.aztecNode, this.contractStore, anchorBlockHeader);
 
     const entryPointArtifact = await this.contractStore.getFunctionArtifactWithDebugMetadata(contractAddress, selector);
@@ -169,6 +175,9 @@ export class ContractFunctionSimulator {
       request.txContext,
       callContext,
       anchorBlockHeader,
+      async call => {
+        await this.runUtility(call, [], anchorBlockHeader, scopes, jobId);
+      },
       request.authWitnesses,
       request.capsules,
       HashedValuesCache.create(request.argsOfCalls),
@@ -185,6 +194,7 @@ export class ContractFunctionSimulator {
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      jobId,
       0, // totalPublicArgsCount
       startSideEffectCounter,
       undefined, // log
@@ -208,8 +218,9 @@ export class ContractFunctionSimulator {
         request.functionSelector,
       );
       const simulatorTeardownTimer = new Timer();
-      const { usedProtocolNullifierForNonces } = noteCache.finish();
-      const firstNullifierHint = usedProtocolNullifierForNonces ? Fr.ZERO : noteCache.getAllNullifiers()[0];
+
+      noteCache.finish();
+      const firstNullifierHint = noteCache.getNonceGenerator();
 
       const publicCallRequests = collectNested([executionResult], r =>
         r.publicInputs.publicCallRequests
@@ -254,7 +265,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);
 
@@ -279,6 +291,7 @@ export class ContractFunctionSimulator {
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      jobId,
       undefined,
       scopes,
     );
@@ -319,7 +332,12 @@ export class ContractFunctionSimulator {
    */
   getStats() {
     const nodeRPCCalls =
-      typeof (this.aztecNode as ProxiedNode).getStats === 'function' ? (this.aztecNode as ProxiedNode).getStats() : {};
+      typeof (this.aztecNode as BenchmarkedNode).getStats === 'function'
+        ? (this.aztecNode as BenchmarkedNode).getStats()
+        : {
+            perMethod: {},
+            roundTrips: { roundTrips: 0, totalBlockingTime: 0, roundTripDurations: [], roundTripMethods: [] },
+          };
 
     return { nodeRPCCalls };
   }
@@ -342,15 +360,15 @@ class OrderedSideEffect<T> {
  * (allowing state overrides) and is much faster, while still generating a valid
  * output that can be sent to the node for public simulation
  * @param privateExecutionResult - The result of the private execution.
- * @param nonceGenerator - A nonce generator for note hashes. According to the protocol rules,
- * it can either be the first nullifier in the tx or the hash of the initial tx request if there are none.
  * @param contractStore - A provider for contract data in order to get function names and debug info.
+ * @param minRevertibleSideEffectCounterOverride - Optional override for the min revertible side effect counter.
+ * Used by TXE to simulate account contract behavior (setting the counter before app execution).
  * @returns The simulated proving result.
  */
 export async function generateSimulatedProvingResult(
   privateExecutionResult: PrivateExecutionResult,
-  nonceGenerator: Fr,
   contractStore: ContractStore,
+  minRevertibleSideEffectCounterOverride?: number,
 ): Promise<PrivateKernelExecutionProofOutput<PrivateKernelTailCircuitPublicInputs>> {
   const siloedNoteHashes: OrderedSideEffect<Fr>[] = [];
   const nullifiers: OrderedSideEffect<Fr>[] = [];
@@ -454,16 +472,18 @@ export async function generateSimulatedProvingResult(
   const getEffect = <T>(orderedSideEffect: OrderedSideEffect<T>) => orderedSideEffect.sideEffect;
 
   const isPrivateOnlyTx = privateExecutionResult.publicFunctionCalldata.length === 0;
-  const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(privateExecutionResult);
+  const minRevertibleSideEffectCounter =
+    minRevertibleSideEffectCounterOverride ?? getFinalMinRevertibleSideEffectCounter(privateExecutionResult);
 
   const [nonRevertibleNullifiers, revertibleNullifiers] = splitOrderedSideEffects(
     nullifiers.sort(sortByCounter),
     minRevertibleSideEffectCounter,
   );
-  if (nonRevertibleNullifiers.length > 0 && !nonRevertibleNullifiers[0].equals(nonceGenerator)) {
+  const nonceGenerator = privateExecutionResult.firstNullifier;
+  if (nonRevertibleNullifiers.length === 0) {
+    nonRevertibleNullifiers.push(nonceGenerator);
+  } else if (!nonRevertibleNullifiers[0].equals(nonceGenerator)) {
     throw new Error('The first non revertible nullifier should be equal to the nonce generator. This is a bug!');
-  } else {
-    nonRevertibleNullifiers.unshift(nonceGenerator);
   }
 
   if (isPrivateOnlyTx) {
@@ -499,6 +519,12 @@ export async function generateSimulatedProvingResult(
       siloedNoteHashes.sort(sortByCounter),
       minRevertibleSideEffectCounter,
     );
+    const nonRevertibleUniqueNoteHashes = await Promise.all(
+      nonRevertibleNoteHashes.map(async (noteHash, i) => {
+        const nonce = await computeNoteHashNonce(nonceGenerator, i);
+        return await computeUniqueNoteHash(nonce, noteHash);
+      }),
+    );
     const [nonRevertibleL2ToL1Messages, revertibleL2ToL1Messages] = splitOrderedSideEffects(
       l2ToL1Messages.sort(sortByCounter),
       minRevertibleSideEffectCounter,
@@ -517,7 +543,7 @@ export async function generateSimulatedProvingResult(
     );
 
     const nonRevertibleData = new PrivateToPublicAccumulatedData(
-      padArrayEnd(nonRevertibleNoteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
+      padArrayEnd(nonRevertibleUniqueNoteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
       padArrayEnd(nonRevertibleNullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX),
       padArrayEnd(nonRevertibleL2ToL1Messages, ScopedL2ToL1Message.empty(), MAX_L2_TO_L1_MSGS_PER_TX),
       padArrayEnd(nonRevertibleTaggedPrivateLogs, PrivateLog.empty(), MAX_PRIVATE_LOGS_PER_TX),
@@ -565,7 +591,7 @@ function splitOrderedSideEffects<T>(effects: OrderedSideEffect<T>[], minRevertib
   const revertibleSideEffects: T[] = [];
   const nonRevertibleSideEffects: T[] = [];
   effects.forEach(effect => {
-    if (effect.counter < minRevertibleSideEffectCounter) {
+    if (minRevertibleSideEffectCounter === 0 || effect.counter < minRevertibleSideEffectCounter) {
       nonRevertibleSideEffects.push(effect.sideEffect);
     } else {
       revertibleSideEffects.push(effect.sideEffect);