@aztec/pxe 5.2.0 → 5.3.0-nightly.20260818

package/src/contract/contract_sync_service.ts

@@ -2,15 +2,17 @@ import type { Logger } from '@aztec/foundation/log';
 import { allToCompletion } from '@aztec/foundation/promise';
 import { Semaphore } from '@aztec/foundation/queue';
 import { isProtocolContract } from '@aztec/protocol-contracts';
-import type { FunctionCall, FunctionSelector } from '@aztec/stdlib/abi';
+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 { ContractSyncConfig } from '../config/index.js';
 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 CallKey, ContractCallGraph, type ContractFunction, toCallKey } from './contract_call_graph.js';
 import type { ContractClassService } from './contract_class_service.js';
 import { syncScope } from './helpers.js';
 
@@ -20,6 +22,12 @@ import { syncScope } from './helpers.js';
  */
 export const MAX_CONCURRENT_SCOPE_SYNCS = 5;
 
+/**
+ * Selector of the macro-generated `sync_state` utility function, which is the same for every contract.
+ * Pinned against a compiled artifact in tests, so a macro signature change fails there.
+ */
+export const SYNC_STATE_SELECTOR = FunctionSelector.fromString('0x418ef5da');
+
 /**
  * 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.
@@ -30,9 +38,14 @@ 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();
+  // (`contractAddress:scopeAddress`). The value is a promise that resolves when the contract is synced.
+  private readonly syncedContracts: Map<SyncKey, Promise<void>> = new Map();
+
+  // Per-job speculation state, dropped when the job commits or discards.
+  private readonly speculationByJob: Map<JobId, JobSpeculation> = new Map();
+
+  // Predicts a function's callees from the calls observed in past jobs, driving speculative sync.
+  private readonly callGraph: ContractCallGraph;
 
   constructor(
     private aztecNode: AztecNode,
@@ -40,38 +53,63 @@ export class ContractSyncService implements StagedStore {
     private contractClassService: ContractClassService,
     private noteStore: NoteStore,
     private log: Logger,
-  ) {}
+    { concurrentContractSyncEnabled }: ContractSyncConfig,
+  ) {
+    this.callGraph = new ContractCallGraph(concurrentContractSyncEnabled);
+  }
 
   /**
    * 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,
-      ),
+  async ensureContractSynced({
+    contract,
+    functionToInvokeAfterSync,
+    utilityExecutor,
+    anchorBlockHeader,
+    jobId,
+    scopes,
+    triggeredBy,
+  }: ContractSyncRequest): Promise<void> {
+    // A call is recorded only when both functions are known: the invoked callee and the caller that triggered it.
+    if (functionToInvokeAfterSync && triggeredBy) {
+      this.callGraph.recordCall({
+        jobId,
+        caller: triggeredBy,
+        callee: { address: contract, selector: functionToInvokeAfterSync },
+      });
+    }
+
+    await this.#startSyncIfNeeded(
+      contract,
+      functionToInvokeAfterSync,
+      utilityExecutor,
+      anchorBlockHeader,
+      jobId,
+      scopes,
     );
+  }
 
-    await this.#awaitSync(contractAddress, scopes);
+  /**
+   * Waits until every speculative sync the job fired has finished, then rejects if any failed, so the job discards
+   * instead of committing. This is needed because a sync that fails midway can leave partial staged writes, and a
+   * speculative failure might not be surfaced by any request.
+   */
+  async settle(jobId: JobId): Promise<void> {
+    // A speculative sync's execution can fire more speculative syncs mid-drain, so loop until no new promises
+    // appear, and only escalate once nothing is still writing.
+    const { syncs } = this.#speculationForJob(jobId);
+    const failures: unknown[] = [];
+    while (syncs.length > 0) {
+      const results = await Promise.allSettled(syncs.splice(0));
+      failures.push(...results.filter(result => result.status === 'rejected').map(result => result.reason));
+    }
+    if (failures.length > 0) {
+      throw new AggregateError(
+        failures,
+        'Speculative syncs failed, so the job must discard its staged writes instead of committing',
+      );
+    }
   }
 
   /** Clears sync cache entries for the given scopes of a contract. */
