refactor(pxe): skip storage reads for never-updated contracts

Author: Nico Chamo

yarn-project/pxe/src/private_kernel/private_kernel_oracle.test.ts

@@ -0,0 +1,73 @@
+import { PUBLIC_DATA_TREE_HEIGHT } from '@aztec/constants';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { SiblingPath } from '@aztec/foundation/trees';
+import type { KeyStore } from '@aztec/key-store';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { DelayedPublicMutableValuesWithHash } from '@aztec/stdlib/delayed-public-mutable';
+import { computePublicDataTreeLeafSlot } from '@aztec/stdlib/hash';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import { PublicDataTreeLeaf, PublicDataTreeLeafPreimage, PublicDataWitness } from '@aztec/stdlib/trees';
+import { BlockHeader } from '@aztec/stdlib/tx';
+
+import { mock } from 'jest-mock-extended';
+
+import type { ContractStore } from '../storage/contract_store/contract_store.js';
+import { PrivateKernelOracle } from './private_kernel_oracle.js';
+
+describe('PrivateKernelOracle', () => {
+  let oracle: PrivateKernelOracle;
+  let node: ReturnType<typeof mock<AztecNode>>;
+
+  beforeEach(() => {
+    node = mock<AztecNode>();
+    oracle = new PrivateKernelOracle(mock<ContractStore>(), mock<KeyStore>(), node, BlockHeader.empty());
+  });
+
+  describe('getUpdatedClassIdHints', () => {
+    it('skips storage reads when contract class was never updated', async () => {
+      const contractAddress = await AztecAddress.random();
+      const hashLeafSlot = await getHashLeafSlot(contractAddress);
+
+      // Non-matching slot simulates a low-leaf witness (slot was never written)
+      const unrelatedSlot = new Fr(hashLeafSlot.toBigInt() - 1n);
+      node.getPublicDataWitness.mockResolvedValue(makeWitness(unrelatedSlot));
+
+      const result = await oracle.getUpdatedClassIdHints(contractAddress);
+
+      expect(result.updatedClassIdValues.isEmpty()).toBe(true);
+      expect(node.getPublicStorageAt).not.toHaveBeenCalled();
+    });
+
+    it('reads storage when contract class was updated', async () => {
+      const contractAddress = await AztecAddress.random();
+      const hashLeafSlot = await getHashLeafSlot(contractAddress);
+
+      // Matching slot means the contract class was updated
+      node.getPublicDataWitness.mockResolvedValue(makeWitness(hashLeafSlot, Fr.random()));
+      node.getPublicStorageAt.mockResolvedValue(new Fr(42));
+
+      const result = await oracle.getUpdatedClassIdHints(contractAddress);
+
+      expect(result.updatedClassIdValues.isEmpty()).toBe(false);
+      expect(node.getPublicStorageAt).toHaveBeenCalled();
+    });
+  });
+
+  async function getHashLeafSlot(contractAddress: AztecAddress) {
+    const { delayedPublicMutableHashSlot } =
+      await DelayedPublicMutableValuesWithHash.getContractUpdateSlots(contractAddress);
+    return computePublicDataTreeLeafSlot(
+      ProtocolContractAddress.ContractInstanceRegistry,
+      delayedPublicMutableHashSlot,
+    );
+  }
+
+  function makeWitness(slot: Fr, value: Fr = Fr.ZERO) {
+    return new PublicDataWitness(
+      0n,
+      new PublicDataTreeLeafPreimage(new PublicDataTreeLeaf(slot, value), Fr.ZERO, 0n),
+      SiblingPath.random(PUBLIC_DATA_TREE_HEIGHT),
+    );
+  }
+});

yarn-project/pxe/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';
@@ -132,12 +138,16 @@ export class PrivateKernelOracle {
       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 delayedPublicMutableValues = await DelayedPublicMutableValues.readFromTree(
-      delayedPublicMutableSlot,
-      readStorage,
-    );
+    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(