@aztec/pxe 0.0.0-test.1 → 0.0.1-commit.1142ef1

package/src/private_kernel/private_kernel_execution_prover.ts

@@ -0,0 +1,430 @@
+import { vkAsFieldsMegaHonk } from '@aztec/foundation/crypto/keys';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { 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';
+import { ProtocolContractsList } from '@aztec/protocol-contracts';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { PrivateKernelProver } from '@aztec/stdlib/interfaces/client';
+import {
+  HidingKernelToPublicPrivateInputs,
+  HidingKernelToRollupPrivateInputs,
+  PaddedSideEffectAmounts,
+  PrivateCallData,
+  type PrivateExecutionStep,
+  PrivateKernelCircuitPublicInputs,
+  PrivateKernelData,
+  type PrivateKernelExecutionProofOutput,
+  PrivateKernelInitCircuitPrivateInputs,
+  PrivateKernelInnerCircuitPrivateInputs,
+  type PrivateKernelSimulateOutput,
+  PrivateKernelTailCircuitPrivateInputs,
+  type PrivateKernelTailCircuitPublicInputs,
+  PrivateVerificationKeyHints,
+} from '@aztec/stdlib/kernel';
+import { ChonkProof, ChonkProofWithPublicInputs } from '@aztec/stdlib/proofs';
+import {
+  type PrivateCallExecutionResult,
+  type PrivateExecutionResult,
+  TxRequest,
+  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 type { PrivateKernelOracle } from './private_kernel_oracle.js';
+
+const NULL_SIMULATE_OUTPUT: PrivateKernelSimulateOutput<PrivateKernelCircuitPublicInputs> = {
+  publicInputs: PrivateKernelCircuitPublicInputs.empty(),
+  verificationKey: VerificationKeyData.empty(),
+  outputWitness: new Map(),
+  bytecode: Buffer.from([]),
+};
+
+export interface PrivateKernelExecutionProverConfig {
+  simulate: boolean;
+  skipFeeEnforcement: boolean;
+  profileMode: 'gates' | 'execution-steps' | 'full' | 'none';
+}
+
+/**
+ * The PrivateKernelExecutionProver class is responsible for taking a transaction request and sequencing the
+ * the execution of the private functions within, sequenced with private kernel "glue" to check protocol rules.
+ * The result can be a chonk proof of the private transaction portion, or just a simulation that can e.g.
+ * inform state tree updates.
+ */
+export class PrivateKernelExecutionProver {
+  private log = createLogger('pxe:private-kernel-execution-prover');
+
+  constructor(
+    private oracle: PrivateKernelOracle,
+    private proofCreator: PrivateKernelProver,
+    private fakeProofs = false,
+  ) {}
+
+  /**
+   * Generate a proof for a given transaction request and execution result.
+   * The function iterates through the nested executions in the execution result, creates private call data,
+   * and generates a proof using the provided ProofCreator instance. It also maintains an index of new notes
+   * created during the execution and returns them as a part of the KernelProverOutput.
+   *
+   * @param txRequest - The authenticated transaction request object.
+   * @param executionResult - The execution result object containing nested executions and preimages.
+   * @param profile - Set true to profile the gate count for each circuit
+   * @returns A Promise that resolves to a KernelProverOutput object containing proof, public inputs, and output notes.
+   */
+  async proveWithKernels(
+    txRequest: TxRequest,
+    executionResult: PrivateExecutionResult,
+    { simulate, skipFeeEnforcement, profileMode }: PrivateKernelExecutionProverConfig = {
+      simulate: false,
+      skipFeeEnforcement: false,
+      profileMode: 'none',
+    },
+  ): Promise<PrivateKernelExecutionProofOutput<PrivateKernelTailCircuitPublicInputs>> {
+    const skipProofGeneration = this.fakeProofs || simulate;
+    const generateWitnesses = !skipProofGeneration || profileMode !== 'none';
+
+    const timer = new Timer();
+
+    const isPrivateOnlyTx = executionResult.publicFunctionCalldata.length === 0;
+
+    // Initialize an executionStack, beginning with the PrivateCallExecutionResult
+    // of the entrypoint function of the tx.
+    const executionStack = [executionResult.entrypoint];
+    let firstIteration = true;
+
+    let output = NULL_SIMULATE_OUTPUT;
+
+    const executionSteps: PrivateExecutionStep[] = [];
+
+    const noteHashNullifierCounterMap = collectNoteHashNullifierCounterMap(executionResult);
+    const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(executionResult);
+    const splitCounter = isPrivateOnlyTx ? 0 : minRevertibleSideEffectCounter;
+
+    while (executionStack.length) {
+      if (!firstIteration) {
+        let resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
+          output,
+          executionStack,
+          noteHashNullifierCounterMap,
+          splitCounter,
+        );
+        while (resetBuilder.needsReset()) {
+          const witgenTimer = new Timer();
+          const privateInputs = await resetBuilder.build(this.oracle);
+          output = generateWitnesses
+            ? await this.proofCreator.generateResetOutput(privateInputs)
+            : await this.proofCreator.simulateReset(privateInputs);
+          executionSteps.push({
+            functionName: 'private_kernel_reset',
+            bytecode: output.bytecode,
+            witness: output.outputWitness,
+            vk: output.verificationKey.keyAsBytes,
+            timings: {
+              witgen: witgenTimer.ms(),
+            },
+          });
+          resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
+            output,
+            executionStack,
+            noteHashNullifierCounterMap,
+            splitCounter,
+          );
+        }
+      }
+
+      const currentExecution = executionStack.pop()!;
+
+      executionStack.push(...[...currentExecution.nestedExecutionResults].reverse());
+
+      const functionName = await this.oracle.getDebugFunctionName(
+        currentExecution.publicInputs.callContext.contractAddress,
+        currentExecution.publicInputs.callContext.functionSelector,
+      );
+
+      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,
+        },
+      });
+
+      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(
+      output,
+      [],
+      noteHashNullifierCounterMap,
+      splitCounter,
+    );
+    while (resetBuilder.needsReset()) {
+      const witgenTimer = new Timer();
+      const privateInputs = await resetBuilder.build(this.oracle);
+      output = generateWitnesses
+        ? await this.proofCreator.generateResetOutput(privateInputs)
+        : await this.proofCreator.simulateReset(privateInputs);
+
+      executionSteps.push({
+        functionName: 'private_kernel_reset',
+        bytecode: output.bytecode,
+        witness: output.outputWitness,
+        vk: output.verificationKey.keyAsBytes,
+        timings: {
+          witgen: witgenTimer.ms(),
+        },
+      });
+
+      resetBuilder = new PrivateKernelResetPrivateInputsBuilder(output, [], noteHashNullifierCounterMap, splitCounter);
+    }
+
+    if (output.publicInputs.feePayer.isZero() && skipFeeEnforcement) {
+      if (!skipProofGeneration) {
+        throw new Error('Fee payment must be enforced when creating real proof.');
+      }
+      output.publicInputs.feePayer = new AztecAddress(Fr.MAX_FIELD_VALUE);
+    }
+    // Private tail.
+    const vkData = await this.getVkData(output.verificationKey);
+    const previousKernelData = new PrivateKernelData(output.publicInputs, vkData);
+
+    this.log.debug(
+      `Calling private kernel tail with hwm ${previousKernelData.publicInputs.minRevertibleSideEffectCounter}`,
+    );
+
+    // 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}.`,
+      );
+    }
+
+    const privateInputs = new PrivateKernelTailCircuitPrivateInputs(
+      previousKernelData,
+      paddedSideEffectAmounts,
+      includeByTimestampUpperBound,
+    );
+
+    const witgenTimer = new Timer();
+    const tailOutput = generateWitnesses
+      ? await this.proofCreator.generateTailOutput(privateInputs)
+      : await this.proofCreator.simulateTail(privateInputs);
+
+    executionSteps.push({
+      functionName: 'private_kernel_tail',
+      bytecode: tailOutput.bytecode,
+      witness: tailOutput.outputWitness,
+      vk: tailOutput.verificationKey.keyAsBytes,
+      timings: {
+        witgen: witgenTimer.ms(),
+      },
+    });
+
+    // Hiding kernel is only executed if we are generating witnesses.
+    // For simulation, we can end with the tail, since the Hiding kernel will simply return the same tail output.
+    if (generateWitnesses) {
+      const previousKernelVkData = await this.getVkData(tailOutput.verificationKey);
+
+      const witgenTimer = new Timer();
+      let hidingOutput: PrivateKernelSimulateOutput<PrivateKernelTailCircuitPublicInputs>;
+      if (tailOutput.publicInputs.forPublic) {
+        const privateInputs = new HidingKernelToPublicPrivateInputs(
+          tailOutput.publicInputs.toPrivateToPublicKernelCircuitPublicInputs(),
+          previousKernelVkData,
+        );
+        hidingOutput = await this.proofCreator.generateHidingToPublicOutput(privateInputs);
+      } else {
+        const privateInputs = new HidingKernelToRollupPrivateInputs(
+          tailOutput.publicInputs.toPrivateToRollupKernelCircuitPublicInputs(),
+          previousKernelVkData,
+        );
+        hidingOutput = await this.proofCreator.generateHidingToRollupOutput(privateInputs);
+      }
+
+      executionSteps.push({
+        functionName: 'hiding_kernel',
+        bytecode: hidingOutput.bytecode,
+        witness: hidingOutput.outputWitness,
+        vk: hidingOutput.verificationKey.keyAsBytes,
+        timings: {
+          witgen: witgenTimer.ms(),
+        },
+      });
+    }
+
+    if (profileMode == 'gates' || profileMode == 'full') {
+      for (const entry of executionSteps) {
+        const gateCountTimer = new Timer();
+        const gateCount = await this.proofCreator.computeGateCountForCircuit(entry.bytecode, entry.functionName);
+        entry.gateCount = gateCount;
+        entry.timings.gateCount = gateCountTimer.ms();
+      }
+    }
+    if (profileMode === 'gates') {
+      for (const entry of executionSteps) {
+        // These buffers are often a few megabytes in size - prevent accidentally serializing them if not requested.
+        entry.bytecode = Buffer.from([]);
+        entry.witness = new Map();
+      }
+    }
+
+    if (generateWitnesses) {
+      this.log.info(`Private kernel witness generation took ${timer.ms()}ms`);
+    }
+
+    let chonkProof: ChonkProof;
+    // TODO(#7368) how do we 'bincode' encode these inputs?
+    let provingTime;
+    if (!skipProofGeneration) {
+      const provingTimer = new Timer();
+      const proofWithPublicInputs = await this.proofCreator.createChonkProof(executionSteps);
+      provingTime = provingTimer.ms();
+      this.ensurePublicInputsMatch(proofWithPublicInputs, tailOutput.publicInputs);
+      chonkProof = proofWithPublicInputs.removePublicInputs();
+    } else {
+      chonkProof = ChonkProof.random();
+    }
+
+    return {
+      publicInputs: tailOutput.publicInputs,
+      executionSteps,
+      chonkProof,
+      timings: provingTime ? { proving: provingTime } : undefined,
+    };
+  }
+
+  /**
+   * Checks that the public inputs of the chonk proof match the public inputs of the tail circuit.
+   * This can only mismatch if there is a circuit / noir / bb bug.
+   * @param chonkProof - The chonk proof with public inputs.
+   * @param tailPublicInputs - The public inputs resulting from witness generation of the tail circuit.
+   */
+  private ensurePublicInputsMatch(
+    chonkProof: ChonkProofWithPublicInputs,
+    tailPublicInputs: PrivateKernelTailCircuitPublicInputs,
+  ) {
+    const serializedChonkProofPublicInputs = chonkProof.getPublicInputs();
+    const serializedTailPublicInputs = tailPublicInputs.publicInputs().toFields();
+    if (serializedChonkProofPublicInputs.length !== serializedTailPublicInputs.length) {
+      throw new Error(
+        `Public inputs length mismatch: ${serializedChonkProofPublicInputs.length} !== ${serializedTailPublicInputs.length}`,
+      );
+    }
+    if (
+      !serializedChonkProofPublicInputs.every((input: Fr, index: number) =>
+        input.equals(serializedTailPublicInputs[index]),
+      )
+    ) {
+      throw new Error(`Public inputs mismatch between kernel and chonk proof`);
+    }
+  }
+
+  private async getVkData(verificationKey: VerificationKeyData) {
+    const previousVkMembershipWitness = await this.oracle.getVkMembershipWitness(verificationKey.keyAsFields);
+    return new VkData(
+      verificationKey,
+      Number(previousVkMembershipWitness.leafIndex),
+      previousVkMembershipWitness.siblingPath,
+    );
+  }
+
+  private async createPrivateCallData({ publicInputs, vk: vkAsBuffer }: PrivateCallExecutionResult) {
+    const { contractAddress, functionSelector } = publicInputs.callContext;
+
+    const vkAsFields = await vkAsFieldsMegaHonk(vkAsBuffer);
+    const vk = await VerificationKeyAsFields.fromKey(vkAsFields);
+
+    const { currentContractClassId, publicKeys, saltedInitializationHash } =
+      await this.oracle.getContractAddressPreimage(contractAddress);
+    const functionLeafMembershipWitness = await this.oracle.getFunctionMembershipWitness(
+      currentContractClassId,
+      functionSelector,
+    );
+
+    const { artifactHash: contractClassArtifactHash, publicBytecodeCommitment: contractClassPublicBytecodeCommitment } =
+      await this.oracle.getContractClassIdPreimage(currentContractClassId);
+
+    const updatedClassIdHints = await this.oracle.getUpdatedClassIdHints(contractAddress);
+
+    return PrivateCallData.from({
+      publicInputs,
+      vk,
+      verificationKeyHints: PrivateVerificationKeyHints.from({
+        publicKeys,
+        contractClassArtifactHash,
+        contractClassPublicBytecodeCommitment,
+        saltedInitializationHash,
+        functionLeafMembershipWitness,
+        updatedClassIdHints,
+      }),
+    });
+  }
+}

