feat: merge-train/fairies

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

@@ -214,6 +214,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,25 @@ 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 +92,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_graph.test.ts

@@ -0,0 +1,212 @@
+import { FunctionSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+import {
+  ContractCallGraph,
+  type ContractFunction,
+  MAX_CONFIDENCE,
+  PREDICTION_THRESHOLD,
+} from './contract_call_graph.js';
+
+describe('ContractCallGraph', () => {
+  let callGraph: ContractCallGraph;
+
+  const accountEntrypoint = fn(1, 1);
+  const accountClaim = fn(1, 2);
+  const tokenTransfer = fn(2, 1);
+  const tokenBalance = fn(2, 2);
+  const fpcFee = fn(3, 1);
+
+  beforeEach(() => {
+    callGraph = new ContractCallGraph(true);
+  });
+
+  it('returns nothing for a function it has never seen', () => {
+    expect(calleesOf(accountEntrypoint)).toEqual([]);
+  });
+
+  it('does not predict a callee until enough committed jobs observe the call', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD - 1,
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+
+    expect(calleesOf(accountEntrypoint)).toEqual([]);
+  });
+
+  it('predicts a callee once enough committed jobs observe the call', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee]));
+  });
+
+  it('predicts only direct callees, not callees of callees', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: accountEntrypoint, callee: fpcFee },
+        { caller: fpcFee, callee: tokenTransfer },
+      ],
+    });
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([fpcFee]));
+    expect(calleesOf(fpcFee)).toEqual(callKeys([tokenTransfer]));
+  });
+
+  it('keys calls per function, so a sibling function of the same contract predicts nothing', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
+    expect(calleesOf(accountClaim)).toEqual([]);
+  });
+
+  it("predicts a function's callees even when its own callers rarely call it", () => {
+    runJob({
+      jobId: 'rare',
+      calls: [{ caller: accountEntrypoint, callee: tokenTransfer }],
+    });
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: tokenTransfer, callee: fpcFee }] });
+
+    expect(calleesOf(accountEntrypoint)).toEqual([]);
+    expect(calleesOf(tokenTransfer)).toEqual(callKeys([fpcFee]));
+  });
+
+  it('ignores same-contract calls', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: tokenTransfer, callee: tokenBalance }] });
+
+    expect(calleesOf(tokenTransfer)).toEqual([]);
+  });
+
+  it('does not learn from discarded jobs', () => {
+    runJobs({ count: PREDICTION_THRESHOLD - 1, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+    callGraph.recordCall({ jobId: 'discarded', caller: accountEntrypoint, callee: tokenTransfer });
+    callGraph.discardJob('discarded');
+
+    expect(calleesOf(accountEntrypoint)).toEqual([]);
+  });
+
+  it('leaves confidence untouched by jobs in which the caller makes no calls', () => {
+    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+
+    // 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']) {
+      callGraph.commitJob(jobId);
+    }
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
+  });
+
+  it('keeps predicting a callee at full confidence through every miss it tolerates', () => {
+    runJobs({
+      count: MAX_CONFIDENCE,
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+    runJobs({
+      count: MAX_CONFIDENCE - PREDICTION_THRESHOLD,
+      calls: [{ caller: accountEntrypoint, callee: tokenTransfer }],
+    });
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee]));
+  });
+
+  it('caps confidence, so a heavily called callee stops being predicted one miss past that tolerance', () => {
+    runJobs({
+      count: MAX_CONFIDENCE * 2,
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+    runJobs({
+      count: MAX_CONFIDENCE - PREDICTION_THRESHOLD + 1,
+      calls: [{ caller: accountEntrypoint, callee: tokenTransfer }],
+    });
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
+  });
+
+  it('drops a callee below the threshold on a miss and predicts it again after one hit', () => {
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee]));
+
+    runJob({ jobId: 'miss', calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
+
+    runJob({
+      jobId: 'refresh',
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+
+    expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee]));
+  });
+
+  it('never records calls when disabled', () => {
+    callGraph = new ContractCallGraph(false);
+    runJobs({
+      count: PREDICTION_THRESHOLD,
+      calls: [
+        { caller: accountEntrypoint, callee: tokenTransfer },
+        { caller: accountEntrypoint, callee: fpcFee },
+      ],
+    });
+
+    expect(calleesOf(accountEntrypoint)).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: records each direct call and commits. */
+  function runJob({ jobId, calls }: { jobId: string; calls: Call[] }) {
+    for (const { caller, callee } of calls) {
+      callGraph.recordCall({ jobId, caller, callee });
+    }
+    callGraph.commitJob(jobId);
+  }
+
+  /** Returns the direct callees predicted for the given function, as sorted `address:selector` strings. */
+  function calleesOf(caller: ContractFunction): string[] {
+    return callKeys(callGraph.predictDirectCallees(caller));
+  }
+});
+
+/** A direct call observed by a job. */
+type Call = { caller: ContractFunction; callee: ContractFunction };
+
+function fn(contractIndex: number, functionIndex: number): ContractFunction {
+  return { address: makeAddress(contractIndex), selector: new FunctionSelector(0x1000 + functionIndex) };
+}
+
+function makeAddress(index: number): AztecAddress {
+  return AztecAddress.fromNumberUnsafe(0x1000 + index);
+}
+
+/** Flattens functions to sorted `address:selector` strings, so sets of predictions can be compared. */
+function callKeys(functions: ContractFunction[]): string[] {
+  return functions.map(({ address, selector }) => `${address.toString()}:${selector.toString()}`).sort();
+}

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