@@ -88,14 +126,18 @@ export class ContractSyncService implements StagedStore {
     this.syncedContracts.clear();
   }
 
-  commit(_jobId: string): Promise<void> {
+  commit(jobId: JobId): Promise<void> {
+    this.callGraph.commitJob(jobId);
+    this.speculationByJob.delete(jobId);
     return Promise.resolve();
   }
 
-  discardStaged(_jobId: string): Promise<void> {
+  discardStaged(jobId: JobId): 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();
+    this.callGraph.discardJob(jobId);
+    this.speculationByJob.delete(jobId);
     return Promise.resolve();
   }
 
@@ -103,38 +145,113 @@ export class ContractSyncService implements StagedStore {
    * 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).
+   * When concurrent contract sync is enabled, the predicted direct callees of the invoked function and of the
+   * contract's `sync_state` start speculatively too, once the contract's own syncs have started (see
+   * {@link #speculativelySync}).
+   * @returns A promise that resolves once every requested scope is synced, including syncs already in flight from
+   * concurrent calls. Speculative syncs are not included: those are only awaited by a later request that needs
+   * their contract, or by the job's {@link settle}.
    */
-  #startSyncIfNeeded(
+  async #startSyncIfNeeded(
     contractAddress: AztecAddress,
+    functionToInvokeAfterSync: FunctionSelector | null,
+    utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise<any>,
+    anchorBlockHeader: BlockHeader,
+    jobId: JobId,
     scopes: AztecAddress[],
+  ): Promise<void> {
+    const scopesToSync = scopes.filter(scope => !this.syncedContracts.has(toKey(contractAddress, scope)));
+    if (scopesToSync.length > 0) {
+      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 syncScopePromise = runBounded(syncSlot, () =>
+          syncScope(
+            contractAddress,
+            this.contractStore,
+            this.contractClassService,
+            anchorBlockHeader,
+            functionToInvokeAfterSync,
+            utilityExecutor,
+            scope,
+          ),
+        );
+        // This cached promise is what every later request for this scope awaits, and both branches write staged data,
+        // so it must run both to completion even when one fails.
+        const promise = allToCompletion([syncNullifiersPromise, syncScopePromise])
+          .then(() => {})
+          .catch(err => {
+            this.syncedContracts.delete(key);
+            throw err;
+          });
+        this.syncedContracts.set(key, promise);
+      }
+    }
+
+    // `sync_state` itself calls other contracts (e.g. most contract syncs query the handshake registry), so its
+    // predicted callees start syncing alongside the contract's own syncs.
+    this.#speculativelySync(contractAddress, SYNC_STATE_SELECTOR, utilityExecutor, anchorBlockHeader, jobId, scopes);
+    this.#speculativelySync(
+      contractAddress,
+      functionToInvokeAfterSync,
+      utilityExecutor,
+      anchorBlockHeader,
+      jobId,
+      scopes,
+    );
+
+    await this.#awaitSync(contractAddress, scopes);
+  }
+
+  /**
+   * Starts the syncs of the contracts the given function is predicted to call (see {@link ContractCallGraph} for how
+   * predictions are learned). Each started sync speculates from its own function in turn, so the whole predicted call
+   * tree syncs in parallel with the contract instead of one contract at a time as execution reaches it.
+   *
+   * A wrong prediction is cheap: the extra node requests are batched into round trips the job already makes, and the
+   * synced data simply goes unused.
+   */
+  #speculativelySync(
+    contractAddress: AztecAddress,
+    functionToInvokeAfterSync: FunctionSelector | null,
+    utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise<any>,
     anchorBlockHeader: BlockHeader,
