feat!: extend storage read oracle to receive address and block number (#7243)

Follow up to #7237, closes
#7230. I only
changed the oracles and not the PXE interface to keep this change as
small as possible.

I did change the node interface, but made it so you can still do it the
old way by passing `'latest'`, which I had to do in a couple places.

Finally, I added getters for `UnconstrainedContext`, mirroring the work
in #7320, which I imagine are the ones we'll use in the vast majority of
cases.

Author: nventuro

noir-projects/aztec-nr/aztec/src/context/unconstrained_context.nr

@@ -1,4 +1,5 @@
-use dep::protocol_types::address::AztecAddress;
+use dep::protocol_types::{address::AztecAddress, traits::Deserialize};
+use crate::oracle::storage::{raw_storage_read, storage_read};
 
 struct UnconstrainedContext {
     block_number: u32,
@@ -35,6 +36,14 @@ impl UnconstrainedContext {
     fn chain_id(self) -> Field {
         self.chain_id
     }
+
+    unconstrained fn raw_storage_read<N>(self: Self, storage_slot: Field) -> [Field; N] {
+        storage_read(self.this_address(), storage_slot, self.block_number())
+    }
+
+    unconstrained fn storage_read<T, N>(self, storage_slot: Field) -> T where T: Deserialize<N> {
+        T::deserialize(self.raw_storage_read(storage_slot))
+    }
 }
 
 #[oracle(getContractAddress)]

noir-projects/aztec-nr/aztec/src/oracle/storage.nr

@@ -1,42 +1,58 @@
-use dep::protocol_types::traits::Deserialize;
+use dep::protocol_types::{address::AztecAddress, traits::Deserialize};
 
 #[oracle(storageRead)]
-unconstrained fn storage_read_oracle<N>(storage_slot: Field, length: Field) -> [Field; N] {}
-
-unconstrained pub fn raw_storage_read<N>(storage_slot: Field) -> [Field; N] {
-    storage_read_oracle(storage_slot, N)
+unconstrained fn storage_read_oracle<N>(
+    address: Field,
+    storage_slot: Field,
+    block_number: Field,
+    length: Field
+) -> [Field; N] {}
+
+unconstrained pub fn raw_storage_read<N>(
+    address: AztecAddress,
+    storage_slot: Field,
+    block_number: u32
+) -> [Field; N] {
+    storage_read_oracle(address.to_field(), storage_slot, block_number as Field, N)
 }
 
-unconstrained pub fn storage_read<T, N>(storage_slot: Field) -> T where T: Deserialize<N> {
-    T::deserialize(raw_storage_read(storage_slot))
+unconstrained pub fn storage_read<T, N>(
+    address: AztecAddress,
+    storage_slot: Field,
+    block_number: u32
+) -> T where T: Deserialize<N> {
+    T::deserialize(raw_storage_read(address, storage_slot, block_number))
 }
 
 mod tests {
     use crate::oracle::storage::{raw_storage_read, storage_read};
+    use dep::protocol_types::address::AztecAddress;
 
     use std::test::OracleMock;
     use crate::test::mocks::mock_struct::MockStruct;
 
+    global address = AztecAddress::from_field(29);
+    global slot = 7;
+    global block_number = 17;
+
     #[test]
     fn test_raw_storage_read() {
-        let slot = 7;
         let written = MockStruct { a: 13, b: 42 };
 
-        let _ = OracleMock::mock("storageRead").with_params((slot, 2)).returns(written.serialize());
+        let _ = OracleMock::mock("storageRead").returns(written.serialize());
 
-        let read: [Field; 2] = raw_storage_read(slot);
+        let read: [Field; 2] = raw_storage_read(address, slot, block_number);
         assert_eq(read[0], 13);
         assert_eq(read[1], 42);
     }
 
     #[test]
     fn test_storage_read() {
-        let slot = 7;
         let written = MockStruct { a: 13, b: 42 };
 
-        let _ = OracleMock::mock("storageRead").with_params((slot, 2)).returns(written.serialize());
+        let _ = OracleMock::mock("storageRead").returns(written.serialize());
 
-        let read: MockStruct = storage_read(slot);
+        let read: MockStruct = storage_read(address, slot, block_number);
         assert_eq(read.a, 13);
         assert_eq(read.b, 42);
     }

noir-projects/aztec-nr/aztec/src/state_vars/public_immutable.nr

@@ -56,6 +56,6 @@ impl <T> PublicImmutable<T, &mut PublicContext> {
 
 impl<T> PublicImmutable<T, UnconstrainedContext> {
     unconstrained pub fn read<T_SERIALIZED_LEN>(self) -> T where T: Deserialize<T_SERIALIZED_LEN> {
-        storage_read(self.storage_slot)
+        self.context.storage_read(self.storage_slot)
     }
 }

noir-projects/aztec-nr/aztec/src/state_vars/public_mutable.nr

@@ -41,6 +41,6 @@ impl<T> PublicMutable<T, &mut PublicContext> {
 
 impl<T> PublicMutable<T, UnconstrainedContext> {
     unconstrained pub fn read<T_SERIALIZED_LEN>(self) -> T where T: Deserialize<T_SERIALIZED_LEN> {
-        storage_read(self.storage_slot)
+        self.context.storage_read(self.storage_slot)
     }
 }

noir-projects/aztec-nr/aztec/src/state_vars/shared_immutable.nr

@@ -48,7 +48,7 @@ impl<T> SharedImmutable<T, &mut PublicContext> {
 
 impl<T> SharedImmutable<T, UnconstrainedContext> {
     unconstrained pub fn read_public<T_SERIALIZED_LEN>(self) -> T where T: Deserialize<T_SERIALIZED_LEN> {
-        storage_read(self.storage_slot)
+        self.context.storage_read(self.storage_slot)
     }
 }
 

noir-projects/noir-contracts/contracts/token_contract/src/test/utils.nr

@@ -72,10 +72,11 @@ pub fn setup_and_mint(with_account_contracts: bool) -> (&mut TestEnvironment, Az
 pub fn check_public_balance(token_contract_address: AztecAddress, address: AztecAddress, address_amount: Field) {
     let current_contract_address = cheatcodes::get_contract_address();
     cheatcodes::set_contract_address(token_contract_address);
+    let block_number = cheatcodes::get_block_number();
 
     let balances_slot = Token::storage().public_balances.slot;
     let address_slot = derive_storage_slot_in_map(balances_slot, address);
-    let amount: U128 = storage_read(address_slot);
+    let amount: U128 = storage_read(token_contract_address, address_slot, block_number);
     assert(amount.to_field() == address_amount, "Public balance is not correct");
     cheatcodes::set_contract_address(current_contract_address);
 }

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -704,10 +704,11 @@ export class AztecNodeService implements AztecNode {
    *
    * @param contract - Address of the contract to query.
    * @param slot - Slot to query.
+   * @param blockNumber - The block number at which to get the data or 'latest'.
    * @returns Storage value at the given contract slot.
    */
-  public async getPublicStorageAt(contract: AztecAddress, slot: Fr): Promise<Fr> {
-    const committedDb = await this.#getWorldState('latest');
+  public async getPublicStorageAt(contract: AztecAddress, slot: Fr, blockNumber: L2BlockNumber): Promise<Fr> {
+    const committedDb = await this.#getWorldState(blockNumber);
     const leafSlot = computePublicDataTreeLeafSlot(contract, slot);
 
     const lowLeafResult = await committedDb.getPreviousValueIndex(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot.toBigInt());

yarn-project/circuit-types/src/interfaces/aztec-node.ts

@@ -291,9 +291,10 @@ export interface AztecNode {
    *
    * @param contract - Address of the contract to query.
    * @param slot - Slot to query.
+   * @param blockNumber - The block number at which to get the data or 'latest'.
    * @returns Storage value at the given contract slot.
    */
-  getPublicStorageAt(contract: AztecAddress, slot: Fr): Promise<Fr>;
+  getPublicStorageAt(contract: AztecAddress, slot: Fr, blockNumber: L2BlockNumber): Promise<Fr>;
 
   /**
    * Returns the currently committed block header.

yarn-project/end-to-end/src/benchmarks/bench_process_history.test.ts

@@ -54,7 +54,7 @@ describe('benchmarks/process_history', () => {
           const node = await AztecNodeService.createAndSync(nodeConfig);
           // call getPublicStorageAt (which calls #getWorldState, which calls #syncWorldState) to force a sync with
           // world state to ensure the node has caught up
-          await node.getPublicStorageAt(AztecAddress.random(), Fr.random());
+          await node.getPublicStorageAt(AztecAddress.random(), Fr.random(), 'latest');
           return node;
         });
 

yarn-project/end-to-end/src/benchmarks/bench_publish_rollup.test.ts

@@ -37,7 +37,7 @@ describe('benchmarks/publish_rollup', () => {
       // world state to ensure the node has caught up
       context.logger.info(`Starting new aztec node`);
       const node = await AztecNodeService.createAndSync({ ...context.config, disableSequencer: true });
-      await node.getPublicStorageAt(AztecAddress.random(), Fr.random());
+      await node.getPublicStorageAt(AztecAddress.random(), Fr.random(), 'latest');
 
       // Spin up a new pxe and sync it, we'll use it to test sync times of new accounts for the last block
       context.logger.info(`Starting new pxe`);