@aztec/pxe 0.0.1-commit.03f7ef2 → 0.0.1-commit.04d373f

package/src/private_kernel/private_kernel_execution_prover.ts

@@ -1,6 +1,8 @@
+import { MAX_APPS_PER_KERNEL } from '@aztec/constants';
+import { uniqueBy } from '@aztec/foundation/collection';
 import { vkAsFieldsMegaHonk } from '@aztec/foundation/crypto/keys';
 import { Fr } from '@aztec/foundation/curves/bn254';
-import { createLogger } from '@aztec/foundation/log';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { pushTestData } from '@aztec/foundation/testing';
 import { Timer } from '@aztec/foundation/timer';
 import { getVKTreeRoot } from '@aztec/noir-protocol-circuits-types/vk-tree';
@@ -16,25 +18,32 @@ import {
   PrivateKernelCircuitPublicInputs,
   PrivateKernelData,
   type PrivateKernelExecutionProofOutput,
+  PrivateKernelInit2CircuitPrivateInputs,
+  PrivateKernelInit3CircuitPrivateInputs,
   PrivateKernelInitCircuitPrivateInputs,
+  PrivateKernelInner2CircuitPrivateInputs,
+  PrivateKernelInner3CircuitPrivateInputs,
   PrivateKernelInnerCircuitPrivateInputs,
   type PrivateKernelSimulateOutput,
   PrivateKernelTailCircuitPrivateInputs,
   type PrivateKernelTailCircuitPublicInputs,
   PrivateVerificationKeyHints,
+  type UpdatedClassIdHints,
 } from '@aztec/stdlib/kernel';
 import { ChonkProof, ChonkProofWithPublicInputs } from '@aztec/stdlib/proofs';
 import {
   type PrivateCallExecutionResult,
   type PrivateExecutionResult,
   TxRequest,
-  collectNoteHashLeafIndexMap,
+  collectNested,
   collectNoteHashNullifierCounterMap,
   getFinalMinRevertibleSideEffectCounter,
 } from '@aztec/stdlib/tx';
 import { VerificationKeyAsFields, VerificationKeyData, VkData } from '@aztec/stdlib/vks';
 