-    jobId: string,
-    syncScopeFn: (scope: AztecAddress) => Promise<void>,
+    jobId: JobId,
+    scopes: AztecAddress[],
   ): void {
-    const scopesToSync = scopes.filter(scope => !this.syncedContracts.has(toKey(contractAddress, scope)));
-    if (scopesToSync.length === 0) {
+    // Without a function there is no key to predict from (the request is a direct read).
+    if (!functionToInvokeAfterSync) {
       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);
-      // This cached promise is what every later request for this scope awaits, and both branches write staged data,
-      // so it must run both to completion even when one fails.
-      const promise = allToCompletion([syncNullifiersPromise, runBounded(syncSlot, () => syncScopeFn(scope))])
-        .then(() => {})
-        .catch(err => {
-          this.syncedContracts.delete(key);
-          throw err;
-        });
-      this.syncedContracts.set(key, promise);
+    const speculation = this.#speculationForJob(jobId);
+    const caller: ContractFunction = { address: contractAddress, selector: functionToInvokeAfterSync };
+    for (const callee of this.callGraph.predictDirectCallees(caller)) {
+      // The job's set of already-speculated functions stops the recursion when the predicted graph has a cycle.
+      if (speculation.speculated.has(toCallKey(callee))) {
+        continue;
+      }
+      speculation.speculated.add(toCallKey(callee));
+      const syncPromise = this.#startSyncIfNeeded(
+        callee.address,
+        callee.selector,
+        utilityExecutor,
+        anchorBlockHeader,
+        jobId,
+        scopes,
+      );
+      speculation.syncs.push(syncPromise);
+      // `settle` only escalates these failures at the end of the job: catch here so one does not become an unhandled
+      // rejection before then, and log it.
+      syncPromise.catch(err => {
+        this.log.warn(`Speculative sync of ${callee.address} failed`, { jobId, error: err?.message });
+      });
     }
   }
 
@@ -142,7 +259,7 @@ export class ContractSyncService implements StagedStore {
   async #syncNoteNullifiers(
     contractAddress: AztecAddress,
     anchorBlockHeader: BlockHeader,
-    jobId: string,
+    jobId: JobId,
     scopes: AztecAddress[],
   ): Promise<void> {
     // Protocol contracts don't have private state to sync
@@ -155,6 +272,15 @@ export class ContractSyncService implements StagedStore {
     await noteService.syncNoteNullifiers(contractAddress, scopes);
   }
 
+  #speculationForJob(jobId: JobId): JobSpeculation {
+    let speculation = this.speculationByJob.get(jobId);
+    if (!speculation) {
+      speculation = { speculated: new Set(), syncs: [] };
+      this.speculationByJob.set(jobId, speculation);
+    }
+    return speculation;
+  }
+
   /** 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
@@ -164,7 +290,44 @@ export class ContractSyncService implements StagedStore {
   }
 }
 
-function toKey(contract: AztecAddress, scope: AztecAddress) {
+/** A request to synchronize a contract's private state. */
+type ContractSyncRequest = {
+  /** The contract to sync. */
+  contract: AztecAddress;
+  /**
+   * The function that will be invoked after the sync, or null when nothing will be invoked (e.g. reading
+   * notes/events directly).
+   */
+  functionToInvokeAfterSync: FunctionSelector | null;
+  /** Executes a utility function call under the given scopes. Syncs run each contract's sync_state through it. */
+  utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise<any>;
+  /** The anchor block to sync at. */
+  anchorBlockHeader: BlockHeader;
+  /** The job requesting the sync. */
+  jobId: JobId;
+  /** Access scopes to pass through to the utility executor (affects whose account's private state is discovered). */
+  scopes: AztecAddress[];
+  /**
+   * The function whose execution triggered this sync request, or undefined when the request is a job's top-level use
+   * (an entry call or a direct read) rather than a nested call.
+   */
+  triggeredBy: ContractFunction | undefined;
+};
+
+type JobId = string;
+
+/** A job's speculation state. */
+type JobSpeculation = {
+  /** Functions prediction already ran for, so the recursion stops on cycles in the predicted graph. */
+  speculated: Set<CallKey>;
+  /** Every sync fired by prediction, awaited by {@link settle} before the job commits or discards. */
+  syncs: Promise<void>[];
+};
+
+/** Key of a contract's sync cache entry for a single scope: `contractAddress:scopeAddress`. */
+type SyncKey = `0x${string}:0x${string}`;
+
+function toKey(contract: AztecAddress, scope: AztecAddress): SyncKey {
   return `${contract.toString()}:${scope.toString()}`;
 }
 

