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,211 @@
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+import { ContractCallDependencies, MAX_CONFIDENCE, PREDICTION_THRESHOLD } from './contract_call_dependencies.js';
+
+describe('ContractCallDependencies', () => {
+  let callDependencies: ContractCallDependencies;
+
+  const account = makeAddress(1);
+  const token = makeAddress(2);
+  const fpc = makeAddress(3);
+
+  beforeEach(() => {
+    callDependencies = new ContractCallDependencies(true);
+  });
+
+  it('returns nothing for a contract it has never seen', () => {
+    expect(dependenciesOf(account)).toEqual([]);
+  });
+
+  it('does not predict a callee until enough committed jobs observe the call', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD - 1,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+
+    expect(dependenciesOf(account)).toEqual([]);
+  });
+
+  it('predicts a callee once enough committed jobs observe the call', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([token, fpc]));
+  });
+
+  it('predicts only direct callees, not callees of callees', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: account, callee: fpc },
+        { caller: fpc, callee: token },
+      ],
+    });
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([fpc]));
+    expect(dependenciesOf(fpc)).toEqual(addressStrings([token]));
+  });
+
+  it("predicts a contract's callees even when its own callers rarely call it", () => {
+    runJob({
+      jobId: 'rare',
+      calls: [{ caller: account, callee: token }],
+    });
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: token, callee: fpc }] });
+
+    // The account rarely calls the token, so the account predicts nothing. But the token's own dependencies are known,
+    // ready for the moment a job actually uses it.
+    expect(dependenciesOf(account)).toEqual([]);
+    expect(dependenciesOf(token)).toEqual(addressStrings([fpc]));
+  });
+
+  it('ignores self-calls', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: token, callee: token }] });
+
+    expect(dependenciesOf(token)).toEqual([]);
+  });
+
+  it('does not learn from discarded jobs', () => {
+    runJobs({ count: PREDICTION_THRESHOLD - 1, calls: [{ caller: account, callee: token }] });
+    callDependencies.onContractUsed('discarded', token, account);
+    callDependencies.discardJob('discarded');
+
+    expect(dependenciesOf(account)).toEqual([]);
+  });
+
+  it('leaves confidence untouched by jobs in which the caller makes no calls', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: account, callee: token }] });
+
+    // The account calls no one in these jobs, so the confidence of the callees it did not call is unaffected.
+    for (const jobId of ['read1', 'read2']) {
+      callDependencies.onContractUsed(jobId, account, undefined);
+      callDependencies.commitJob(jobId);
+    }
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([token]));
+  });
+
+  it('keeps predicting a dependency at full confidence through every miss it tolerates', () => {
+    runJobs({
+      count: MAX_CONFIDENCE,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+    runJobs({ count: MAX_CONFIDENCE - PREDICTION_THRESHOLD, calls: [{ caller: account, callee: token }] });
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([token, fpc]));
+  });
+
+  it('caps confidence, so a heavily called dependency stops being predicted one miss past that tolerance', () => {
+    runJobs({
+      count: MAX_CONFIDENCE * 2,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+    runJobs({ count: MAX_CONFIDENCE - PREDICTION_THRESHOLD + 1, calls: [{ caller: account, callee: token }] });
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([token]));
+  });
+
+  it('drops a dependency below the threshold on a miss and predicts it again after one hit', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+    expect(dependenciesOf(account)).toEqual(addressStrings([token, fpc]));
+
+    runJob({ jobId: 'miss', calls: [{ caller: account, callee: token }] });
+    expect(dependenciesOf(account)).toEqual(addressStrings([token]));
+
+    runJob({
+      jobId: 'refresh',
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([token, fpc]));
+  });
+
+  it('stops predicting a forgotten contract from every caller that called it', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+        { caller: token, callee: fpc },
+      ],
+    });
+    expect(dependenciesOf(account)).toEqual(addressStrings([token, fpc]));
+    expect(dependenciesOf(token)).toEqual(addressStrings([fpc]));
+
+    callDependencies.forget(fpc);
+
+    expect(dependenciesOf(account)).toEqual(addressStrings([token]));
+    expect(dependenciesOf(token)).toEqual([]);
+  });
+
+  it('never records dependencies when disabled', () => {
+    callDependencies = new ContractCallDependencies(false);
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: account, callee: token },
+        { caller: account, callee: fpc },
+      ],
+    });
+
+    expect(dependenciesOf(account)).toEqual([]);
+  });
+
+  /** Runs `count` whole jobs, each observing the given direct calls. */
+  function runJobs({ count, calls }: { count: number; calls: Call[] }) {
+    for (let i = 0; i < count; i++) {
+      runJob({ jobId: `job${i}`, calls });
+    }
+  }
+
+  /** Runs a whole job: uses the first caller as the entry, then each callee with its caller, and commits. */
+  function runJob({ jobId, calls }: { jobId: string; calls: Call[] }) {
+    callDependencies.onContractUsed(jobId, calls[0].caller, undefined);
+    for (const { caller, callee } of calls) {
+      callDependencies.onContractUsed(jobId, callee, caller);
+    }
+    callDependencies.commitJob(jobId);
+  }
+
+  /** Returns the direct dependencies predicted for the given contract, as sorted address strings. */
+  function dependenciesOf(contract: AztecAddress): string[] {
+    return callDependencies
+      .predictDirectDependencies(contract)
+      .map(address => address.toString())
+      .sort();
+  }
+});
+
+/** A direct call observed by a job. */
+type Call = { caller: AztecAddress; callee: AztecAddress };
+
+function makeAddress(index: number): AztecAddress {
+  return AztecAddress.fromNumberUnsafe(0x1000 + index);
+}
+
+function addressStrings(addresses: AztecAddress[]): string[] {
+  return addresses.map(address => address.toString()).sort();
+}

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