package/src/{kernel_prover/proving_data_oracle.ts → private_kernel/private_kernel_oracle.ts}

@@ -1,6 +1,7 @@
-import type { FUNCTION_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT, VK_TREE_HEIGHT } from '@aztec/constants';
-import type { Fr, GrumpkinScalar, Point } from '@aztec/foundation/fields';
-import type { MembershipWitness } from '@aztec/foundation/trees';
+import { FUNCTION_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT, VK_TREE_HEIGHT } from '@aztec/constants';
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import type { GrumpkinScalar, Point } from '@aztec/foundation/curves/grumpkin';
+import { MembershipWitness } from '@aztec/foundation/trees';
 import type { FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { UpdatedClassIdHints } from '@aztec/stdlib/kernel';
@@ -9,10 +10,10 @@ import type { NullifierMembershipWitness } from '@aztec/stdlib/trees';
 import type { VerificationKeyAsFields } from '@aztec/stdlib/vks';
 
 /**
- * Provides functionality to fetch membership witnesses for verification keys,
- * contract addresses, and function selectors in their respective merkle trees.
+ * Provides functionality needed by the private kernel for interacting with our state trees.
+ * This is either PrivateKernelOracleImpl, or a mocked test implementation.
  */
-export interface ProvingDataOracle {
+export interface PrivateKernelOracle {
   /** Retrieves the preimage of a contract address from the registered contract instances db. */
   getContractAddressPreimage(address: AztecAddress): Promise<{
     saltedInitializationHash: Fr;
@@ -27,13 +28,7 @@ export interface ProvingDataOracle {
   ): Promise<{ artifactHash: Fr; publicBytecodeCommitment: Fr; privateFunctionsRoot: Fr }>;
 
   /**
-   * Retrieve the function membership witness for the given contract class and function selector.
-   * The function membership witness represents a proof that the function belongs to the specified contract.
-   * Throws an error if the contract address or function selector is unknown.
-   *
-   * @param contractClassId - The id of the class.
-   * @param selector - The function selector.
-   * @returns A promise that resolves with the MembershipWitness instance for the specified contract's function.
+   * Returns a membership witness with the sibling path and leaf index in our private functions tree.
    */
   getFunctionMembershipWitness(
     contractClassId: Fr,
@@ -41,30 +36,22 @@ export interface ProvingDataOracle {
   ): Promise<MembershipWitness<typeof FUNCTION_TREE_HEIGHT>>;
 
   /**
-   * Retrieve the membership witness corresponding to a verification key.
-   * This function currently returns a random membership witness of the specified height,
-   * which is a placeholder implementation until a concrete membership witness calculation
-   * is implemented.
-   *
-   * @param vk - The VerificationKey for which the membership witness is needed.
-   * @returns A Promise that resolves to the MembershipWitness instance.
+   * 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>>;
 
   /**
-   * Get the note membership witness for a note in the note hash tree at the given leaf index.
-   *
-   * @param leafIndex - The leaf index of the note in the note hash tree.
-   * @returns the MembershipWitness for the note.
+   * 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>>;
+  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined>;
 
+  /**
+   * Returns a membership witness with the sibling path and leaf index in our nullifier indexed merkle tree.
+   */
   getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitness | undefined>;
-
   /**
-   * Get the root of the note hash tree.
-   *
-   * @returns the root of the note hash tree.
+   * Returns the root of our note hash merkle tree.
    */
   getNoteHashTreeRoot(): Promise<Fr>;
 
@@ -77,7 +64,10 @@ export interface ProvingDataOracle {
    */
   getMasterSecretKey(masterPublicKey: Point): Promise<GrumpkinScalar>;
 
+  /** Use debug data to get the function name corresponding to a selector. */
   getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector): Promise<string | undefined>;
 
+  /** 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>;
 }

package/src/{kernel_oracle/index.ts → private_kernel/private_kernel_oracle_impl.ts}

@@ -1,41 +1,45 @@
-import { type NOTE_HASH_TREE_HEIGHT, PUBLIC_DATA_TREE_HEIGHT, VK_TREE_HEIGHT } from '@aztec/constants';
-import type { Fr, GrumpkinScalar, Point } from '@aztec/foundation/fields';
+import { NOTE_HASH_TREE_HEIGHT, PUBLIC_DATA_TREE_HEIGHT, VK_TREE_HEIGHT } from '@aztec/constants';
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import type { GrumpkinScalar, Point } from '@aztec/foundation/curves/grumpkin';
 import { createLogger } from '@aztec/foundation/log';
-import type { Tuple } from '@aztec/foundation/serialize';
 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 { L2BlockNumber } from '@aztec/stdlib/block';
+import type { BlockParameter } from '@aztec/stdlib/block';
 import { computeContractClassIdPreimage, 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 { SharedMutableValues, SharedMutableValuesWithHash } from '@aztec/stdlib/shared-mutable';
 import type { NullifierMembershipWitness } from '@aztec/stdlib/trees';
 import type { VerificationKeyAsFields } from '@aztec/stdlib/vks';
 
-import type { ContractDataProvider } from '../storage/contract_data_provider/contract_data_provider.js';
-import type { ProvingDataOracle } from './../kernel_prover/proving_data_oracle.js';
+import type { ContractStore } from '../storage/contract_store/contract_store.js';
+import type { PrivateKernelOracle } from './private_kernel_oracle.js';
 
 // TODO: Block number should not be "latest".
 // It should be fixed at the time the proof is being simulated. I.e., it should be the same as the value defined in the constant data.
+
 /**
  * A data oracle that provides information needed for simulating a transaction.
  */
-export class KernelOracle implements ProvingDataOracle {
+export class PrivateKernelOracleImpl implements PrivateKernelOracle {
   constructor(
-    private contractDataProvider: ContractDataProvider,
+    private contractStore: ContractStore,
     private keyStore: KeyStore,
     private node: AztecNode,
-    private blockNumber: L2BlockNumber = 'latest',
+    private blockNumber: BlockParameter = 'latest',
     private log = createLogger('pxe:kernel_oracle'),
   ) {}
 
   public async getContractAddressPreimage(address: AztecAddress) {
-    const instance = await this.contractDataProvider.getContractInstance(address);
+    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,
@@ -43,12 +47,21 @@ export class KernelOracle implements ProvingDataOracle {
   }
 
   public async getContractClassIdPreimage(contractClassId: Fr) {
-    const contractClass = await this.contractDataProvider.getContractClass(contractClassId);
+    const contractClass = await this.contractStore.getContractClass(contractClassId);
+    if (!contractClass) {
+      throw new Error(`Contract class not found when getting class id preimage. Class id: ${contractClassId}.`);
+    }
     return computeContractClassIdPreimage(contractClass);
   }
 
   public async getFunctionMembershipWitness(contractClassId: Fr, selector: FunctionSelector) {
-    return await this.contractDataProvider.getFunctionMembershipWitness(contractClassId, selector);
+    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;
   }
 
   public getVkMembershipWitness(vk: VerificationKeyAsFields) {
@@ -56,13 +69,8 @@ export class KernelOracle implements ProvingDataOracle {
     return Promise.resolve(new MembershipWitness(VK_TREE_HEIGHT, BigInt(leafIndex), getVKSiblingPath(leafIndex)));
   }
 
-  async getNoteHashMembershipWitness(leafIndex: bigint): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>> {
-    const path = await this.node.getNoteHashSiblingPath(this.blockNumber, leafIndex);
-    return new MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>(
-      path.pathSize,
-      leafIndex,
-      path.toFields() as Tuple<Fr, typeof NOTE_HASH_TREE_HEIGHT>,
-    );
+  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined> {
+    return this.node.getNoteHashMembershipWitness(this.blockNumber, noteHash);
   }
 
   getNullifierMembershipWitness(nullifier: Fr): Promise<NullifierMembershipWitness | undefined> {
@@ -82,27 +90,29 @@ export class KernelOracle implements ProvingDataOracle {
   }
 
   public getDebugFunctionName(contractAddress: AztecAddress, selector: FunctionSelector): Promise<string> {
-    return this.contractDataProvider.getDebugFunctionName(contractAddress, selector);
+    return this.contractStore.getDebugFunctionName(contractAddress, selector);
   }
 
   public async getUpdatedClassIdHints(contractAddress: AztecAddress): Promise<UpdatedClassIdHints> {
-    const { sharedMutableSlot, sharedMutableHashSlot } = await SharedMutableValuesWithHash.getContractUpdateSlots(
-      contractAddress,
-    );
+    const { delayedPublicMutableSlot, delayedPublicMutableHashSlot } =
+      await DelayedPublicMutableValuesWithHash.getContractUpdateSlots(contractAddress);
 
     const hashLeafSlot = await computePublicDataTreeLeafSlot(
-      ProtocolContractAddress.ContractInstanceDeployer,
-      sharedMutableHashSlot,
+      ProtocolContractAddress.ContractInstanceRegistry,
+      delayedPublicMutableHashSlot,
     );
-    const updatedClassIdWitness = await this.node.getPublicDataTreeWitness(this.blockNumber, hashLeafSlot);
+    const updatedClassIdWitness = await this.node.getPublicDataWitness(this.blockNumber, hashLeafSlot);
 
     if (!updatedClassIdWitness) {
       throw new Error(`No public data tree witness found for ${hashLeafSlot}`);
     }
 
     const readStorage = (storageSlot: Fr) =>
-      this.node.getPublicStorageAt(this.blockNumber, ProtocolContractAddress.ContractInstanceDeployer, storageSlot);
-    const sharedMutableValues = await SharedMutableValues.readFromTree(sharedMutableSlot, readStorage);
+      this.node.getPublicStorageAt(this.blockNumber, ProtocolContractAddress.ContractInstanceRegistry, storageSlot);
+    const delayedPublicMutableValues = await DelayedPublicMutableValues.readFromTree(
+      delayedPublicMutableSlot,
+      readStorage,
+    );
 
     return new UpdatedClassIdHints(
       new MembershipWitness(
@@ -111,7 +121,7 @@ export class KernelOracle implements ProvingDataOracle {
         updatedClassIdWitness.siblingPath.toTuple(),
       ),
       updatedClassIdWitness.leafPreimage,
-      sharedMutableValues,
+      delayedPublicMutableValues,
     );
   }
 }