-import { PrivateKernelResetPrivateInputsBuilder } from './hints/build_private_kernel_reset_private_inputs.js';
+import { BatchPlanner } from './batch_planner.js';
+import { computeTxExpirationTimestamp } from './hints/compute_tx_expiration_timestamp.js';
+import { PrivateKernelResetPrivateInputsBuilder } from './hints/private_kernel_reset_private_inputs_builder.js';
 import type { PrivateKernelOracle } from './private_kernel_oracle.js';
 
 const NULL_SIMULATE_OUTPUT: PrivateKernelSimulateOutput<PrivateKernelCircuitPublicInputs> = {
@@ -57,13 +66,17 @@ export interface PrivateKernelExecutionProverConfig {
  * inform state tree updates.
  */
 export class PrivateKernelExecutionProver {
-  private log = createLogger('pxe:private-kernel-execution-prover');
+  private log: Logger;
 
   constructor(
     private oracle: PrivateKernelOracle,
     private proofCreator: PrivateKernelProver,
     private fakeProofs = false,
-  ) {}
+    bindings?: LoggerBindings,
+    private maxBatchSize: number = MAX_APPS_PER_KERNEL,
+  ) {
+    this.log = createLogger('pxe:private-kernel-execution-prover', bindings);
+  }
 
   /**
    * Generate a proof for a given transaction request and execution result.
@@ -101,11 +114,17 @@ export class PrivateKernelExecutionProver {
 
     const executionSteps: PrivateExecutionStep[] = [];
 
-    const noteHashLeafIndexMap = collectNoteHashLeafIndexMap(executionResult);
     const noteHashNullifierCounterMap = collectNoteHashNullifierCounterMap(executionResult);
     const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(executionResult);
     const splitCounter = isPrivateOnlyTx ? 0 : minRevertibleSideEffectCounter;
 
+    // Each kernel iteration absorbs up to `maxBatchSize` apps. The planner walks the upcoming
+    // apps and decides where the next reset must fall by using an accumulator and
+    // reusing the existing single-app `needsReset()` check.
+    const planner = new BatchPlanner(noteHashNullifierCounterMap, splitCounter, this.maxBatchSize);
+
+    const updatedClassIdHintsMap = await this.prefetchUpdatedClassIdHints(executionResult);
+
     while (executionStack.length) {
       if (!firstIteration) {
         let resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
@@ -115,8 +134,9 @@ export class PrivateKernelExecutionProver {
           splitCounter,
         );
         while (resetBuilder.needsReset()) {
+          // Inner reset: without siloing.
           const witgenTimer = new Timer();
-          const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
+          const privateInputs = await resetBuilder.build(this.oracle);
           output = generateWitnesses
             ? await this.proofCreator.generateResetOutput(privateInputs)
             : await this.proofCreator.simulateReset(privateInputs);
@@ -138,93 +158,45 @@ export class PrivateKernelExecutionProver {
         }
       }
 
-      const currentExecution = executionStack.pop()!;
-
-      executionStack.push(...[...currentExecution.nestedExecutionResults].reverse());
-
-      const functionName = await this.oracle.getDebugFunctionName(
-        currentExecution.publicInputs.callContext.contractAddress,
-        currentExecution.publicInputs.callContext.functionSelector,
-      );
+      const batchSize = planner.decideBatchSize(output.publicInputs, executionStack);
+      const apps: PrivateCallData[] = [];
+      for (let i = 0; i < batchSize; i++) {
+        apps.push(await this.consumeNextApp(executionStack, executionSteps, updatedClassIdHintsMap));
+      }
 
-      executionSteps.push({
-        functionName: functionName!,
-        bytecode: currentExecution.acir,
-        witness: currentExecution.partialWitness,
-        vk: currentExecution.vk,
-        timings: {
-          witgen: currentExecution.profileResult?.timings.witgen ?? 0,
-          oracles: currentExecution.profileResult?.timings.oracles,
-        },
+      output = await this.runBatchedKernel({
+        apps,
+        firstIteration,
+        previousOutput: output,
+        txRequest,
+        isPrivateOnlyTx,
+        firstNullifierHint: executionResult.firstNullifier,
+        minRevertibleSideEffectCounter,
+        executionSteps,
+        generateWitnesses,
       });
 
-      const privateCallData = await this.createPrivateCallData(currentExecution);
-
-      if (firstIteration) {
-        const witgenTimer = new Timer();
-
-        const proofInput = new PrivateKernelInitCircuitPrivateInputs(
-          txRequest,
-          getVKTreeRoot(),
-          ProtocolContractsList,
-          privateCallData,
-          isPrivateOnlyTx,
-          executionResult.firstNullifier,
-          minRevertibleSideEffectCounter,
-        );
-        this.log.debug(
-          `Calling private kernel init with isPrivateOnly ${isPrivateOnlyTx} and firstNullifierHint ${proofInput.firstNullifierHint}`,
-        );
-
-        pushTestData('private-kernel-inputs-init', proofInput);
-
-        output = generateWitnesses
-          ? await this.proofCreator.generateInitOutput(proofInput)
-          : await this.proofCreator.simulateInit(proofInput);
-
-        executionSteps.push({
-          functionName: 'private_kernel_init',
-          bytecode: output.bytecode,
-          witness: output.outputWitness,
-          vk: output.verificationKey.keyAsBytes,
-          timings: {
-            witgen: witgenTimer.ms(),
-          },
-        });
-      } else {
-        const witgenTimer = new Timer();
-        const vkData = await this.getVkData(output.verificationKey);
-        const previousKernelData = new PrivateKernelData(output.publicInputs, vkData);
-        const proofInput = new PrivateKernelInnerCircuitPrivateInputs(previousKernelData, privateCallData);
-
-        pushTestData('private-kernel-inputs-inner', proofInput);
-        output = generateWitnesses
-          ? await this.proofCreator.generateInnerOutput(proofInput)
-          : await this.proofCreator.simulateInner(proofInput);
-
-        executionSteps.push({
-          functionName: 'private_kernel_inner',
-          bytecode: output.bytecode,
-          witness: output.outputWitness,
-          vk: output.verificationKey.keyAsBytes,
-          timings: {
-            witgen: witgenTimer.ms(),
-          },
-        });
-      }
       firstIteration = false;
     }
 
-    // Reset.
-    let resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
+    // Final reset: include siloing of note hashes, nullifiers and private logs.
+    const finalResetBuilder = new PrivateKernelResetPrivateInputsBuilder(
       output,
       [],
       noteHashNullifierCounterMap,
       splitCounter,
     );
-    while (resetBuilder.needsReset()) {
+    if (!finalResetBuilder.needsReset()) {
+      // The final reset must be performed exactly once, because each tx has at least one nullifier that requires
+      // siloing, and siloing cannot be done multiple times.
+      // While, in theory, it might be possible to silo note hashes first and then run another reset to silo nullifiers
+      // and/or private logs, we currently don't have standalone dimensions for the arrays that require siloing. As a
+      // result, all necessary siloing must be done together in a single reset.
+      // Refer to the possible combinations of dimensions in private_kernel_reset_config.json.
+      throw new Error('Nothing to reset for the final reset.');
+    } else {
       const witgenTimer = new Timer();
-      const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
+      const privateInputs = await finalResetBuilder.build(this.oracle);
       output = generateWitnesses
         ? await this.proofCreator.generateResetOutput(privateInputs)
         : await this.proofCreator.simulateReset(privateInputs);
@@ -238,8 +210,6 @@ export class PrivateKernelExecutionProver {
           witgen: witgenTimer.ms(),
         },
       });
-
-      resetBuilder = new PrivateKernelResetPrivateInputsBuilder(output, [], noteHashNullifierCounterMap, splitCounter);
     }
 
     if (output.publicInputs.feePayer.isZero() && skipFeeEnforcement) {
@@ -259,20 +229,14 @@ export class PrivateKernelExecutionProver {
     // TODO: Enable padding once we better understand the final amounts to pad to.
     const paddedSideEffectAmounts = PaddedSideEffectAmounts.empty();
 
-    // Use the aggregated includeByTimestamp set throughout the tx execution.
-    // TODO: Call `computeTxIncludeByTimestamp` to round the value down and reduce precision, improving privacy.
-    const includeByTimestampUpperBound = previousKernelData.publicInputs.includeByTimestamp;
-    const anchorBlockTimestamp = previousKernelData.publicInputs.constants.anchorBlockHeader.globalVariables.timestamp;
-    if (includeByTimestampUpperBound <= anchorBlockTimestamp) {
-      throw new Error(
-        `Include-by timestamp must be greater than the anchor block timestamp. Anchor block timestamp: ${anchorBlockTimestamp}. Include-by timestamp: ${includeByTimestampUpperBound}.`,
-      );
-    }
+    // Round the aggregated expirationTimestamp down to reduce precision and avoid leaking which private
+    // functions were called via their exact expiration offsets.
+    const expirationTimestampUpperBound = computeTxExpirationTimestamp(previousKernelData.publicInputs);
 
     const privateInputs = new PrivateKernelTailCircuitPrivateInputs(
       previousKernelData,
       paddedSideEffectAmounts,
-      includeByTimestampUpperBound,
+      expirationTimestampUpperBound,
     );
 
     const witgenTimer = new Timer();
@@ -398,7 +362,60 @@ export class PrivateKernelExecutionProver {
     );
   }
 
-  private async createPrivateCallData({ publicInputs, vk: vkAsBuffer }: PrivateCallExecutionResult) {
+  /**
+   * Pops the next app off the execution stack, pushes its nested calls back on (preserving DFS
+   * order), records the app's execution step, and returns its constructed `PrivateCallData`.
+   * Caller is responsible for ensuring the stack is non-empty.
+   */
+  private async consumeNextApp(
+    executionStack: PrivateCallExecutionResult[],
+    executionSteps: PrivateExecutionStep[],
+    updatedClassIdHintsMap: Map<string, UpdatedClassIdHints>,
+  ): Promise<PrivateCallData> {
+    const next = executionStack.pop()!;
+    executionStack.push(...[...next.nestedExecutionResults].reverse());
+
+    const functionName = await this.oracle.getDebugFunctionName(
+      next.publicInputs.callContext.contractAddress,
+      next.publicInputs.callContext.functionSelector,
+    );
+
+    executionSteps.push({
+      functionName: functionName!,
+      bytecode: next.acir,
+      witness: next.partialWitness,
+      vk: next.vk,
+      timings: {
+        witgen: next.profileResult?.timings.witgen ?? 0,
+        oracles: next.profileResult?.timings.oracles,
+      },
+    });
+
+    return await this.createPrivateCallData(next, updatedClassIdHintsMap);
+  }
+
+  /** Prefetches updated class id hints for all unique contracts in the execution tree in parallel. */
+  private async prefetchUpdatedClassIdHints(
+    executionResult: PrivateExecutionResult,
+  ): Promise<Map<string, UpdatedClassIdHints>> {
+    const allAddresses = collectNested([executionResult.entrypoint], exec => [
+      exec.publicInputs.callContext.contractAddress,
+    ]);
+    const uniqueAddresses = uniqueBy(allAddresses, a => a.toString());
+    return new Map<string, UpdatedClassIdHints>(
+      await Promise.all(
+        uniqueAddresses.map(
+          async addr =>
+            [addr.toString(), await this.oracle.getUpdatedClassIdHints(addr)] as [string, UpdatedClassIdHints],
+        ),
+      ),
+    );
+  }
+
+  private async createPrivateCallData(
+    { publicInputs, vk: vkAsBuffer }: PrivateCallExecutionResult,
+    updatedClassIdHintsMap: Map<string, UpdatedClassIdHints>,
+  ) {
     const { contractAddress, functionSelector } = publicInputs.callContext;
 
     const vkAsFields = await vkAsFieldsMegaHonk(vkAsBuffer);
@@ -414,7 +431,7 @@ export class PrivateKernelExecutionProver {
     const { artifactHash: contractClassArtifactHash, publicBytecodeCommitment: contractClassPublicBytecodeCommitment } =
       await this.oracle.getContractClassIdPreimage(currentContractClassId);
 
-    const updatedClassIdHints = await this.oracle.getUpdatedClassIdHints(contractAddress);
+    const updatedClassIdHints = updatedClassIdHintsMap.get(contractAddress.toString())!;
 
     return PrivateCallData.from({
       publicInputs,
@@ -429,4 +446,153 @@ export class PrivateKernelExecutionProver {
       }),
     });
   }
+
+  /**
+   * Runs one batched kernel iteration. Picks the right circuit variant based on whether this is
+   * the first iteration or a later one and the size of the batch.
+   */
+  private async runBatchedKernel(args: {
+    apps: PrivateCallData[];
+    firstIteration: boolean;
+    previousOutput: PrivateKernelSimulateOutput<PrivateKernelCircuitPublicInputs>;
+    txRequest: TxRequest;
+    isPrivateOnlyTx: boolean;
+    firstNullifierHint: Fr;
+    minRevertibleSideEffectCounter: number;
+    executionSteps: PrivateExecutionStep[];
+    generateWitnesses: boolean;
+  }): Promise<PrivateKernelSimulateOutput<PrivateKernelCircuitPublicInputs>> {
+    const {
+      apps,
+      firstIteration,
+      previousOutput,
+      txRequest,
+      isPrivateOnlyTx,
+      firstNullifierHint,
+      minRevertibleSideEffectCounter,
+      executionSteps,
+      generateWitnesses,
+    } = args;
+
+    const witgenTimer = new Timer();
+    let output: PrivateKernelSimulateOutput<PrivateKernelCircuitPublicInputs>;
+    let functionName: string;
+
+    if (firstIteration) {
+      const vkTreeRoot = getVKTreeRoot();
+      switch (apps.length) {
+        case 1: {
+          const proofInput = new PrivateKernelInitCircuitPrivateInputs(
+            txRequest,
+            vkTreeRoot,
+            ProtocolContractsList,
+            apps[0],
+            isPrivateOnlyTx,
+            firstNullifierHint,
+            minRevertibleSideEffectCounter,
+          );
+          this.log.debug(
+            `Calling private kernel init with isPrivateOnly ${isPrivateOnlyTx} and firstNullifierHint ${proofInput.firstNullifierHint}`,
+          );
+          pushTestData('private-kernel-inputs-init', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInitOutput(proofInput)
+            : await this.proofCreator.simulateInit(proofInput);
+          functionName = 'private_kernel_init';
+          break;
+        }
+        case 2: {
+          const proofInput = new PrivateKernelInit2CircuitPrivateInputs(
+            txRequest,
+            vkTreeRoot,
+            ProtocolContractsList,
+            apps[0],
+            apps[1],
+            isPrivateOnlyTx,
+            firstNullifierHint,
+            minRevertibleSideEffectCounter,
+          );
+          this.log.debug(
+            `Calling private kernel init_2 with isPrivateOnly ${isPrivateOnlyTx} and firstNullifierHint ${proofInput.firstNullifierHint}`,
+          );
+          pushTestData('private-kernel-inputs-init-2', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInit2Output(proofInput)
+            : await this.proofCreator.simulateInit2(proofInput);
+          functionName = 'private_kernel_init_2';
+          break;
+        }
+        case 3: {
+          const proofInput = new PrivateKernelInit3CircuitPrivateInputs(
+            txRequest,
+            vkTreeRoot,
+            ProtocolContractsList,
+            apps[0],
+            apps[1],
+            apps[2],
+            isPrivateOnlyTx,
+            firstNullifierHint,
+            minRevertibleSideEffectCounter,
+          );
+          this.log.debug(
+            `Calling private kernel init_3 with isPrivateOnly ${isPrivateOnlyTx} and firstNullifierHint ${proofInput.firstNullifierHint}`,
+          );
+          pushTestData('private-kernel-inputs-init-3', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInit3Output(proofInput)
+            : await this.proofCreator.simulateInit3(proofInput);
+          functionName = 'private_kernel_init_3';
+          break;
+        }
+        default:
+          throw new Error(`Unsupported init kernel batch size: ${apps.length}`);
+      }
+    } else {
+      const vkData = await this.getVkData(previousOutput.verificationKey);
+      const previousKernelData = new PrivateKernelData(previousOutput.publicInputs, vkData);
+      switch (apps.length) {
+        case 1: {
+          const proofInput = new PrivateKernelInnerCircuitPrivateInputs(previousKernelData, apps[0]);
+          pushTestData('private-kernel-inputs-inner', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInnerOutput(proofInput)
+            : await this.proofCreator.simulateInner(proofInput);
+          functionName = 'private_kernel_inner';
+          break;
+        }
+        case 2: {
+          const proofInput = new PrivateKernelInner2CircuitPrivateInputs(previousKernelData, apps[0], apps[1]);
+          pushTestData('private-kernel-inputs-inner-2', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInner2Output(proofInput)
+            : await this.proofCreator.simulateInner2(proofInput);
+          functionName = 'private_kernel_inner_2';
+          break;
+        }
+        case 3: {
+          const proofInput = new PrivateKernelInner3CircuitPrivateInputs(previousKernelData, apps[0], apps[1], apps[2]);
+          pushTestData('private-kernel-inputs-inner-3', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInner3Output(proofInput)
+            : await this.proofCreator.simulateInner3(proofInput);
+          functionName = 'private_kernel_inner_3';
+          break;
+        }
+        default:
+          throw new Error(`Unsupported inner kernel batch size: ${apps.length}`);
+      }
+    }
+
+    executionSteps.push({
+      functionName,
+      bytecode: output.bytecode,
+      witness: output.outputWitness,
+      vk: output.verificationKey.keyAsBytes,
+      timings: {
+        witgen: witgenTimer.ms(),
+      },
+    });
+
+    return output;
+  }
 }

package/src/private_kernel/private_kernel_oracle.ts

@@ -1,72 +1,162 @@
-import { FUNCTION_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT, VK_TREE_HEIGHT } from '@aztec/constants';
+import {
+  FUNCTION_TREE_HEIGHT,
+  NOTE_HASH_TREE_HEIGHT,
+  PUBLIC_DATA_TREE_HEIGHT,
+  UPDATES_VALUE_SIZE,
+  VK_TREE_HEIGHT,
+} from '@aztec/constants';
 import type { Fr } from '@aztec/foundation/curves/bn254';
-import type { GrumpkinScalar, Point } from '@aztec/foundation/curves/grumpkin';
+import type { GrumpkinScalar } from '@aztec/foundation/curves/grumpkin';
 import { MembershipWitness } from '@aztec/foundation/trees';
+import type { KeyStore } from '@aztec/key-store';
+import { getVKIndex, getVKSiblingPath } from '@aztec/noir-protocol-circuits-types/vk-tree';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
 import type { FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { type ContractInstanceWithAddress, computeSaltedInitializationHash } from '@aztec/stdlib/contract';
+import { DelayedPublicMutableValues, DelayedPublicMutableValuesWithHash } from '@aztec/stdlib/delayed-public-mutable';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import { UpdatedClassIdHints } from '@aztec/stdlib/kernel';
-import type { PublicKeys } from '@aztec/stdlib/keys';
 import type { NullifierMembershipWitness } from '@aztec/stdlib/trees';
+import type { BlockHeader } from '@aztec/stdlib/tx';
 import type { VerificationKeyAsFields } from '@aztec/stdlib/vks';
 
+import type { ContractStore } from '../storage/contract_store/contract_store.js';
+
 /**
  * Provides functionality needed by the private kernel for interacting with our state trees.
- * This is either PrivateKernelOracleImpl, or a mocked test implementation.
  */
-export interface PrivateKernelOracle {
+export class PrivateKernelOracle {
+  constructor(
+    private contractStore: ContractStore,
+    private keyStore: KeyStore,
+    private node: AztecNode,
+    private blockHeader: BlockHeader,
+  ) {}
+
   /** Retrieves the preimage of a contract address from the registered contract instances db. */
-  getContractAddressPreimage(address: AztecAddress): Promise<{
-    saltedInitializationHash: Fr;
-    publicKeys: PublicKeys;
-    currentContractClassId: Fr;
-    originalContractClassId: Fr;
-  }>;
+  public async getContractAddressPreimage(
+    address: AztecAddress,
+  ): Promise<ContractInstanceWithAddress & { saltedInitializationHash: Fr }> {
+    const instance = await this.contractStore.getContractInstance(address);
+    if (!instance) {
+      throw new Error(`Contract instance not found when getting address preimage. Contract address: ${address}.`);
+    }
+    return {
+      saltedInitializationHash: await computeSaltedInitializationHash(instance),
+      ...instance,
+    };
+  }
 
   /** Retrieves the preimage of a contract class id from the contract classes db. */
-  getContractClassIdPreimage(
-    contractClassId: Fr,
-  ): Promise<{ artifactHash: Fr; publicBytecodeCommitment: Fr; privateFunctionsRoot: Fr }>;
+  public async getContractClassIdPreimage(contractClassId: Fr) {
+    const contractClass = await this.contractStore.getContractClassWithPreimage(contractClassId);
+    if (!contractClass) {
+      throw new Error(`Contract class not found when getting class id preimage. Class id: ${contractClassId}.`);
+    }
+    return {
+      artifactHash: contractClass.artifactHash,
+      privateFunctionsRoot: contractClass.privateFunctionsRoot,
+      publicBytecodeCommitment: contractClass.publicBytecodeCommitment,
+    };
+  }
 
-  /**
-   * Returns a membership witness with the sibling path and leaf index in our private functions tree.
-   */
-  getFunctionMembershipWitness(
+  /** Returns a membership witness with the sibling path and leaf index in our private functions tree. */
+  public async getFunctionMembershipWitness(
     contractClassId: Fr,
     selector: FunctionSelector,
-  ): Promise<MembershipWitness<typeof FUNCTION_TREE_HEIGHT>>;
+  ): Promise<MembershipWitness<typeof FUNCTION_TREE_HEIGHT>> {
+    const membershipWitness = await this.contractStore.getFunctionMembershipWitness(contractClassId, selector);
+    if (!membershipWitness) {
+      throw new Error(
+        `Membership witness not found for contract class id ${contractClassId} and selector ${selector}.`,
+      );
+    }
+    return membershipWitness;
+  }
 
   /**
    * Returns a membership witness with the sibling path and leaf index in our protocol VK indexed merkle tree.
    * Used to validate the previous kernel's verification key.
    */
-  getVkMembershipWitness(vk: VerificationKeyAsFields): Promise<MembershipWitness<typeof VK_TREE_HEIGHT>>;
+  public getVkMembershipWitness(vk: VerificationKeyAsFields): Promise<MembershipWitness<typeof VK_TREE_HEIGHT>> {
+    const leafIndex = getVKIndex(vk);
+    return Promise.resolve(new MembershipWitness(VK_TREE_HEIGHT, BigInt(leafIndex), getVKSiblingPath(leafIndex)));
+  }
 
-  /**
-   * Returns a membership witness with the sibling path and leaf index in our private function indexed merkle tree.
-   */ getNoteHashMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>>;
+  /** Returns a membership witness with the sibling path and leaf index in our note hash tree. */
+  async getNoteHashMembershipWitness(
+    noteHash: Fr,
+  ): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined> {
+    return this.node.getNoteHashMembershipWitness(await this.blockHeader.hash(), noteHash);
+  }
 
-  /**
-   * Returns a membership witness with the sibling path and leaf index in our nullifier indexed merkle tree.
-   */
-  getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitness | undefined>;
-  /**
-   * Returns the root of our note hash merkle tree.
-   */
-  getNoteHashTreeRoot(): Promise<Fr>;
+  /** Returns a membership witness with the sibling path and leaf index in our nullifier indexed merkle tree. */
+  async getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitness | undefined> {
+    return this.node.getNullifierMembershipWitness(await this.blockHeader.hash(), nullifier);
+  }
+
+  /** Returns the root of our note hash merkle tree. */
+  getNoteHashTreeRoot(): Fr {
+    return this.blockHeader.state.partial.noteHashTree.root;
+  }
 
   /**
-   * Retrieves the sk_m corresponding to the pk_m.
-   * @throws If the provided public key is not associated with any of the registered accounts.
-   * @param pkM - The master public key to get secret key for.
+   * Retrieves the sk_m corresponding to the pk_m hash.
+   * @throws If the provided hash is not associated with any of the registered accounts.
+   * @param masterPublicKeyHash - The master public key hash to get secret key for.
    * @returns A Promise that resolves to sk_m.
    * @dev Used when feeding the sk_m to the kernel circuit for keys verification.
    */
-  getMasterSecretKey(masterPublicKey: Point): Promise<GrumpkinScalar>;
+  public getMasterSecretKey(masterPublicKeyHash: Fr): Promise<GrumpkinScalar> {
+    return this.keyStore.getMasterSecretKey(masterPublicKeyHash);
+  }
 
   /** Use debug data to get the function name corresponding to a selector. */
-  getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector): Promise<string | undefined>;
+  public getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector): Promise<string | undefined> {
+    return this.contractStore.getDebugFunctionName(contractAddress, selector);
+  }
+
+  /**
+   * Returns a membership witness and leaf index to our public data indexed merkle tree,
+   * along with an associated DelayedPublicMutable containing the class ID to update.
+   */
+  public async getUpdatedClassIdHints(contractAddress: AztecAddress): Promise<UpdatedClassIdHints> {
+    const { delayedPublicMutableSlot, delayedPublicMutableHashSlot } =
+      await DelayedPublicMutableValuesWithHash.getContractUpdateSlots(contractAddress);
+
+    const hashLeafSlot = await computePublicDataTreeLeafSlot(
+      ProtocolContractAddress.ContractInstanceRegistry,
+      delayedPublicMutableHashSlot,
+    );
+    const blockHash = await this.blockHeader.hash();
+
+    const updatedClassIdWitness = await this.node.getPublicDataWitness(blockHash, hashLeafSlot);
+
+    if (!updatedClassIdWitness) {
+      throw new Error(`No public data tree witness found for ${hashLeafSlot}`);
+    }
+
+    // In an indexed merkle tree, getPublicDataWitness returns a leaf whose slot matches our query
+    // only if the slot has been written to. Otherwise, it returns the "low leaf" predecessor, whose
+    // slot will differ. Most contracts are never updated, so we can skip the readFromTree call
+    // (which triggers multiple RPC calls) and return empty values directly.
+    const readStorage = (storageSlot: Fr) =>
+      this.node.getPublicStorageAt(blockHash, ProtocolContractAddress.ContractInstanceRegistry, storageSlot);
+    const slotExists = updatedClassIdWitness.leafPreimage.leaf.slot.equals(hashLeafSlot);
+    const delayedPublicMutableValues = slotExists
+      ? await DelayedPublicMutableValues.readFromTree(delayedPublicMutableSlot, readStorage)
+      : DelayedPublicMutableValues.empty(UPDATES_VALUE_SIZE);
 
-  /** Returns a membership witness and leaf index to our public data indexed merkle tree,
-   * along with an associated DelayedPublicMutable containing the class ID to update. */
-  getUpdatedClassIdHints(contractAddress: AztecAddress): Promise<UpdatedClassIdHints>;
+    return new UpdatedClassIdHints(
+      new MembershipWitness(
+        PUBLIC_DATA_TREE_HEIGHT,
+        updatedClassIdWitness.index,
+        updatedClassIdWitness.siblingPath.toTuple(),
+      ),
+      updatedClassIdWitness.leafPreimage,
+      delayedPublicMutableValues,
+    );
+  }
 }