@aztec/pxe 0.0.1-commit.d3ec352c → 0.0.1-commit.d58ff9d0

package/src/contract/contract_sync_service.ts

@@ -0,0 +1,176 @@
+import type { Logger } from '@aztec/foundation/log';
+import { Semaphore } from '@aztec/foundation/queue';
+import { isProtocolContract } from '@aztec/protocol-contracts';
+import type { FunctionCall, FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { BlockHeader } from '@aztec/stdlib/tx';
+
+import type { StagedStore } from '../job_coordinator/job_coordinator.js';
+import { NoteService } from '../notes/note_service.js';
+import type { ContractStore } from '../storage/contract_store/contract_store.js';
+import type { NoteStore } from '../storage/note_store/note_store.js';
+import type { ContractClassService } from './contract_class_service.js';
+import { syncScope } from './helpers.js';
+
+/**
+ * Maximum number of scope syncs running concurrently within a single sync call. Sized to trade off parallelism
+ * on non-ACIR work (node RPC, note store reads) against memory pressure from concurrent circuit execution.
+ */
+export const MAX_CONCURRENT_SCOPE_SYNCS = 5;
+
+/**
+ * Service for syncing the private state of contracts. It uses a cache to avoid redundant sync operations - the cache
+ * is wiped when the anchor block changes.
+ *
+ * TODO: The StagedStore naming is broken here. Figure out a better name.
+ */
+export class ContractSyncService implements StagedStore {
+  readonly storeName = 'contract_sync';
+
+  // Tracks contracts synced since last wipe. The cache is keyed per individual scope address
+  // (`contractAddress:scopeAddress`), or `contractAddress:*` for all scopes (all accounts).
+  // The value is a promise that resolves when the contract is synced.
+  private syncedContracts: Map<string, Promise<void>> = new Map();
+
+  constructor(
+    private aztecNode: AztecNode,
+    private contractStore: ContractStore,
+    private contractClassService: ContractClassService,
+    private noteStore: NoteStore,
+    private log: Logger,
+  ) {}
+
+  /**
+   * Ensures a contract's private state is synchronized.
+   * Uses a cache to avoid redundant sync operations - the cache is wiped when the anchor block changes.
+   * @param contractAddress - The address of the contract to sync.
+   * @param functionToInvokeAfterSync - The function selector that will be called after sync (used to validate it's
+   * not sync_state itself).
+   * @param utilityExecutor - Executor function for running the sync_state utility function.
+   * @param scopes - Access scopes to pass through to the utility executor (affects whose account's private state is discovered).
+   */
+  async ensureContractSynced(
+    contractAddress: AztecAddress,
+    functionToInvokeAfterSync: FunctionSelector | null,
+    utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise<any>,
+    anchorBlockHeader: BlockHeader,
+    jobId: string,
+    scopes: AztecAddress[],
+  ): Promise<void> {
+    this.#startSyncIfNeeded(contractAddress, scopes, anchorBlockHeader, jobId, scope =>
+      syncScope(
+        contractAddress,
+        this.contractStore,
+        this.contractClassService,
+        anchorBlockHeader,
+        functionToInvokeAfterSync,
+        utilityExecutor,
+        scope,
+      ),
+    );
+
+    await this.#awaitSync(contractAddress, scopes);
+  }
+
+  /** Clears sync cache entries for the given scopes of a contract. */
+  invalidateContractForScopes(contractAddress: AztecAddress, scopes: AztecAddress[]): void {
+    if (scopes.length === 0) {
+      return;
+    }
+    scopes.forEach(scope => this.syncedContracts.delete(toKey(contractAddress, scope)));
+  }
+
+  /** Clears sync cache. Called by BlockSynchronizer when anchor block changes. */
+  wipe(): void {
+    this.log.debug(`Wiping contract sync cache (${this.syncedContracts.size} entries)`);
+    this.syncedContracts.clear();
+  }
+
+  commit(_jobId: string): Promise<void> {
+    return Promise.resolve();
+  }
+
+  discardStaged(_jobId: string): Promise<void> {
+    // We clear the synced contracts cache here because, when the job is discarded, any associated database writes from
+    // the sync are also undone.
+    this.syncedContracts.clear();
+    return Promise.resolve();
+  }
+
+  /**
+   * For each unsynced scope, creates a promise that waits on:
+   *  1. Note nullifier sync (shared, batched across all unsynced scopes).
+   *  2. Per-scope sync (individual, semaphore-bounded).
+   */
+  #startSyncIfNeeded(
+    contractAddress: AztecAddress,
+    scopes: AztecAddress[],
+    anchorBlockHeader: BlockHeader,
+    jobId: string,
+    syncScopeFn: (scope: AztecAddress) => Promise<void>,
+  ): void {
+    const scopesToSync = scopes.filter(scope => !this.syncedContracts.has(toKey(contractAddress, scope)));
+    if (scopesToSync.length === 0) {
+      return;
+    }
+
+    this.log.debug(`Syncing contract ${contractAddress} for ${scopesToSync.length} scope(s)`);
+
+    const syncNullifiersPromise = this.#syncNoteNullifiers(contractAddress, anchorBlockHeader, jobId, scopesToSync);
+
+    // We build a new semaphore for each sync call, so it rate-limits the scopes within that single call. We do
+    // this so that if these scope syncs trigger nested syncs, the nested ones can execute without causing a deadlock.
+    const syncSlot = new Semaphore(MAX_CONCURRENT_SCOPE_SYNCS);
+
+    for (const scope of scopesToSync) {
+      const key = toKey(contractAddress, scope);
+      const promise = Promise.all([syncNullifiersPromise, runBounded(syncSlot, () => syncScopeFn(scope))])
+        .then(() => {})
+        .catch(err => {
+          this.syncedContracts.delete(key);
+          throw err;
+        });
+      this.syncedContracts.set(key, promise);
+    }
+  }
+
+  /** Syncs note nullifiers across all unsynced scopes in a single batched call. */
+  async #syncNoteNullifiers(
+    contractAddress: AztecAddress,
+    anchorBlockHeader: BlockHeader,
+    jobId: string,
+    scopes: AztecAddress[],
+  ): Promise<void> {
+    // Protocol contracts don't have private state to sync
+    if (isProtocolContract(contractAddress)) {
+      return;
+    }
+    // This runs in parallel with per-scope sync (which also writes to the note store). That's safe because
+    // the note store handles concurrent operations.
+    const noteService = new NoteService(this.noteStore, this.aztecNode, anchorBlockHeader, jobId);
+    await noteService.syncNoteNullifiers(contractAddress, scopes);
+  }
+
+  /** Collects all relevant scope promises (including in-flight ones from concurrent calls) and awaits them. */
+  async #awaitSync(contractAddress: AztecAddress, scopes: AztecAddress[]): Promise<void> {
+    const promises = scopes
+      .map(scope => this.syncedContracts.get(toKey(contractAddress, scope)))
+      .filter(p => p !== undefined);
+    await Promise.all(promises);
+  }
+}
+
+function toKey(contract: AztecAddress, scope: AztecAddress) {
+  return `${contract.toString()}:${scope.toString()}`;
+}
+
+/** Runs fn while holding a slot in the given semaphore, bounding concurrent scope syncs within a single call. */
+async function runBounded<T>(syncSlot: Semaphore, fn: () => Promise<T>): Promise<T> {
+  await syncSlot.acquire();
+  try {
+    return await fn();
+  } finally {
+    syncSlot.release();
+  }
+}