@@ -0,0 +1,113 @@
+import { FunctionSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+/** Confidence a call must reach to be predicted. */
+export const PREDICTION_THRESHOLD = 2;
+
+/** Cap on a call's confidence, so a function called by many jobs is still dropped within a few missed ones. */
+export const MAX_CONFIDENCE = 5;
+
+/**
+ * A call graph over contract functions - who calls whom - learned from the direct calls observed in past jobs, so
+ * a function's predicted callees can sync their contracts before execution reaches them.
+ *
+ * A function's direct calls tend to repeat across jobs: constrained delivery calls the handshake registry, a transfer
+ * may call an authwit, an AMM calls its tokens. The same function does not always make the same calls, though: they
+ * can depend on context or storage state. A call must therefore repeat often enough to earn confidence before it is
+ * predicted. Calls are keyed per function, not per contract, since different functions of a contract call different
+ * contracts. See {@link commitJob} for how each call's confidence is learned from committed jobs.
+ *
+ * Purely in-memory bookkeeping: the graph is lost when PXE is rebuilt (e.g. on restart).
+ */
+export class ContractCallGraph {
+  // job -> caller function -> functions it called directly
+  private readonly activeJobs: Map<JobId, Map<CallKey, Set<CallKey>>> = new Map();
+
+  // caller function -> function it calls directly -> confidence score
+  private readonly callConfidence: Map<CallKey, Map<CallKey, number>> = new Map();
+
+  constructor(private readonly enabled: boolean) {}
+
+  /** Records that `caller` directly called `callee` in the given job. */
+  recordCall({ jobId, caller, callee }: { jobId: JobId; caller: ContractFunction; callee: ContractFunction }): void {
+    // Same-contract calls are ignored: our goal is to warm a callee's contract ahead of use, and the target of such
+    // a call is already warm.
+    if (!this.enabled || caller.address.equals(callee.address)) {
+      return;
+    }
+    let callsInJob = this.activeJobs.get(jobId);
+    if (!callsInJob) {
+      callsInJob = new Map();
+      this.activeJobs.set(jobId, callsInJob);
+    }
+    let callees = callsInJob.get(toCallKey(caller));
+    if (!callees) {
+      callees = new Set();
+      callsInJob.set(toCallKey(caller), callees);
+    }
+    callees.add(toCallKey(callee));
+  }
+
+  /** Predicts the functions `caller` will call directly. */
+  predictDirectCallees(caller: ContractFunction): ContractFunction[] {
+    const callees = this.callConfidence.get(toCallKey(caller)) ?? new Map<CallKey, number>();
+    return [...callees.entries()]
+      .filter(([, confidence]) => confidence >= PREDICTION_THRESHOLD)
+      .map(([callee]) => fromCallKey(callee));
+  }
+
+  /**
+   * Commits the job so the calls it observed are recorded and learned from. A function that called nothing keeps its
+   * callees untouched, so read-only uses (e.g. reading notes or events) erode nothing.
+   */
+  commitJob(jobId: JobId): void {
+    const callsInJob = this.activeJobs.get(jobId);
+    this.activeJobs.delete(jobId);
+    if (!callsInJob) {
+      return;
+    }
+
+    for (const [caller, observed] of callsInJob) {
+      const callees = this.callConfidence.get(caller) ?? new Map<CallKey, number>();
+      for (const [callee, confidence] of callees) {
+        const delta = observed.has(callee) ? 1 : -1;
+        const updated = Math.min(confidence + delta, MAX_CONFIDENCE);
+        if (updated === 0) {
+          callees.delete(callee);
+        } else {
+          callees.set(callee, updated);
+        }
+      }
+      // First-time callees enter at 1, below PREDICTION_THRESHOLD: a call must repeat before it is predicted.
+      [...observed].filter(callee => !callees.has(callee)).forEach(callee => callees.set(callee, 1));
+      this.callConfidence.set(caller, callees);
+    }
+  }
+
+  /** Drops a discarded job without learning. */
+  discardJob(jobId: JobId): void {
+    this.activeJobs.delete(jobId);
+  }
+}
+
+/** A specific function of a contract, as observed in a call. */
+export type ContractFunction = {
+  /** The address of the contract the function belongs to. */
+  address: AztecAddress;
+  /** The selector of the function. */
+  selector: FunctionSelector;
+};
+
+type JobId = string;
+
+/** A {@link ContractFunction} flattened to a `contractAddress:selector` string, so maps can key on it. */
+export type CallKey = `0x${string}:${string}`;
+
+export function toCallKey({ address, selector }: ContractFunction): CallKey {
+  return `${address.toString()}:${selector.toString()}`;
+}
+
+function fromCallKey(key: CallKey): ContractFunction {
+  const [address, selector] = key.split(':');
+  return { address: AztecAddress.fromStringUnsafe(address), selector: FunctionSelector.fromString(selector) };
+}