feat(pxe): speculatively sync predicted contract calls

yarn-project/foundation/src/config/env_var.ts

@@ -213,6 +213,7 @@ export type EnvVar =
   | 'PROVER_TEST_DELAY_TYPE'
   | 'PROVER_TEST_VERIFICATION_DELAY_MS'
   | 'PXE_AUTO_SYNC'
+  | 'PXE_CONCURRENT_CONTRACT_SYNC_ENABLED'
   | 'PXE_L2_BLOCK_BATCH_SIZE'
   | 'PXE_PROVER_ENABLED'
   | 'PXE_SYNC_CHAIN_TIP'

yarn-project/pxe/src/config/index.ts

@@ -36,7 +36,23 @@ export interface BlockSynchronizerConfig {
   autoSync: boolean;
 }
 
-export type PXEConfig = KernelProverConfig & DataStoreConfig & ChainConfig & BlockSynchronizerConfig;
+/**
+ * Configuration settings for the contract sync service.
+ */
+export interface ContractSyncConfig {
+  /**
+   * Whether PXE speculatively syncs contracts it predicts will follow the one requested, running them concurrently
+   * with it instead of waiting for execution to reach them. When enabled, repeated flows sync faster, but a wrong
+   * prediction spends unnecessary node requests syncing contracts the job never uses. Experimental; off by default.
+   */
+  concurrentContractSyncEnabled: boolean;
+}
+
+export type PXEConfig = KernelProverConfig &
+  DataStoreConfig &
+  ChainConfig &
+  BlockSynchronizerConfig &
+  ContractSyncConfig;
 
 export type CliPXEOptions = {
   /** Custom Aztec Node URL to connect to  */
@@ -74,6 +90,12 @@ export const pxeConfigMappings: ConfigMappingsType<PXEConfig> = {
       'Whether PXE syncs with the node automatically before each operation. Disable to let the caller (e.g. a wallet) drive syncs explicitly via pxe.sync().',
     ...booleanConfigHelper(true),
   },
+  concurrentContractSyncEnabled: {
+    env: 'PXE_CONCURRENT_CONTRACT_SYNC_ENABLED',
+    description:
+      'Whether PXE speculatively syncs contracts it predicts will follow the one requested, running them concurrently with it. Repeated flows sync faster, but a wrong prediction spends unnecessary node requests. Experimental; off by default.',
+    ...booleanConfigHelper(false),
+  },
 };
 
 /**

yarn-project/pxe/src/contract/contract_call_dependencies.test.ts

@@ -0,0 +1,190 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { FunctionSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+import { ContractCallDependencies, MAX_CONFIDENCE, PREDICTION_THRESHOLD } from './contract_call_dependencies.js';
+
+/** Missed jobs a dependency at full confidence survives while still being predicted. */
+const TOLERATED_MISSES = MAX_CONFIDENCE - PREDICTION_THRESHOLD;
+
+describe('ContractCallDependencies', () => {
+  let callDependencies: ContractCallDependencies;
+
+  const account = makeAddress(1);
+  const token = makeAddress(2);
+  const fpc = makeAddress(3);
+  const otherAccount = makeAddress(4);
+  const entrypoint = makeSelector(0x11223344);
+  const otherFunction = makeSelector(0x55667788);
+  const scopes = [account];
+  const accountEntryCall = { contract: account, functionToInvoke: entrypoint };
+
+  beforeEach(() => {
+    callDependencies = new ContractCallDependencies(true);
+  });
+
+  it('returns nothing for an entry call it has never seen', () => {
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual([]);
+  });
+
+  it('does not return a contract until enough committed jobs have used it', () => {
+    runJobs({ count: PREDICTION_THRESHOLD - 1, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual([]);
+  });
+
+  it('returns a contract once enough committed jobs of the same entry call use it', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token, fpc]));
+  });
+
+  it('does not share dependencies across different entry calls', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token], scopes });
+
+    const otherEntryContract = { contract: token, functionToInvoke: entrypoint };
+    const otherEntryFunction = { contract: account, functionToInvoke: otherFunction };
+    expect(knownDependencies({ entryCall: otherEntryContract, scopes })).toEqual([]);
+    expect(knownDependencies({ entryCall: otherEntryFunction, scopes })).toEqual([]);
+  });
+
+  it('does not share dependencies across different scope sets', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes: [otherAccount] })).toEqual([]);
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes: [account, otherAccount] })).toEqual([]);
+  });
+
+  it('ignores scope order when identifying an entry call', () => {
+    const bothAccounts = [account, otherAccount];
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token], scopes: bothAccounts });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes: [otherAccount, account] })).toEqual(
+      addressStrings([token]),
+    );
+  });
+
+  it('distinguishes an entry call with no function from one with a function', () => {
+    const noFunctionEntryCall = { contract: account, functionToInvoke: null };
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: noFunctionEntryCall, uses: [token], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual([]);
+    expect(knownDependencies({ entryCall: noFunctionEntryCall, scopes })).toEqual(addressStrings([token]));
+  });
+
+  it('does not learn from discarded jobs', () => {
+    runJobs({ count: PREDICTION_THRESHOLD - 1, entryCall: accountEntryCall, uses: [token], scopes });
+    callDependencies.onContractUsed('discarded', account, entrypoint, scopes);
+    callDependencies.onContractUsed('discarded', token, otherFunction, scopes);
+    callDependencies.discardJob('discarded');
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual([]);
+  });
+
+  it('keeps returning a dependency at full confidence through every miss it tolerates', () => {
+    runJobs({ count: MAX_CONFIDENCE, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+    runJobs({ count: TOLERATED_MISSES, entryCall: accountEntryCall, uses: [token], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token, fpc]));
+  });
+
+  it('caps confidence, so a heavily used dependency stops being returned one miss past that tolerance', () => {
+    runJobs({ count: MAX_CONFIDENCE * 2, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+    runJobs({ count: TOLERATED_MISSES + 1, entryCall: accountEntryCall, uses: [token], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token]));
+  });
+
+  it('stops returning a dependency whose confidence falls below the threshold', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+    runJob({ jobId: 'miss', entryCall: accountEntryCall, uses: [token], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token]));
+  });
+
+  it('returns a dependency that fell below the threshold as soon as one job uses it again', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+    runJob({ jobId: 'miss', entryCall: accountEntryCall, uses: [token], scopes });
+    runJob({ jobId: 'refresh', entryCall: accountEntryCall, uses: [token, fpc], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token, fpc]));
+  });
+
+  it('stops returning a contract as soon as it is forgotten', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token, fpc]));
+
+    // Forgetting needs an active job, since that is what identifies the entry call.
+    callDependencies.onContractUsed('forgetting', account, entrypoint, scopes);
+    callDependencies.forget('forgetting', fpc);
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token]));
+  });
+
+  it('ignores forget calls for unknown jobs', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token], scopes });
+    callDependencies.forget('unknownJob', token);
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual(addressStrings([token]));
+  });
+
+  it('never returns or records dependencies when disabled', () => {
+    callDependencies = new ContractCallDependencies(false);
+    runJobs({ count: PREDICTION_THRESHOLD, entryCall: accountEntryCall, uses: [token, fpc], scopes });
+
+    expect(knownDependencies({ entryCall: accountEntryCall, scopes })).toEqual([]);
+  });
+
+  /** Runs `count` whole jobs of the given entry call, each using the given contracts. */
+  function runJobs({ count, ...job }: { count: number } & Omit<JobRun, 'jobId'>) {
+    for (let i = 0; i < count; i++) {
+      runJob({ jobId: `job${i}`, ...job });
+    }
+  }
+
+  /** Runs a whole job: starts it with the given entry call, uses the given contracts, and commits. */
+  function runJob({ jobId, entryCall, uses, scopes: jobScopes }: JobRun) {
+    callDependencies.onContractUsed(jobId, entryCall.contract, entryCall.functionToInvoke, jobScopes);
+    for (const contract of uses) {
+      callDependencies.onContractUsed(jobId, contract, otherFunction, jobScopes);
+    }
+    callDependencies.commitJob(jobId);
+  }
+
+  /**
+   * Returns the dependencies known for the given entry call, through a job that starts with it and uses nothing
+   * else. Committing such a job records no dependencies, so probing does not change what is known.
+   */
+  function knownDependencies({ entryCall, scopes: jobScopes }: JobStart): string[] {
+    const known = callDependencies.onContractUsed('probe', entryCall.contract, entryCall.functionToInvoke, jobScopes);
+    callDependencies.commitJob('probe');
+    return known.map(address => address.toString()).sort();
+  }
+});
+
+/** How a job starts: the entry call it makes first, and the scopes every one of its uses runs under. */
+type JobStart = {
+  entryCall: EntryCall;
+  scopes: AztecAddress[];
+};
+
+/** A job to run: how it starts, the id it runs under, and the contracts it uses after the start. */
+type JobRun = JobStart & {
+  jobId: string;
+  uses: AztecAddress[];
+};
+
+/** The (contract, function) a job starts with. */
+type EntryCall = { contract: AztecAddress; functionToInvoke: FunctionSelector | null };
+
+function makeAddress(index: number): AztecAddress {
+  return AztecAddress.fromNumberUnsafe(0x1000 + index);
+}
+
+function makeSelector(value: number): FunctionSelector {
+  return FunctionSelector.fromField(new Fr(value));
+}
+
+function addressStrings(addresses: AztecAddress[]): string[] {
+  return addresses.map(address => address.toString()).sort();
+}

