feat!: introduce WithHash<T> + use it in PublicImmutable (#8022)

docs/docs/migration_notes.md

@@ -6,6 +6,27 @@ keywords: [sandbox, aztec, notes, migration, updating, upgrading]
 
 Aztec is in full-speed development. Literally every version breaks compatibility with the previous ones. This page attempts to target errors and difficulties you might encounter when upgrading, and how to resolve them.
 
+### TBD
+
+### [Aztec.nr] Introduction of `WithHash<T>`
+`WithHash<T>` is a struct that allows for efficient reading of value `T` from public storage in private.
+This is achieved by storing the value with its hash, then obtaining the values via an oracle and verifying them against the hash.
+This results in in a fewer tree inclusion proofs for values `T` that are packed into more than a single field.
+
+`WithHash<T>` is leveraged by state variables like `PublicImmutable`.
+This is a breaking change because now we require values stored in `PublicImmutable` and `SharedMutable` to implement the `Eq` trait.
+
+To implement the `Eq` trait you can use the `#[derive(Eq)]` macro:
+
+```diff
++ use std::meta::derive;
+
++ #[derive(Eq)]
+pub struct YourType {
+    ...
+}
+```
+
 ## 0.73.0
 
 ### [Token, FPC] Moving fee-related complexity from the Token to the FPC

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

@@ -1,22 +1,31 @@
 use crate::{
     context::{PrivateContext, PublicContext, UnconstrainedContext},
-    history::public_storage::PublicStorageHistoricalRead,
     state_vars::storage::Storage,
+    utils::with_hash::WithHash,
 };
 use dep::protocol_types::{constants::INITIALIZATION_SLOT_SEPARATOR, traits::Packable};
 
 /// Stores an immutable value in public state which can be read from public, private and unconstrained execution
 /// contexts.
+///
+/// Leverages `WithHash<T>` to enable efficient private reads of public storage. `WithHash` wrapper allows for
+/// efficient reads by verifying large values through a single hash check and then proving inclusion only of the hash
+/// in the public storage. This reduces the number of required tree inclusion proofs from O(M) to O(1).
+///
+/// This is valuable when T packs to multiple fields, as it maintains "almost constant" verification overhead
+/// regardless of the original data size.
 // docs:start:public_immutable_struct
 pub struct PublicImmutable<T, Context> {
     context: Context,
     storage_slot: Field,
 }
 // docs:end:public_immutable_struct
 
-impl<T, Context, let N: u32> Storage<N> for PublicImmutable<T, Context>
+/// `WithHash<T>` stores both the packed value (using N fields) and its hash (1 field), requiring N = M + 1 total
+/// fields.
+impl<T, Context, let M: u32, let N: u32> Storage<N> for PublicImmutable<T, Context>
 where
-    T: Packable<N>,
+    WithHash<T, M>: Packable<N>,
 {
     fn get_storage_slot(self) -> Field {
         self.storage_slot
@@ -38,7 +47,7 @@ impl<T, Context> PublicImmutable<T, Context> {
 
 impl<T, let T_PACKED_LEN: u32> PublicImmutable<T, &mut PublicContext>
 where
-    T: Packable<T_PACKED_LEN>,
+    T: Packable<T_PACKED_LEN> + Eq,
 {
     // docs:start:public_immutable_struct_write
     pub fn initialize(self, value: T) {
@@ -49,41 +58,36 @@ where
 
         // We populate the initialization slot with a non-zero value to indicate that the struct is initialized
         self.context.storage_write(initialization_slot, 0xdead);
-        self.context.storage_write(self.storage_slot, value);
+        self.context.storage_write(self.storage_slot, WithHash::new(value));
     }
     // docs:end:public_immutable_struct_write
 
     // Note that we don't access the context, but we do call oracles that are only available in public
     // docs:start:public_immutable_struct_read
     pub fn read(self) -> T {
-        self.context.storage_read(self.storage_slot)
+        WithHash::public_storage_read(*self.context, self.storage_slot)
     }
     // docs:end:public_immutable_struct_read
 }
 
 impl<T, let T_PACKED_LEN: u32> PublicImmutable<T, UnconstrainedContext>
 where
-    T: Packable<T_PACKED_LEN>,
+    T: Packable<T_PACKED_LEN> + Eq,
 {
     pub unconstrained fn read(self) -> T {
-        self.context.storage_read(self.storage_slot)
+        WithHash::unconstrained_public_storage_read(self.context, self.storage_slot)
     }
 }
 
 impl<T, let T_PACKED_LEN: u32> PublicImmutable<T, &mut PrivateContext>
 where
-    T: Packable<T_PACKED_LEN>,
+    T: Packable<T_PACKED_LEN> + Eq,
 {
     pub fn read(self) -> T {
-        let header = self.context.get_block_header();
-        let mut fields = [0; T_PACKED_LEN];
-
-        for i in 0..fields.len() {
-            fields[i] = header.public_storage_historical_read(
-                self.storage_slot + i as Field,
-                (*self.context).this_address(),
-            );
-        }
-        T::unpack(fields)
+        WithHash::historical_public_storage_read(
+            self.context.get_block_header(),
+            self.context.this_address(),
+            self.storage_slot,
+        )
     }
 }

noir-projects/aztec-nr/aztec/src/utils/mod.nr

@@ -5,3 +5,4 @@ pub mod field;
 pub mod point;
 pub mod to_bytes;
 pub mod secrets;
+pub mod with_hash;

noir-projects/aztec-nr/aztec/src/utils/with_hash.nr

@@ -0,0 +1,238 @@
+use crate::{
+    context::{PublicContext, UnconstrainedContext},
+    history::public_storage::PublicStorageHistoricalRead,
+    oracle,
+};
+use dep::protocol_types::{
+    address::AztecAddress, block_header::BlockHeader, hash::poseidon2_hash, traits::Packable,
+};
+
+/// A struct that allows for efficient reading of value `T` from public storage in private.
+///
+/// The efficient reads are achieved by verifying large values through a single hash check
+/// and then proving inclusion only of the hash in public storage. This reduces the number
+/// of required tree inclusion proofs from `N` to 1.
+///
+/// # Type Parameters
+/// - `T`: The underlying type being wrapped, must implement `Packable<N>`
+/// - `N`: The number of field elements required to pack values of type `T`
+pub struct WithHash<T, let N: u32> {
+    value: T,
+    packed: [Field; N],
+    hash: Field,
+}
+
+impl<T, let N: u32> WithHash<T, N>
+where
+    T: Packable<N> + Eq,
+{
+    pub fn new(value: T) -> Self {
+        let packed = value.pack();
+        Self { value, packed, hash: poseidon2_hash(packed) }
+    }
+
+    pub fn get_value(self) -> T {
+        self.value
+    }
+
+    pub fn get_hash(self) -> Field {
+        self.hash
+    }
+
+    pub fn public_storage_read(context: PublicContext, storage_slot: Field) -> T {
+        context.storage_read(storage_slot)
+    }
+
+    pub unconstrained fn unconstrained_public_storage_read(
+        context: UnconstrainedContext,
+        storage_slot: Field,
+    ) -> T {
+        context.storage_read(storage_slot)
+    }
+
+    pub fn historical_public_storage_read(
+        header: BlockHeader,
+        address: AztecAddress,
+        storage_slot: Field,
+    ) -> T {
+        let historical_block_number = header.global_variables.block_number as u32;
+
+        // We could simply produce historical inclusion proofs for each field in `packed`, but that would require one
+        // full sibling path per storage slot (since due to kernel siloing the storage is not contiguous). Instead, we
+        // get an oracle to provide us the values, and instead we prove inclusion of their hash, which is both a much
+        // smaller proof (a single slot), and also independent of the size of T (except in that we need to pack and hash T).
+        let hint = WithHash::new(
+            /// Safety: We verify that a hash of the hint/packed data matches the stored hash.
+            unsafe {
+                oracle::storage::storage_read(address, storage_slot, historical_block_number)
+            },
+        );
+
+        let hash = header.public_storage_historical_read(storage_slot + N as Field, address);
+
+        if hash != 0 {
+            assert_eq(hash, hint.get_hash(), "Hint values do not match hash");
+        } else {
+            // The hash slot can only hold a zero if it is uninitialized. Therefore, the hints must then be zero
+            // (i.e. the default value for public storage) as well.
+            assert_eq(
+                hint.get_value(),
+                T::unpack(std::mem::zeroed()),
+                "Non-zero hint for zero hash",
+            );
+        };
+
+        hint.get_value()
+    }
+}
+
+impl<T, let N: u32> Packable<N + 1> for WithHash<T, N>
+where
+    T: Packable<N>,
+{
+    fn pack(self) -> [Field; N + 1] {
+        let mut result: [Field; N + 1] = std::mem::zeroed();
+        for i in 0..N {
+            result[i] = self.packed[i];
+        }
+        result[N] = self.hash;
+
+        result
+    }
+
+    fn unpack(packed: [Field; N + 1]) -> Self {
+        let mut value_packed: [Field; N] = std::mem::zeroed();
+        for i in 0..N {
+            value_packed[i] = packed[i];
+        }
+        let hash = packed[N];
+
+        Self { value: T::unpack(value_packed), packed: value_packed, hash }
+    }
+}
+
+mod test {
+    use crate::{
+        oracle::random::random,
+        test::{
+            helpers::{cheatcodes, test_environment::TestEnvironment},
+            mocks::mock_struct::MockStruct,
+        },
+        utils::with_hash::WithHash,
+    };
+    use dep::protocol_types::hash::poseidon2_hash;
+    use dep::std::{mem, test::OracleMock};
+
+    global storage_slot: Field = 47;
+
+    #[test]
+    unconstrained fn create_and_recover() {
+        let value = MockStruct { a: 5, b: 3 };
+        let value_with_hash = WithHash::new(value);
+        let recovered = WithHash::unpack(value_with_hash.pack());
+
+        assert_eq(recovered.value, value);
+        assert_eq(recovered.packed, value.pack());
+        assert_eq(recovered.hash, poseidon2_hash(value.pack()));
+    }
+
+    #[test]
+    unconstrained fn read_uninitialized_value() {
+        let mut env = TestEnvironment::new();
+
+        let block_header = env.private().historical_header;
+        let address = env.contract_address();
+
+        let result = WithHash::<MockStruct, _>::historical_public_storage_read(
+            block_header,
+            address,
+            storage_slot,
+        );
+
+        // We should get zeroed value
+        let expected: MockStruct = mem::zeroed();
+        assert_eq(result, expected);
+    }
+
+    #[test]
+    unconstrained fn read_initialized_value() {
+        let mut env = TestEnvironment::new();
+
+        let value = MockStruct { a: 5, b: 3 };
+        let value_with_hash = WithHash::new(value);
+
+        // We write the value with hash to storage
+        cheatcodes::direct_storage_write(
+            env.contract_address(),
+            storage_slot,
+            value_with_hash.pack(),
+        );
+
+        // We advance block by 1 because env.private() currently returns context at latest_block - 1
+        env.advance_block_by(1);
+
+        let result = WithHash::<MockStruct, _>::historical_public_storage_read(
+            env.private().historical_header,
+            env.contract_address(),
+            storage_slot,
+        );
+
+        assert_eq(result, value);
+    }
+
+    #[test(should_fail_with = "Non-zero hint for zero hash")]
+    unconstrained fn test_bad_hint_uninitialized_value() {
+        let mut env = TestEnvironment::new();
+
+        env.advance_block_to(6);
+
+        let value_packed = MockStruct { a: 1, b: 1 }.pack();
+
+        let block_header = env.private().historical_header;
+        let address = env.contract_address();
+
+        // Mock the oracle to return a non-zero hint/packed value
+        let _ = OracleMock::mock("storageRead")
+            .with_params((
+                address.to_field(), storage_slot, block_header.global_variables.block_number as u32,
+                value_packed.len(),
+            ))
+            .returns(value_packed)
+            .times(1);
+
+        // This should revert because the hint value is non-zero and the hash is zero (default value of storage)
+        let _ = WithHash::<MockStruct, _>::historical_public_storage_read(
+            block_header,
+            address,
+            storage_slot,
+        );
+    }
+
+    #[test(should_fail_with = "Hint values do not match hash")]
+    unconstrained fn test_bad_hint_initialized_value() {
+        let mut env = TestEnvironment::new();
+
+        let value_packed = MockStruct { a: 5, b: 3 }.pack();
+
+        // We write the value to storage
+        cheatcodes::direct_storage_write(env.contract_address(), storage_slot, value_packed);
+
+        // Now we write incorrect hash to the hash storage slot
+        let incorrect_hash = random();
+        let hash_storage_slot = storage_slot + (value_packed.len() as Field);
+        cheatcodes::direct_storage_write(
+            env.contract_address(),
+            hash_storage_slot,
+            [incorrect_hash],
+        );
+
+        // We advance block by 1 because env.private() currently returns context at latest_block - 1
+        env.advance_block_by(1);
+
+        let _ = WithHash::<MockStruct, _>::historical_public_storage_read(
+            env.private().historical_header,
+            env.contract_address(),
+            storage_slot,
+        );
+    }
+}

noir-projects/aztec-nr/compressed-string/src/field_compressed_string.nr

@@ -6,7 +6,7 @@ use std::meta::derive;
 
 // A Fixedsize Compressed String.
 // Essentially a special version of Compressed String for practical use.
-#[derive(Deserialize, Packable, Serialize)]
+#[derive(Deserialize, Eq, Packable, Serialize)]
 pub struct FieldCompressedString {
     value: Field,
 }

noir-projects/noir-contracts/contracts/amm_contract/src/config.nr

@@ -4,7 +4,7 @@ use std::meta::derive;
 /// We store the tokens of the pool in a struct such that to load it from SharedImmutable asserts only a single
 /// merkle proof.
 /// (Once we actually do the optimization. WIP in https://github.com/AztecProtocol/aztec-packages/pull/8022).
-#[derive(Deserialize, Packable, Serialize)]
+#[derive(Deserialize, Eq, Packable, Serialize)]
 pub struct Config {
     pub token0: AztecAddress,
     pub token1: AztecAddress,

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

@@ -18,6 +18,7 @@ pub contract AppSubscription {
     use router::utils::privately_check_block_number;
     use token::Token;
 
+    // TODO: This can be optimized by storing the values in Config struct in 1 PublicImmutable (less merkle proofs).
     #[storage]
     struct Storage<Context> {
         target_address: PublicImmutable<AztecAddress, Context>,