package/src/contract/helpers.ts

@@ -0,0 +1,35 @@
+import type { FunctionCall, FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { BlockHeader } from '@aztec/stdlib/tx';
+
+import type { ContractStore } from '../storage/contract_store/contract_store.js';
+import type { ContractClassService } from './contract_class_service.js';
+import { isSkipSyncContract } from './skip_sync_contracts.js';
+
+export async function syncScope(
+  contractAddress: AztecAddress,
+  contractStore: ContractStore,
+  contractClassService: ContractClassService,
+  anchorBlockHeader: BlockHeader,
+  functionToInvokeAfterSync: FunctionSelector | null,
+  utilityExecutor: (privateSyncCall: FunctionCall, scopes: AztecAddress[]) => Promise<any>,
+  scope: AztecAddress,
+) {
+  // Some canonical contracts hold no private state, so there is nothing to sync (see `skipSyncContracts`).
+  if (isSkipSyncContract(contractAddress)) {
+    return;
+  }
+
+  const classId = await contractClassService.getCurrentClassId(contractAddress, anchorBlockHeader);
+  if (!classId) {
+    throw new Error(`Cannot sync contract ${contractAddress}: its instance is not registered nor published.`);
+  }
+  const syncStateFunctionCall = await contractStore.getFunctionCall('sync_state', [scope], contractAddress, classId);
+  if (functionToInvokeAfterSync && functionToInvokeAfterSync.equals(syncStateFunctionCall.selector)) {
+    throw new Error(
+      'Forbidden `sync_state` invocation. `sync_state` can only be invoked by PXE, manual execution can lead to inconsistencies.',
+    );
+  }
+
+  await utilityExecutor(syncStateFunctionCall, [scope]);
+}

package/src/contract/skip_sync_contracts.ts

@@ -0,0 +1,23 @@
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { STANDARD_AUTH_REGISTRY_ADDRESS } from '@aztec/standard-contracts/auth-registry/constants';
+import { STANDARD_MULTI_CALL_ENTRYPOINT_ADDRESS } from '@aztec/standard-contracts/multi-call-entrypoint/constants';
+import { STANDARD_PUBLIC_CHECKS_ADDRESS } from '@aztec/standard-contracts/public-checks/constants';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+/**
+ * Canonical contracts that hold no private state and are therefore never synced.
+ *
+ * The protocol contracts (registries, fee juice) plus the standard AuthRegistry, MultiCallEntrypoint and PublicChecks
+ * declare no notes and no events, so their macro-generated `sync_state` has nothing to discover.
+ */
+export const skipSyncContracts: AztecAddress[] = [
+  ...Object.values(ProtocolContractAddress),
+  STANDARD_AUTH_REGISTRY_ADDRESS,
+  STANDARD_MULTI_CALL_ENTRYPOINT_ADDRESS,
+  STANDARD_PUBLIC_CHECKS_ADDRESS,
+];
+
+/** Returns whether the given contract should be skipped during private state synchronization. */
+export function isSkipSyncContract(address: AztecAddress): boolean {
+  return skipSyncContracts.some(a => a.equals(address));
+}

package/src/contract_function_simulator/anchored_contract_data.ts

@@ -0,0 +1,74 @@
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import type { FunctionArtifactWithContractName, FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractInstancePreimageWithAddress } from '@aztec/stdlib/contract';
+import type { BlockHeader, ContractOverrides } from '@aztec/stdlib/tx';
+
+import type { ContractClassService } from '../contract/contract_class_service.js';
+import type { ContractStore } from '../storage/contract_store/contract_store.js';
+
+/**
+ * Per-run view of contract data for a single simulation, bound to its anchor block.
+ *
+ * The {@link ContractStore} is pure class-id-keyed storage and has no notion of which class an address runs. This
+ * class bridges that gap: it resolves an address to its current class id (via the {@link ContractClassService}
+ * at the run's anchor block) and then serves artifacts from the store. It is also the single place contract
+ * overrides are applied — when an address is overridden, both its instance and its class id come from the override
+ * rather than the chain, so a simulation can execute different bytecode at that address. Overrides are only set for
+ * `simulateTx` (which skips the kernels), so the override path never reaches proving.
+ */
+export class AnchoredContractData {
+  constructor(
+    private readonly store: ContractStore,
+    private readonly contractClassService: ContractClassService,
+    private readonly anchorBlockHeader: BlockHeader,
+    private readonly overrides?: ContractOverrides,
+  ) {}
+
+  /** Returns the address preimage of the instance at `address`, from the override if any, else from storage. */
+  getContractInstance(address: AztecAddress): Promise<ContractInstancePreimageWithAddress | undefined> {
+    const override = this.overrides?.[address.toString()];
+    if (override) {
+      return Promise.resolve(override.instance);
+    }
+    return this.store.getContractInstance(address);
+  }
+
+  /**
+   * Resolves the class id `address` runs in this simulation: the override's class if overridden, else resolved against
+   * the chain at the anchor block.
+   */
+  getCurrentClassId(address: AztecAddress): Promise<Fr | undefined> {
+    const override = this.overrides?.[address.toString()];
+    if (override) {
+      return Promise.resolve(override.instance.currentContractClassId);
+    }
+    return this.contractClassService.getCurrentClassId(address, this.anchorBlockHeader);
+  }
+
+  async getFunctionArtifact(
+    address: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionArtifactWithContractName | undefined> {
+    const classId = await this.getCurrentClassId(address);
+    return classId ? this.store.getFunctionArtifact(classId, selector) : undefined;
+  }
+
+  async getFunctionArtifactWithDebugMetadata(
+    address: AztecAddress,
+    selector: FunctionSelector,
+  ): Promise<FunctionArtifactWithContractName | undefined> {
+    const classId = await this.getCurrentClassId(address);
+    return classId ? this.store.getFunctionArtifactWithDebugMetadata(classId, selector) : undefined;
+  }
+
+  async getDebugContractName(address: AztecAddress): Promise<string | undefined> {
+    const classId = await this.getCurrentClassId(address);
+    return classId ? this.store.getDebugContractName(classId) : undefined;
+  }
+
+  async getDebugFunctionName(address: AztecAddress, selector: FunctionSelector): Promise<string> {
+    const classId = await this.getCurrentClassId(address);
+    return classId ? this.store.getDebugFunctionName(classId, selector) : `${address}:${selector}`;
+  }
+}

package/src/contract_function_simulator/aztec_node_read_cache.ts

@@ -0,0 +1,89 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { BlockHash, BlockParameter } from '@aztec/stdlib/block';
+import type { AztecNode } from '@aztec/stdlib/interfaces/server';
+import type { TxHash } from '@aztec/stdlib/tx';
+
+/**
+ * Per-execution cache for immutable Aztec node reads.
+ */
+export class AztecNodeReadCache {
+  private readonly cache = new Map<string, Promise<unknown>>();
+
+  constructor(private readonly aztecNode: AztecNode) {}
+
+  /** Fetches a block without reissuing the same node request. */
+  public getBlock(block: BlockParameter) {
+    return this.#cachedRead(`block:${this.#keyPart(block)}`, () => this.aztecNode.getBlock(block));
+  }
+
+  /** Fetches a transaction receipt with its effect attached. */
+  public getTxReceiptWithEffect(txHash: TxHash) {
+    return this.#cachedRead(`tx-receipt-with-effect:${txHash.toString()}`, () =>
+      this.aztecNode.getTxReceipt(txHash, { includeTxEffect: true } as const),
+    );
+  }
+
+  /** Fetches an archive-tree witness for a block hash. */
+  public getBlockHashMembershipWitness(referenceBlock: BlockParameter, blockHash: BlockHash) {
+    return this.#cachedRead(
+      `block-hash-membership-witness:${this.#keyPart(referenceBlock)}:${blockHash.toString()}`,
+      () => this.aztecNode.getBlockHashMembershipWitness(referenceBlock, blockHash),
+    );
+  }
+
+  /** Fetches a public-data-tree witness for a leaf slot. */
+  public getPublicDataWitness(referenceBlock: BlockParameter, leafSlot: Fr) {
+    return this.#cachedRead(`public-data-witness:${this.#keyPart(referenceBlock)}:${leafSlot.toString()}`, () =>
+      this.aztecNode.getPublicDataWitness(referenceBlock, leafSlot),
+    );
+  }
+
+  /** Fetches public storage for a single slot. */
+  public getPublicStorageAt(referenceBlock: BlockParameter, contractAddress: AztecAddress, storageSlot: Fr) {
+    return this.#cachedRead(
+      `public-storage:${this.#keyPart(referenceBlock)}:${contractAddress.toString()}:${storageSlot.toString()}`,
+      () => this.aztecNode.getPublicStorageAt(referenceBlock, contractAddress, storageSlot),
+    );
+  }
+
+  /** Fetches a contiguous public storage range, reusing cached reads for overlapping slots. */
+  public getPublicStorageRange(
+    referenceBlock: BlockParameter,
+    contractAddress: AztecAddress,
+    startStorageSlot: Fr,
+    numberOfElements: number,
+  ) {
+    const slots = Array(numberOfElements)
+      .fill(0)
+      .map((_, i) => new Fr(startStorageSlot.value + BigInt(i)));
+
+    return Promise.all(slots.map(storageSlot => this.getPublicStorageAt(referenceBlock, contractAddress, storageSlot)));
+  }
+
+  #cachedRead<T>(key: string, fetch: () => Promise<T>): Promise<T> {
+    const cached = this.cache.get(key);
+    if (cached) {
+      return cached as Promise<T>;
+    }
+
+    const promise = fetch();
+    promise.catch(() => this.cache.delete(key));
+    this.cache.set(key, promise);
+    return promise;
+  }
+
+  #keyPart(value: unknown): string {
+    if (['string', 'number', 'bigint', 'boolean'].includes(typeof value)) {
+      return String(value);
+    }
+    if (value && typeof value === 'object') {
+      const toString = (value as { toString?: () => string }).toString;
+      if (toString && toString !== Object.prototype.toString) {
+        return toString.call(value);
+      }
+      return JSON.stringify(value, (_key, nested) => (typeof nested === 'bigint' ? nested.toString() : nested));
+    }
+    return String(value);
+  }
+}

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;
+  }
+}