package/src/contract_function_simulator/oracle/private_execution_oracle.ts

@@ -661,14 +661,15 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
 
     isStaticCall = isStaticCall || this.callContext.isStaticCall;
 
-    await this.contractSyncService.ensureContractSynced(
-      targetContractAddress,
-      functionSelector,
-      this.utilityExecutor,
-      this.anchorBlockHeader,
-      this.jobId,
-      this.scopes,
-    );
+    await this.contractSyncService.ensureContractSynced({
+      contract: targetContractAddress,
+      functionToInvokeAfterSync: functionSelector,
+      utilityExecutor: this.utilityExecutor,
+      anchorBlockHeader: this.anchorBlockHeader,
+      jobId: this.jobId,
+      scopes: this.scopes,
+      triggeredBy: { address: this.callContext.contractAddress, selector: this.callContext.functionSelector },
+    });
 
     const targetArtifact = await this.anchoredContractData.getFunctionArtifactWithDebugMetadata(
       targetContractAddress,

package/src/contract_function_simulator/oracle/utility_execution_oracle.ts

@@ -1056,14 +1056,15 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
         }
       }
 
-      await this.contractSyncService.ensureContractSynced(
-        targetContractAddress,
-        functionSelector,
-        this.utilityExecutor,
-        this.anchorBlockHeader,
-        this.jobId,
-        this.scopes,
-      );
+      await this.contractSyncService.ensureContractSynced({
+        contract: targetContractAddress,
+        functionToInvokeAfterSync: functionSelector,
+        utilityExecutor: this.utilityExecutor,
+        anchorBlockHeader: this.anchorBlockHeader,
+        jobId: this.jobId,
+        scopes: this.scopes,
+        triggeredBy: { address: this.contractAddress, selector: this.callContext.functionSelector },
+      });
     }
 
     this.logger.debug(

package/src/debug/pxe_debug_utils.ts

@@ -69,15 +69,16 @@ export class PXEDebugUtils {
 
       const contractFunctionSimulator = this.#getSimulatorForTx();
 
-      await this.contractSyncService.ensureContractSynced(
-        filter.contractAddress,
-        null,
-        async (privateSyncCall, execScopes) =>
+      await this.contractSyncService.ensureContractSynced({
+        contract: filter.contractAddress,
+        functionToInvokeAfterSync: null,
+        utilityExecutor: async (privateSyncCall, execScopes) =>
           await this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId),
         anchorBlockHeader,
         jobId,
-        filter.scopes,
-      );
+        scopes: filter.scopes,
+        triggeredBy: undefined,
+      });
 
       return this.noteStore.getNotes(filter, jobId);
     });

package/src/job_coordinator/job_coordinator.ts

@@ -1,5 +1,6 @@
 import { randomBytes } from '@aztec/foundation/crypto/random';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
+import { allToCompletion } from '@aztec/foundation/promise';
 import type { AztecAsyncKVStore } from '@aztec/kv-store';
 
 /**
@@ -24,6 +25,16 @@ export interface StagedStore {
    * @param jobId - The job identifier
    */
   discardStaged(jobId: string): Promise<void>;