@@ -0,0 +1,107 @@
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+/** Confidence a dependency must reach to be predicted. */
+export const PREDICTION_THRESHOLD = 2;
+
+/** Cap on a dependency's confidence, so a contract called by many jobs is still forgotten within a few missed ones. */
+export const MAX_CONFIDENCE = 5;
+
+/**
+ * Predicts the contracts a job is about to need from the direct calls observed in past jobs, so the caller can sync
+ * them in parallel.
+ *
+ * Wallet workloads are repetitive: a contract usually calls the same contracts every time it executes. Each committed
+ * job in which a caller calls a callee adds a point of confidence to that dependency (capped at
+ * {@link MAX_CONFIDENCE}); each committed job in which the caller calls other contracts but not that one subtracts a
+ * point, forgetting the dependency at zero. A job in which a contract calls nothing leaves its dependencies untouched,
+ * so read-only uses (e.g. reading notes or events) erode nothing.
+ *
+ * Purely in-memory bookkeeping: the dependencies are lost when PXE is rebuilt (e.g. on restart).
+ */
+export class ContractCallDependencies {
+  // job → caller contract → contracts it called directly
+  private readonly activeJobs: Map<JobId, Map<ContractAddress, Set<ContractAddress>>> = new Map();
+
+  // caller contract → contract it calls directly → confidence score
+  private readonly dependencyConfidence: Map<ContractAddress, Map<ContractAddress, number>> = new Map();
+
+  constructor(
+    private readonly enabled: boolean,
+    private readonly log: Logger = createLogger('pxe:contract_call_dependencies'),
+  ) {}
+
+  /**
+   * Records that `caller` directly called `callee` in the given job. Top-level uses (no caller) only ever appear on the
+   * caller side of a dependency, so they record nothing themselves.
+   */
+  onContractUsed(jobId: JobId, callee: AztecAddress, caller: AztecAddress | undefined): void {
+    if (!this.enabled || !caller || caller.equals(callee)) {
+      return;
+    }
+    let callsInJob = this.activeJobs.get(jobId);
+    if (!callsInJob) {
+      callsInJob = new Map();
+      this.activeJobs.set(jobId, callsInJob);
+    }
+    let callees = callsInJob.get(caller.toString());
+    if (!callees) {
+      callees = new Set();
+      callsInJob.set(caller.toString(), callees);
+    }
+    callees.add(callee.toString());
+  }
+
+  /** Predicts the contracts the given one will call directly: callees at {@link PREDICTION_THRESHOLD}+. */
+  predictDirectDependencies(contractAddress: AztecAddress): AztecAddress[] {
+    const dependencies =
+      this.dependencyConfidence.get(contractAddress.toString()) ?? new Map<ContractAddress, number>();
+    return [...dependencies.entries()]
+      .filter(([, confidence]) => confidence >= PREDICTION_THRESHOLD)
+      .map(([contract]) => AztecAddress.fromStringUnsafe(contract));
+  }
+
+  /** Learns dependencies from the calls the committed job observed. */
+  commitJob(jobId: JobId): void {
+    const callsInJob = this.activeJobs.get(jobId);
+    this.activeJobs.delete(jobId);
+    if (!callsInJob) {
+      return;
+    }
+
+    for (const [caller, callees] of callsInJob) {
+      const dependencies = this.dependencyConfidence.get(caller) ?? new Map<ContractAddress, number>();
+      for (const [contract, confidence] of dependencies) {
+        const delta = callees.has(contract) ? 1 : -1;
+        const updated = Math.min(confidence + delta, MAX_CONFIDENCE);
+        if (updated === 0) {
+          dependencies.delete(contract);
+        } else {
+          dependencies.set(contract, updated);
+        }
+      }
+      // Contracts called for the first time enter at confidence 1, below the prediction threshold.
+      [...callees].filter(contract => !dependencies.has(contract)).forEach(contract => dependencies.set(contract, 1));
+      this.dependencyConfidence.set(caller, dependencies);
+    }
+  }
+
+  /** Drops a discarded job without learning. */
+  discardJob(jobId: JobId): void {
+    this.activeJobs.delete(jobId);
+  }
+
+  /** Stops predicting a contract, dropping it from every caller known to call it. */
+  forget(contractAddress: AztecAddress): void {
+    const contract = contractAddress.toString();
+    for (const [caller, dependencies] of this.dependencyConfidence) {
+      if (dependencies.delete(contract) && dependencies.size === 0) {
+        this.dependencyConfidence.delete(caller);
+      }
+    }
+  }
+}
+
+type JobId = string;
+
+type ContractAddress = string;