@aztec/pxe 0.0.1-commit.fffb133c → 0.0.1-dev

package/src/private_kernel/hints/test_utils.ts

@@ -0,0 +1,325 @@
+import {
+  MAX_KEY_VALIDATION_REQUESTS_PER_CALL,
+  MAX_KEY_VALIDATION_REQUESTS_PER_TX,
+  MAX_NOTE_HASHES_PER_CALL,
+  MAX_NOTE_HASHES_PER_TX,
+  MAX_NOTE_HASH_READ_REQUESTS_PER_CALL,
+  MAX_NOTE_HASH_READ_REQUESTS_PER_TX,
+  MAX_NULLIFIERS_PER_CALL,
+  MAX_NULLIFIERS_PER_TX,
+  MAX_NULLIFIER_READ_REQUESTS_PER_CALL,
+  MAX_NULLIFIER_READ_REQUESTS_PER_TX,
+  MAX_PRIVATE_LOGS_PER_CALL,
+  MAX_PRIVATE_LOGS_PER_TX,
+} from '@aztec/constants';
+import { makeTuple } from '@aztec/foundation/array';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { Point } from '@aztec/foundation/curves/grumpkin';
+import type { Serializable } from '@aztec/foundation/serialize';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import {
+  ClaimedLengthArray,
+  KeyValidationRequest,
+  KeyValidationRequestAndSeparator,
+  NoteHash,
+  Nullifier,
+  PrivateCircuitPublicInputs,
+  PrivateKernelCircuitPublicInputs,
+  type PrivateKernelSimulateOutput,
+  ReadRequest,
+  ScopedKeyValidationRequestAndSeparator,
+  ScopedNoteHash,
+  ScopedNullifier,
+  ScopedReadRequest,
+} from '@aztec/stdlib/kernel';
+import { PrivateLogData, ScopedPrivateLogData } from '@aztec/stdlib/kernel';
+import { PrivateLog } from '@aztec/stdlib/logs';
+import { PrivateCallExecutionResult } from '@aztec/stdlib/tx';
+import { VerificationKeyData } from '@aztec/stdlib/vks';
+
+const DEFAULT_CONTRACT_ADDRESS = AztecAddress.fromBigInt(987654n);
+
+/**
+ * Builds a ClaimedLengthArray from a list of items, padding to the required size.
+ */
+function makeClaimed<T extends Serializable, N extends number>(items: T[], emptyFactory: { empty(): T }, maxSize: N) {
+  const padded = makeTuple(maxSize, i => items[i] ?? emptyFactory.empty());
+  return new ClaimedLengthArray<T, typeof maxSize>(padded, items.length);
+}
+
+/** Builder for PrivateKernelCircuitPublicInputs with fluent API for adding side effects. */
+export class PrivateKernelCircuitPublicInputsBuilder {
+  private noteHashes: ScopedNoteHash[] = [];
+  private nullifiers: ScopedNullifier[] = [];
+  private noteHashReadRequests: ScopedReadRequest[] = [];
+  private nullifierReadRequests: ScopedReadRequest[] = [];
+  private keyValidationRequests: ScopedKeyValidationRequestAndSeparator[] = [];
+  private privateLogs: ScopedPrivateLogData[] = [];
+  private nextCounter: number;
+
+  constructor(
+    private contractAddress: AztecAddress = DEFAULT_CONTRACT_ADDRESS,
+    startCounter = 1,
+  ) {
+    this.nextCounter = startCounter;
+  }
+
+  private getCounter(sideEffectCounter?: number): number {
+    if (sideEffectCounter !== undefined) {
+      this.nextCounter = sideEffectCounter + 1;
+      return sideEffectCounter;
+    }
+    return this.nextCounter++;
+  }
+
+  /** Adds a note hash to the accumulated data. Defaults are generated randomly. */
+  addNoteHash(opts?: { value?: Fr; counter?: number; contractAddress?: AztecAddress }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    const addr = opts?.contractAddress ?? this.contractAddress;
+    this.noteHashes.push(new NoteHash(value, counter).scope(addr));
+    return this;
+  }
+
+  /** Adds a nullifier to the accumulated data. Defaults are generated randomly. */
+  addNullifier(opts?: { value?: Fr; noteHash?: Fr; counter?: number; contractAddress?: AztecAddress }): this {
+    const value = opts?.value ?? Fr.random();
+    const noteHash = opts?.noteHash ?? Fr.ZERO;
+    const counter = this.getCounter(opts?.counter);
+    const addr = opts?.contractAddress ?? this.contractAddress;
+    this.nullifiers.push(new Nullifier(value, noteHash, counter).scope(addr));
+    return this;
+  }
+
+  /** Adds a pending note hash read request (non-empty contract address, can match a pending note hash). */
+  addPendingNoteHashReadRequest(opts?: { value?: Fr; counter?: number; contractAddress?: AztecAddress }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    const addr = opts?.contractAddress ?? this.contractAddress;
+    this.noteHashReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), addr));
+    return this;
+  }
+
+  /** Adds a settled note hash read request (empty contract address, resolved against the note hash tree). */
+  addSettledNoteHashReadRequest(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.noteHashReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), AztecAddress.ZERO));
+    return this;
+  }
+
+  /** Adds a pending nullifier read request (non-empty contract address, can match a pending nullifier). */
+  addPendingNullifierReadRequest(opts?: { value?: Fr; counter?: number; contractAddress?: AztecAddress }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    const addr = opts?.contractAddress ?? this.contractAddress;
+    this.nullifierReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), addr));
+    return this;
+  }
+
+  /** Adds a settled nullifier read request (empty contract address, resolved against the nullifier tree). */
+  addSettledNullifierReadRequest(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.nullifierReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), AztecAddress.ZERO));
+    return this;
+  }
+
+  /** Adds a key validation request to validation requests. */
+  addKeyValidationRequest(opts?: { contractAddress?: AztecAddress }): this {
+    const addr = opts?.contractAddress ?? this.contractAddress;
+    this.keyValidationRequests.push(
+      new ScopedKeyValidationRequestAndSeparator(
+        new KeyValidationRequestAndSeparator(
+          new KeyValidationRequest(new Point(Fr.random(), Fr.random(), false), Fr.random()),
+          Fr.random(),
+        ),
+        addr,
+      ),
+    );
+    return this;
+  }
+
+  /** Adds a private log to the accumulated data. Defaults are generated randomly. */
+  addPrivateLog(opts?: { noteHashCounter?: number; counter?: number; contractAddress?: AztecAddress }): this {
+    const noteHashCounter = opts?.noteHashCounter ?? 0;
+    const counter = this.getCounter(opts?.counter);
+    const addr = opts?.contractAddress ?? this.contractAddress;
+    this.privateLogs.push(
+      new ScopedPrivateLogData(new PrivateLogData(PrivateLog.empty(), noteHashCounter, counter), addr),
+    );
+    return this;
+  }
+
+  /** Builds the PrivateKernelCircuitPublicInputs with all added side effects. */
+  build(): PrivateKernelCircuitPublicInputs {
+    const publicInputs = PrivateKernelCircuitPublicInputs.empty();
+    publicInputs.end.noteHashes = makeClaimed(this.noteHashes, ScopedNoteHash, MAX_NOTE_HASHES_PER_TX);
+    publicInputs.end.nullifiers = makeClaimed(this.nullifiers, ScopedNullifier, MAX_NULLIFIERS_PER_TX);
+    publicInputs.end.privateLogs = makeClaimed(this.privateLogs, ScopedPrivateLogData, MAX_PRIVATE_LOGS_PER_TX);
+    publicInputs.validationRequests.noteHashReadRequests = makeClaimed(
+      this.noteHashReadRequests,
+      ScopedReadRequest,
+      MAX_NOTE_HASH_READ_REQUESTS_PER_TX,
+    );
+    publicInputs.validationRequests.nullifierReadRequests = makeClaimed(
+      this.nullifierReadRequests,
+      ScopedReadRequest,
+      MAX_NULLIFIER_READ_REQUESTS_PER_TX,
+    );
+    publicInputs.validationRequests.scopedKeyValidationRequestsAndSeparators = makeClaimed(
+      this.keyValidationRequests,
+      ScopedKeyValidationRequestAndSeparator,
+      MAX_KEY_VALIDATION_REQUESTS_PER_TX,
+    );
+    return publicInputs;
+  }
+}
+
+/** Builder for PrivateCircuitPublicInputs (call-level) with fluent API for adding side effects. */
+export class PrivateCircuitPublicInputsBuilder {
+  private noteHashes: NoteHash[] = [];
+  private nullifiers: Nullifier[] = [];
+  private noteHashReadRequests: ScopedReadRequest[] = [];
+  private nullifierReadRequests: ScopedReadRequest[] = [];
+  private keyValidationRequests: KeyValidationRequestAndSeparator[] = [];
+  private privateLogs: PrivateLogData[] = [];
+  private nextCounter: number;
+
+  constructor(
+    private contractAddress: AztecAddress = DEFAULT_CONTRACT_ADDRESS,
+    startCounter = 1,
+  ) {
+    this.nextCounter = startCounter;
+  }
+
+  private getCounter(sideEffectCounter?: number): number {
+    if (sideEffectCounter !== undefined) {
+      this.nextCounter = sideEffectCounter + 1;
+      return sideEffectCounter;
+    }
+    return this.nextCounter++;
+  }
+
+  /** Adds a note hash. Defaults are generated randomly. */
+  addNoteHash(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.noteHashes.push(new NoteHash(value, counter));
+    return this;
+  }
+
+  /** Adds a nullifier. Defaults are generated randomly. */
+  addNullifier(opts?: { value?: Fr; noteHash?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const noteHash = opts?.noteHash ?? Fr.ZERO;
+    const counter = this.getCounter(opts?.counter);
+    this.nullifiers.push(new Nullifier(value, noteHash, counter));
+    return this;
+  }
+
+  /** Adds a pending note hash read request (non-empty contract address, can match a pending note hash). */
+  addPendingNoteHashReadRequest(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.noteHashReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), this.contractAddress));
+    return this;
+  }
+
+  /** Adds a settled note hash read request (empty contract address, resolved against the note hash tree). */
+  addSettledNoteHashReadRequest(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.noteHashReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), AztecAddress.ZERO));
+    return this;
+  }
+
+  /** Adds a pending nullifier read request (non-empty contract address, can match a pending nullifier). */
+  addPendingNullifierReadRequest(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.nullifierReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), this.contractAddress));
+    return this;
+  }
+
+  /** Adds a settled nullifier read request (empty contract address, resolved against the nullifier tree). */
+  addSettledNullifierReadRequest(opts?: { value?: Fr; counter?: number }): this {
+    const value = opts?.value ?? Fr.random();
+    const counter = this.getCounter(opts?.counter);
+    this.nullifierReadRequests.push(new ScopedReadRequest(new ReadRequest(value, counter), AztecAddress.ZERO));
+    return this;
+  }
+
+  /** Adds a key validation request. */
+  addKeyValidationRequest(): this {
+    this.keyValidationRequests.push(
+      new KeyValidationRequestAndSeparator(
+        new KeyValidationRequest(new Point(Fr.random(), Fr.random(), false), Fr.random()),
+        Fr.random(),
+      ),
+    );
+    return this;
+  }
+
+  /** Adds a private log. Defaults are generated randomly. */
+  addPrivateLog(opts?: { noteHashCounter?: number; counter?: number }): this {
+    const noteHashCounter = opts?.noteHashCounter ?? 0;
+    const counter = this.getCounter(opts?.counter);
+    this.privateLogs.push(new PrivateLogData(PrivateLog.empty(), noteHashCounter, counter));
+    return this;
+  }
+
+  /** Builds the PrivateCircuitPublicInputs with all added side effects. */
+  build(): PrivateCircuitPublicInputs {
+    const publicInputs = PrivateCircuitPublicInputs.empty();
+    publicInputs.callContext.contractAddress = this.contractAddress;
+    publicInputs.noteHashes = makeClaimed(this.noteHashes, NoteHash, MAX_NOTE_HASHES_PER_CALL);
+    publicInputs.nullifiers = makeClaimed(this.nullifiers, Nullifier, MAX_NULLIFIERS_PER_CALL);
+    publicInputs.privateLogs = makeClaimed(this.privateLogs, PrivateLogData, MAX_PRIVATE_LOGS_PER_CALL);
+    publicInputs.noteHashReadRequests = makeClaimed(
+      this.noteHashReadRequests,
+      ScopedReadRequest,
+      MAX_NOTE_HASH_READ_REQUESTS_PER_CALL,
+    );
+    publicInputs.nullifierReadRequests = makeClaimed(
+      this.nullifierReadRequests,
+      ScopedReadRequest,
+      MAX_NULLIFIER_READ_REQUESTS_PER_CALL,
+    );
+    publicInputs.keyValidationRequestsAndSeparators = makeClaimed(
+      this.keyValidationRequests,
+      KeyValidationRequestAndSeparator,
+      MAX_KEY_VALIDATION_REQUESTS_PER_CALL,
+    );
+    return publicInputs;
+  }
+}
+
+/** Wraps a PrivateKernelCircuitPublicInputs in a PrivateKernelSimulateOutput. */
+export function makeKernelOutput(
+  publicInputs?: PrivateKernelCircuitPublicInputs,
+): PrivateKernelSimulateOutput<PrivateKernelCircuitPublicInputs> {
+  return {
+    publicInputs: publicInputs ?? PrivateKernelCircuitPublicInputs.empty(),
+    verificationKey: VerificationKeyData.empty(),
+    outputWitness: new Map(),
+    bytecode: Buffer.from([]),
+  };
+}
+
+/** Wraps a PrivateCircuitPublicInputs in a PrivateCallExecutionResult. */
+export function makeExecutionResult(publicInputs?: PrivateCircuitPublicInputs): PrivateCallExecutionResult {
+  return new PrivateCallExecutionResult(
+    Buffer.alloc(0),
+    Buffer.alloc(0),
+    new Map(),
+    publicInputs ?? PrivateCircuitPublicInputs.empty(),
+    [],
+    new Map(),
+    [],
+    [],
+    [],
+    [],
+    [],
+  );
+}