+
+  /**
+   * A store may have pending work that must finish before the job's staged writes are committed or discarded, yet
+   * commits run inside a transaction that cannot wait for it. Such stores implement this method: it is called before
+   * every commit and discard, outside the transaction. If settling fails, the commit is cancelled, but a discard
+   * proceeds.
+   *
+   * @param jobId - The job identifier
+   */
+  settle?(jobId: string): Promise<void>;
 }
 
 /**
@@ -108,6 +119,9 @@ export class JobCoordinator {
 
     this.log.debug(`Committing job ${jobId}`);
 
+    // Settling must stay outside the transaction: it can take arbitrarily long.
+    await allToCompletion([...this.#stores.values()].map(store => store.settle?.(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 () => {
@@ -133,6 +147,8 @@ export class JobCoordinator {
 
     this.log.debug(`Aborting job ${jobId}`);
 
+    await this.#settleStoresLoggingFailures(jobId);
+
     for (const store of this.#stores.values()) {
       await store.discardStaged(jobId);
     }
@@ -147,4 +163,18 @@ export class JobCoordinator {
   hasJobInProgress(): boolean {
     return this.#currentJobId !== undefined;
   }
+
+  /**
+   * Settles every store, logging failures instead of propagating them. The abort must run to completion no matter what,
+   * so a store that fails to settle cannot stop the others from discarding or mask the error that aborted the job.
+   */
+  async #settleStoresLoggingFailures(jobId: string): Promise<void> {
+    await allToCompletion(
+      [...this.#stores.values()].map(store =>
+        store.settle?.(jobId).catch(err => {
+          this.log.warn(`Store ${store.storeName} failed to settle while aborting job ${jobId}`, { jobId, err });
+        }),
+      ),
+    );
+  }
 }

package/src/pxe.ts

@@ -327,6 +327,7 @@ export class PXE {
       contractClassService,
       noteStore,
       createLogger('pxe:contract_sync', bindings),
+      config,
     );
     const txResolver = new TxResolverService(readCachedNode);
 
@@ -549,15 +550,16 @@ export class PXE {
     const { origin: contractAddress, functionSelector } = txRequest;
 
     try {
-      await this.contractSyncService.ensureContractSynced(
-        contractAddress,
-        functionSelector,
-        (privateSyncCall, execScopes) =>
+      await this.contractSyncService.ensureContractSynced({
+        contract: contractAddress,
+        functionToInvokeAfterSync: functionSelector,
+        utilityExecutor: (privateSyncCall, execScopes) =>
           this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId),
         anchorBlockHeader,
         jobId,
         scopes,
-      );
+        triggeredBy: undefined,
+      });
 
       const result = await contractFunctionSimulator.run(txRequest, {
         anchorBlockHeader,
@@ -1381,15 +1383,16 @@ export class PXE {
         const contractFunctionSimulator = this.#getSimulatorForTx();
 
         const anchorBlockHeader = await this.anchorBlockStore.getBlockHeader();
-        await this.contractSyncService.ensureContractSynced(
-          call.to,
-          call.selector,
-          (privateSyncCall, execScopes) =>
+        await this.contractSyncService.ensureContractSynced({
+          contract: call.to,
+          functionToInvokeAfterSync: call.selector,
+          utilityExecutor: (privateSyncCall, execScopes) =>
             this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId),
           anchorBlockHeader,
           jobId,
           scopes,
-        );
+          triggeredBy: undefined,
+        });
 
         const { result: executionResult, offchainEffects } = await this.#executeUtility(
           contractFunctionSimulator,
@@ -1459,15 +1462,16 @@ export class PXE {
 
       const contractFunctionSimulator = this.#getSimulatorForTx();
 
-      await this.contractSyncService.ensureContractSynced(
-        filter.contractAddress,
-        null,
-        async (privateSyncCall, execScopes) =>
+      await this.contractSyncService.ensureContractSynced({
+        contract: filter.contractAddress,
+        functionToInvokeAfterSync: null,
+        utilityExecutor: async (privateSyncCall, execScopes) =>
           await this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId),
         anchorBlockHeader,
         jobId,
-        filter.scopes,
-      );
+        scopes: filter.scopes,
+        triggeredBy: undefined,
+      });
     });
 
     // anchorBlockNumber is set during the job and fixed to whatever it is after a block sync