yarn-project/pxe/src/contract/contract_call_dependencies.ts

@@ -0,0 +1,152 @@
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import type { FunctionSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+/** Confidence a dependency must reach to be returned, requiring a second committed job to use its contract. */
+export const PREDICTION_THRESHOLD = 2;
+
+/** Cap on a dependency's confidence, so a contract used by many jobs is still forgotten within a few missed ones. */
+export const MAX_CONFIDENCE = 5;
+
+/**
+ * Predicts a job's contract dependencies from past jobs of its entry call, so the caller can sync them in parallel.
+ *
+ * Wallet workloads are repetitive: jobs starting with the same entry call (see {@link EntryCallId}) usually depend on
+ * the same contracts. Each dependency carries a confidence score: every committed job of the entry call that uses it
+ * adds a point (capped at {@link MAX_CONFIDENCE}) and every one that does not subtracts one, forgetting it at zero.
+ * Only dependencies at {@link PREDICTION_THRESHOLD} or above are returned, so one-off dependencies picked by a call's
+ * parameters (e.g. the tokens of a single swap) are never predicted.
+ *
+ * Purely in-memory bookkeeping: the dependencies are lost when PXE is rebuilt (e.g. on restart).
+ */
+export class ContractCallDependencies {
+  private readonly activeJobs: Map<JobId, ActiveJob> = new Map();
+
+  // entry call → dependency → confidence score
+  private readonly dependencyConfidence: Map<EntryCallId, Map<ContractAddress, number>> = new Map();
+
+  constructor(
+    private readonly enabled: boolean,
+    private readonly log: Logger = createLogger('pxe:contract_call_dependencies'),
+  ) {}
+
+  /**
+   * Records that a job used a contract.
+   * @param functionToInvoke - The function that will be invoked on `contractAddress`, or null when nothing will be
+   * invoked (e.g. reading notes/events directly).
+   * @returns The dependencies of the job's entry call whose confidence reached {@link PREDICTION_THRESHOLD}.
+   */
+  onContractUsed(
+    jobId: JobId,
+    contractAddress: AztecAddress,
+    functionToInvoke: FunctionSelector | null,
+    scopes: AztecAddress[],
+  ): AztecAddress[] {
+    if (!this.enabled) {
+      return [];
+    }
+    let job = this.activeJobs.get(jobId);
+    if (!job) {
+      const entryCallId = toEntryCallId(contractAddress, functionToInvoke, scopes);
+      this.log.debug(`Job started with entry call ${entryCallId}`, { jobId, entryCallId });
+      job = { entryCallId, used: new Set() };
+      this.activeJobs.set(jobId, job);
+    } else {
+      // Only contracts used after the start are considered dependencies.
+      job.used.add(contractAddress.toString());
+    }
+
+    const dependencies = this.dependencyConfidence.get(job.entryCallId);
+    if (!dependencies) {
+      return [];
+    }
+    return [...dependencies.entries()]
+      .filter(([, confidence]) => confidence >= PREDICTION_THRESHOLD)
+      .map(([contract]) => AztecAddress.fromStringUnsafe(contract));
+  }
+
+  /** Raises the confidence of the dependencies the committed job used and lowers the rest, forgetting any at zero. */
+  commitJob(jobId: JobId): void {
+    const job = this.activeJobs.get(jobId);
+    this.activeJobs.delete(jobId);
+    if (!job || job.used.size === 0) {
+      return;
+    }
+
+    const dependencies = this.dependencyConfidence.get(job.entryCallId) ?? new Map<ContractAddress, number>();
+    for (const [contract, confidence] of dependencies) {
+      const delta = job.used.has(contract) ? 1 : -1;
+      const updated = Math.min(confidence + delta, MAX_CONFIDENCE);
+      if (updated === 0) {
+        dependencies.delete(contract);
+      } else {
+        dependencies.set(contract, updated);
+      }
+    }
+    // Contracts used for the first time enter at confidence 1, below the prediction threshold.
+    [...job.used].filter(contract => !dependencies.has(contract)).forEach(contract => dependencies.set(contract, 1));
+    this.dependencyConfidence.set(job.entryCallId, dependencies);
+
+    this.log.debug(`Remembering ${dependencies.size} contract(s) for entry call ${job.entryCallId}`, {
+      jobId,
+      entryCallId: job.entryCallId,
+      confidence: Object.fromEntries(dependencies),
+    });
+  }
+
+  /** Drops a discarded job without learning. */
+  discardJob(jobId: JobId): void {
+    this.activeJobs.delete(jobId);
+  }
+
+  /** Stops returning a contract for the given job's entry call. */
+  forget(jobId: JobId, contractAddress: AztecAddress): void {
+    const entryCallId = this.activeJobs.get(jobId)?.entryCallId;
+    if (!entryCallId) {
+      return;
+    }
+    const dependencies = this.dependencyConfidence.get(entryCallId);
+    if (!dependencies?.delete(contractAddress.toString())) {
+      return;
+    }
+    if (dependencies.size === 0) {
+      this.dependencyConfidence.delete(entryCallId);
+    }
+    this.log.debug(`Dropped ${contractAddress} from the remembered dependencies of entry call ${entryCallId}`, {
+      jobId,
+      entryCallId,
+    });
+  }
+}
+
+/**
+ * Builds an entry call id from a call's contract, function and scopes.
+ * The scopes are sorted first, so their order does not matter.
+ */
+function toEntryCallId(
+  contract: AztecAddress,
+  functionToInvoke: FunctionSelector | null,
+  scopes: AztecAddress[],
+): EntryCallId {
+  const scopeSet = scopes
+    .map(scope => scope.toString())
+    .sort()
+    .join(',');
+  return `${contract.toString()}:${functionToInvoke?.toString() ?? ''}:${scopeSet}`;
+}
+
+type JobId = string;
+
+/**
+ * Identifies how a job started: the first (contract, function) it used, plus that use's scope set. The scope set
+ * separates entry calls that share an entry contract across accounts (e.g. a shared multi-call entrypoint).
+ */
+type EntryCallId = string;
+
+type ContractAddress = string;
+
+/** An active job's entry call, plus the addresses of the contracts it has used since it started. */
+type ActiveJob = {
+  readonly entryCallId: EntryCallId;
+  readonly used: Set<ContractAddress>;
+};