package/src/private_kernel/private_kernel_execution_prover.ts

@@ -1,6 +1,7 @@
+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';
@@ -22,17 +23,20 @@ import {
   PrivateKernelTailCircuitPrivateInputs,
   type PrivateKernelTailCircuitPublicInputs,
   PrivateVerificationKeyHints,
+  type UpdatedClassIdHints,
 } from '@aztec/stdlib/kernel';
 import { ChonkProof, ChonkProofWithPublicInputs } from '@aztec/stdlib/proofs';
 import {
   type PrivateCallExecutionResult,
   type PrivateExecutionResult,
   TxRequest,
+  collectNested,
   collectNoteHashNullifierCounterMap,
   getFinalMinRevertibleSideEffectCounter,
 } from '@aztec/stdlib/tx';
 import { VerificationKeyAsFields, VerificationKeyData, VkData } from '@aztec/stdlib/vks';
 
+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';
 
@@ -56,13 +60,16 @@ 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,
+  ) {
+    this.log = createLogger('pxe:private-kernel-execution-prover', bindings);
+  }
 
   /**
    * Generate a proof for a given transaction request and execution result.
@@ -104,6 +111,8 @@ export class PrivateKernelExecutionProver {
     const minRevertibleSideEffectCounter = getFinalMinRevertibleSideEffectCounter(executionResult);
     const splitCounter = isPrivateOnlyTx ? 0 : minRevertibleSideEffectCounter;
 
+    const updatedClassIdHintsMap = await this.prefetchUpdatedClassIdHints(executionResult);
+
     while (executionStack.length) {
       if (!firstIteration) {
         let resetBuilder = new PrivateKernelResetPrivateInputsBuilder(
@@ -113,6 +122,7 @@ export class PrivateKernelExecutionProver {
           splitCounter,
         );
         while (resetBuilder.needsReset()) {
+          // Inner reset: without siloing.
           const witgenTimer = new Timer();
           const privateInputs = await resetBuilder.build(this.oracle);
           output = generateWitnesses
@@ -156,7 +166,7 @@ export class PrivateKernelExecutionProver {
         },
       });
 
-      const privateCallData = await this.createPrivateCallData(currentExecution);
+      const privateCallData = await this.createPrivateCallData(currentExecution, updatedClassIdHintsMap);
 
       if (firstIteration) {
         const witgenTimer = new Timer();
@@ -213,16 +223,24 @@ export class PrivateKernelExecutionProver {
       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);
+      const privateInputs = await finalResetBuilder.build(this.oracle);
       output = generateWitnesses
         ? await this.proofCreator.generateResetOutput(privateInputs)
         : await this.proofCreator.simulateReset(privateInputs);
@@ -236,8 +254,6 @@ export class PrivateKernelExecutionProver {
           witgen: witgenTimer.ms(),
         },
       });
-
-      resetBuilder = new PrivateKernelResetPrivateInputsBuilder(output, [], noteHashNullifierCounterMap, splitCounter);
     }
 
     if (output.publicInputs.feePayer.isZero() && skipFeeEnforcement) {
@@ -257,20 +273,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();
@@ -396,7 +406,28 @@ export class PrivateKernelExecutionProver {
     );
   }
 
-  private async createPrivateCallData({ publicInputs, vk: vkAsBuffer }: PrivateCallExecutionResult) {
+  /** 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);
@@ -412,7 +443,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,

package/src/private_kernel/private_kernel_oracle.ts

@@ -1,4 +1,10 @@
-import { FUNCTION_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT, PUBLIC_DATA_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 { MembershipWitness } from '@aztec/foundation/trees';
@@ -7,17 +13,13 @@ import { getVKIndex, getVKSiblingPath } from '@aztec/noir-protocol-circuits-type
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
 import type { FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { L2BlockHash } from '@aztec/stdlib/block';
-import {
-  type ContractInstanceWithAddress,
-  computeContractClassIdPreimage,
-  computeSaltedInitializationHash,
-} from '@aztec/stdlib/contract';
+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 { 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';
@@ -30,7 +32,7 @@ export class PrivateKernelOracle {
     private contractStore: ContractStore,
     private keyStore: KeyStore,
     private node: AztecNode,
-    private blockHash: L2BlockHash,
+    private blockHeader: BlockHeader,
   ) {}
 
   /** Retrieves the preimage of a contract address from the registered contract instances db. */
