fix: shared mutable private getter fixes

noir-projects/aztec-nr/aztec/src/keys/getters.nr

@@ -63,12 +63,12 @@ fn fetch_key_from_registry(
     let y_coordinate_derived_slot = derive_storage_slot_in_map(y_coordinate_map_slot, address);
 
     let x_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY> = SharedMutablePrivateGetter::new(
-        *context,
+        context,
         AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS),
         x_coordinate_derived_slot
     );
     let y_coordinate_registry: SharedMutablePrivateGetter<Field, DELAY> = SharedMutablePrivateGetter::new(
-        *context,
+        context,
         AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS),
         y_coordinate_derived_slot
     );

noir-projects/aztec-nr/aztec/src/state_vars/shared_mutable/shared_mutable_private_getter.nr

@@ -8,7 +8,7 @@ use crate::state_vars::{
 };
 
 struct SharedMutablePrivateGetter<T, INITIAL_DELAY> {
-    context: PrivateContext,
+    context: &mut PrivateContext,
     // The contract address of the contract we want to read from
     other_contract_address: AztecAddress,
     // The storage slot where the SharedMutable is stored on the other contract
@@ -22,7 +22,7 @@ struct SharedMutablePrivateGetter<T, INITIAL_DELAY> {
 // Currently the Shared Mutable does not support this. We can adapt SharedMutable at a later date
 impl<T, INITIAL_DELAY> SharedMutablePrivateGetter<T, INITIAL_DELAY> {
     pub fn new(
-        context: PrivateContext,
+        context: &mut PrivateContext,
         other_contract_address: AztecAddress,
         storage_slot: Field
     ) -> Self {
@@ -32,13 +32,11 @@ impl<T, INITIAL_DELAY> SharedMutablePrivateGetter<T, INITIAL_DELAY> {
     }
 
     pub fn get_current_value_in_private(self) -> T where T: FromField {
-        let mut context = self.context;
-
-        let (value_change, delay_change, historical_block_number) = self.historical_read_from_public_storage(context);
+        let (value_change, delay_change, historical_block_number) = self.historical_read_from_public_storage(*self.context);
         let effective_minimum_delay = delay_change.get_effective_minimum_delay_at(historical_block_number);
         let block_horizon = value_change.get_block_horizon(historical_block_number, effective_minimum_delay);
 
-        context.set_tx_max_block_number(block_horizon);
+        self.context.set_tx_max_block_number(block_horizon);
         value_change.get_current_at(historical_block_number)
     }
 
@@ -57,7 +55,7 @@ impl<T, INITIAL_DELAY> SharedMutablePrivateGetter<T, INITIAL_DELAY> {
         }
 
         let delay_change_slot = self.get_delay_change_storage_slot();
-        let raw_delay_change_fields = [header.public_storage_historical_read(delay_change_slot, context.this_address())];
+        let raw_delay_change_fields = [header.public_storage_historical_read(delay_change_slot, self.other_contract_address)];
 
         let value_change = ScheduledValueChange::deserialize(raw_value_change_fields);
         let delay_change = ScheduledDelayChange::deserialize(raw_delay_change_fields);

noir-projects/noir-contracts/contracts/auth_contract/src/main.nr

@@ -33,6 +33,7 @@ contract Auth {
 
     // docs:start:public_getter
     #[aztec(public)]
+    #[aztec(view)]
     fn get_authorized() -> AztecAddress {
         // docs:start:shared_mutable_get_current_public
         storage.authorized.get_current_value_in_public()
@@ -41,13 +42,25 @@ contract Auth {
     // docs:end:public_getter
 
     #[aztec(public)]
+    #[aztec(view)]
     fn get_scheduled_authorized() -> AztecAddress {
         // docs:start:shared_mutable_get_scheduled_public
         let (scheduled_value, _block_of_change): (AztecAddress, u32) = storage.authorized.get_scheduled_value_in_public();
         // docs:end:shared_mutable_get_scheduled_public
         scheduled_value
     }
 
+    #[aztec(public)]
+    #[aztec(view)]
+    fn get_authorized_delay() -> pub u32 {
+        storage.authorized.get_current_delay_in_public()
+    }
+
+    #[aztec(public)]
+    fn set_authorized_delay(new_delay: u32) {
+        storage.authorized.schedule_delay_change(new_delay);
+    }
+
     #[aztec(private)]
     fn do_private_authorized_thing() {
         // Reading a value from authorized in private automatically adds an extra validity condition: the base rollup
@@ -58,4 +71,10 @@ contract Auth {
         // docs:end:shared_mutable_get_current_private
         assert_eq(authorized, context.msg_sender(), "caller is not authorized");
     }
+
+    #[aztec(private)]
+    #[aztec(view)]
+    fn get_authorized_in_private() -> AztecAddress {
+        storage.authorized.get_current_value_in_private()
+    }
 }

noir-projects/noir-contracts/contracts/test_contract/src/main.nr

@@ -478,7 +478,7 @@ contract Test {
     ) -> Field where T: FromField, T: ToField {
         // It's a bit wonky because we need to know the delay for get_current_value_in_private to work correctly
         let test: SharedMutablePrivateGetter<T, 5> = SharedMutablePrivateGetter::new(
-            context,
+            &mut context,
             contract_address_to_read,
             storage_slot_of_shared_mutable
         );
@@ -500,7 +500,7 @@ contract Test {
 
         // It's a bit wonky because we need to know the delay for get_current_value_in_private to work correctly
         let registry_private_getter: SharedMutablePrivateGetter<Field, 5> = SharedMutablePrivateGetter::new(
-            context,
+            &mut context,
             AztecAddress::from_field(CANONICAL_KEY_REGISTRY_ADDRESS),
             derived_slot
         );

yarn-project/end-to-end/src/e2e_state_vars.test.ts

@@ -1,15 +1,16 @@
-import { AztecAddress, BatchCall, Fr, type Wallet } from '@aztec/aztec.js';
+import { AztecAddress, BatchCall, Fr, type PXE, type Wallet } from '@aztec/aztec.js';
 import { AuthContract, DocsExampleContract, TestContract } from '@aztec/noir-contracts.js';
 
 import { jest } from '@jest/globals';
 
 import { setup } from './fixtures/utils.js';
 
-const TIMEOUT = 100_000;
+const TIMEOUT = 180_000;
 
 describe('e2e_state_vars', () => {
   jest.setTimeout(TIMEOUT);
 
+  let pxe: PXE;
   let wallet: Wallet;
 
   let teardown: () => Promise<void>;
@@ -19,7 +20,7 @@ describe('e2e_state_vars', () => {
   const RANDOMNESS = 2n;
 
   beforeAll(async () => {
-    ({ teardown, wallet } = await setup(2));
+    ({ teardown, wallet, pxe } = await setup(2));
     contract = await DocsExampleContract.deploy(wallet).send().deployed();
   });
 
@@ -250,31 +251,107 @@ describe('e2e_state_vars', () => {
       // We use the auth contract here because has a nice, clear, simple implementation of the Shared Mutable,
       // and we will need to read from it to test our private getter.
       authContract = await AuthContract.deploy(wallet, wallet.getAddress()).send().deployed();
+    });
 
+    it('checks authorized in AuthContract from TestContract with our SharedMutablePrivateGetter before and after a value change', async () => {
       // We set the authorized value here, knowing there will be some delay before the value change takes place
       await authContract
         .withWallet(wallet)
         .methods.set_authorized(AztecAddress.fromField(new Fr(6969696969)))
         .send()
         .wait();
-    });
 
-    it("checks authorized in auth contract from test contract and finds the old value because the change hasn't been applied yet", async () => {
       const authorized = await testContract.methods
         .test_shared_mutable_private_getter(authContract.address, 2)
         .simulate();
 
+      // We expect the value to not have been applied yet
       expect(AztecAddress.fromBigInt(authorized)).toEqual(AztecAddress.ZERO);
-    });
 
-    it('checks authorized in auth contract from test contract and finds the correctly set value', async () => {
+      // We wait for the SharedMutable delay
       await delay(5);
 
-      const authorized = await testContract.methods
+      // We check after the delay, expecting to find the value we set.
+      const newAuthorized = await testContract.methods
         .test_shared_mutable_private_getter(authContract.address, 2)
         .simulate();
 
-      expect(AztecAddress.fromBigInt(authorized)).toEqual(AztecAddress.fromBigInt(6969696969n));
+      expect(AztecAddress.fromBigInt(newAuthorized)).toEqual(AztecAddress.fromBigInt(6969696969n));
+    });
+
+    it('checks authorized in AuthContract from TestContract and finds the correctly set max block number', async () => {
+      const lastBlockNumber = await pxe.getBlockNumber();
+
+      const tx = await testContract.methods.test_shared_mutable_private_getter(authContract.address, 2).prove();
+
+      const expectedMaxBlockNumber = lastBlockNumber + 5;
+
+      expect(tx.data.forRollup!.rollupValidationRequests.maxBlockNumber.isSome).toEqual(true);
+      expect(tx.data.forRollup!.rollupValidationRequests.maxBlockNumber.value).toEqual(new Fr(expectedMaxBlockNumber));
+    });
+
+    it('checks propagation of max block number in private', async () => {
+      // Our initial max block number will be 5 because that is the value of INITIAL_DELAY
+      const expectedInitialMaxBlockNumber = (await pxe.getBlockNumber()) + 5;
+
+      // Our SharedMutablePrivateGetter here reads from the SharedMutable authorized storage property in AuthContract
+      const tx = await testContract.methods.test_shared_mutable_private_getter(authContract.address, 2).prove();
+
+      // The validity of our SharedMutable read request is limited to 5 blocks, which is our initial delay.
+      expect(tx.data.forRollup!.rollupValidationRequests.maxBlockNumber.isSome).toEqual(true);
+      expect(tx.data.forRollup!.rollupValidationRequests.maxBlockNumber.value).toEqual(
+        new Fr(expectedInitialMaxBlockNumber),
+      );
+
+      // We change the SharedMutable authorized delay here to 2, this means that a change to the "authorized" value can
+      // only be applied 2 blocks after it is initiated, and thus read requests on a historical state without an initiated change is
+      // valid for at least 2 blocks.
+      await authContract.methods.set_authorized_delay(2).send().wait();
+
+      // Note: Because we are decreasing the delay, we must first wait for the full previous delay - 1 (5 -1).
+      await delay(4);
+
+      const expectedModifiedMaxBlockNumber = (await pxe.getBlockNumber()) + 2;
+
+      // We now call our AuthContract to see if the change in max block number has reflected our delay change
+      const tx2 = await authContract.methods.get_authorized_in_private().prove();
+
+      // The validity of our SharedMutable read request should now be limited to 2 blocks, instead of 5
+      expect(tx2.data.forRollup!.rollupValidationRequests.maxBlockNumber.isSome).toEqual(true);
+      expect(tx2.data.forRollup!.rollupValidationRequests.maxBlockNumber.value).toEqual(
+        new Fr(expectedModifiedMaxBlockNumber),
+      );
+
+      // We check for parity between accessing our SharedMutable directly and from our SharedMutablePrivateGetter. The validity assumptions should remain the same.
+      const tx3 = await testContract.methods.test_shared_mutable_private_getter(authContract.address, 2).prove();
+
+      expect(tx3.data.forRollup!.rollupValidationRequests.maxBlockNumber.isSome).toEqual(true);
+      expect(tx3.data.forRollup!.rollupValidationRequests.maxBlockNumber.value).toEqual(
+        new Fr(expectedModifiedMaxBlockNumber),
+      );
+
+      // We now change the SharedMutable authorized delay here to 100, the same assumptions apply here
+      await authContract.methods.set_authorized_delay(100).send().wait();
+
+      // Note: Because we are increasing the delay, we do not have to wait for this change to apply
+      const expectedLengthenedModifiedMaxBlockNumber = (await pxe.getBlockNumber()) + 100;
+
+      // We now call our AuthContract to see if the change in max block number has reflected our delay change
+      const tx4 = await authContract.methods.get_authorized_in_private().prove();
+
+      // The validity of our SharedMutable read request should now be limited to 100 blocks, instead of 2
+      expect(tx4.data.forRollup!.rollupValidationRequests.maxBlockNumber.isSome).toEqual(true);
+      expect(tx4.data.forRollup!.rollupValidationRequests.maxBlockNumber.value).toEqual(
+        new Fr(expectedLengthenedModifiedMaxBlockNumber),
+      );
+
+      // We check for parity between accessing our SharedMutable directly and from our SharedMutablePrivateGetter. The validity assumptions should remain the same.
+      const tx5 = await testContract.methods.test_shared_mutable_private_getter(authContract.address, 2).prove();
+
+      expect(tx5.data.forRollup!.rollupValidationRequests.maxBlockNumber.isSome).toEqual(true);
+      expect(tx5.data.forRollup!.rollupValidationRequests.maxBlockNumber.value).toEqual(
+        new Fr(expectedLengthenedModifiedMaxBlockNumber),
+      );
     });
   });
 });