@@ -49,11 +51,15 @@ export class PrivateKernelOracle {
 
   /** Retrieves the preimage of a contract class id from the contract classes db. */
   public async getContractClassIdPreimage(contractClassId: Fr) {
-    const contractClass = await this.contractStore.getContractClass(contractClassId);
+    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 computeContractClassIdPreimage(contractClass);
+    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. */
@@ -80,22 +86,20 @@ export class PrivateKernelOracle {
   }
 
   /** Returns a membership witness with the sibling path and leaf index in our note hash tree. */
-  getNoteHashMembershipWitness(noteHash: Fr): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined> {
-    return this.node.getNoteHashMembershipWitness(this.blockHash, noteHash);
+  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> {
-    return this.node.getNullifierMembershipWitness(this.blockHash, nullifier);
+  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. */
-  async getNoteHashTreeRoot(): Promise<Fr> {
-    const header = await this.node.getBlockHeader(this.blockHash);
-    if (!header) {
-      throw new Error(`No block header found for block hash ${this.blockHash}`);
-    }
-    return header.state.partial.noteHashTree.root;
+  getNoteHashTreeRoot(): Fr {
+    return this.blockHeader.state.partial.noteHashTree.root;
   }
 
   /**
@@ -126,18 +130,24 @@ export class PrivateKernelOracle {
       ProtocolContractAddress.ContractInstanceRegistry,
       delayedPublicMutableHashSlot,
     );
-    const updatedClassIdWitness = await this.node.getPublicDataWitness(this.blockHash, hashLeafSlot);
+    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(this.blockHash, ProtocolContractAddress.ContractInstanceRegistry, storageSlot);
-    const delayedPublicMutableValues = await DelayedPublicMutableValues.readFromTree(
-      delayedPublicMutableSlot,
-      readStorage,
-    );
+      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);
 
     return new UpdatedClassIdHints(
       new MembershipWitness(