feat!: improve storage slot allocation (#10320)

Author: nventuro

docs/docs/migration_notes.md

@@ -8,6 +8,15 @@ Aztec is in full-speed development. Literally every version breaks compatibility
 
 ## TBD
 
+### [Aztec.nr] Improved storage slot allocation
+State variables are no longer assumed to be generic over a type that implements the `Serialize` trait: instead, they must implement the `Storage` trait with an `N` value equal to the number of slots they need to reserve.
+
+For the vast majority of state variables, this simply means binding the serialization length to this trait:
+
+```diff
++ impl<T, let N: u32> Storage<N> for MyStateVar<T> where T: Serialize<N> { };
+```
+
 ### [Aztec.nr] Introduction of `Packable` trait
 We have introduced a `Packable` trait that allows types to be serialized and deserialized with a focus on minimizing the size of the resulting Field array.
 This is in contrast to the `Serialize` and `Deserialize` traits, which follows Noir's intrinsic serialization format.

noir-projects/aztec-nr/aztec/src/macros/functions/mod.nr

@@ -332,7 +332,12 @@ pub comptime fn transform_unconstrained(f: FunctionDefinition) {
     let module_has_storage = module_has_storage(f.module());
 
     let storage_init = if module_has_storage {
-        quote { let storage = Storage::init(context); }
+        quote {
+            // Some functions don't access storage, but it'd be quite difficult to only inject this variable if it is
+            // referenced. We instead ignore 'unused variable' warnings for it.
+            #[allow(unused_variables)]
+            let storage = Storage::init(context);
+        }
     } else {
         quote {}
     };

noir-projects/aztec-nr/aztec/src/macros/storage/mod.nr

@@ -1,6 +1,6 @@
 use std::{collections::umap::UHashMap, hash::{BuildHasherDefault, poseidon2::Poseidon2Hasher}};
 
-use super::utils::get_packed_size;
+use super::utils::get_storage_size;
 use super::utils::is_note;
 
 /// Stores a map from a module to the name of the struct that describes its storage layout.
@@ -30,22 +30,22 @@ pub comptime fn storage(s: StructDefinition) -> Quoted {
     for field in s.fields_as_written() {
         // FIXME: This doesn't handle field types with generics
         let (name, typ) = field;
-        let (storage_field_constructor, serialized_size) =
+        let (storage_field_constructor, storage_size) =
             generate_storage_field_constructor(typ, quote { $slot }, false);
         storage_vars_constructors =
             storage_vars_constructors.push_back(quote { $name: $storage_field_constructor });
         // We have `Storable` in a separate `.nr` file instead of defining it in the last quote of this function
         // because that way a dev gets a more reasonable error if he defines a struct with the same name in
         // a contract.
         storage_layout_fields =
-            storage_layout_fields.push_back(quote { $name: dep::aztec::prelude::Storable });
+            storage_layout_fields.push_back(quote { pub $name: dep::aztec::prelude::Storable });
         storage_layout_constructors = storage_layout_constructors.push_back(
             quote { $name: dep::aztec::prelude::Storable { slot: $slot } },
         );
         //let with_context_generic = add_context_generic(typ, context_generic);
         //println(with_context_generic);
         //new_storage_fields = new_storage_fields.push_back((name,  with_context_generic ));
-        slot += serialized_size;
+        slot += storage_size;
     }
 
     //s.set_fields(new_storage_fields);
@@ -119,28 +119,25 @@ comptime fn generate_storage_field_constructor(
             generate_storage_field_constructor(generics[1], quote { slot }, true);
         (quote { $struct_name::new(context, $slot, | context, slot | { $value_constructor }) }, 1)
     } else {
-        let (container_struct, container_struct_generics) = typ.as_struct().unwrap();
-        let container_struct_name = container_struct.name();
-
-        let serialized_size = if parent_is_map {
+        let storage_size = if parent_is_map {
             // Variables inside a map do not require contiguous slots since the map slot derivation is assumed to result
             // in slots very far away from one another.
             1
         } else {
+            let (_, container_struct_generics) = typ.as_struct().unwrap();
             let stored_struct = container_struct_generics[0];
-            if is_note(stored_struct) & (container_struct_name != quote { PublicMutable }) {
+
+            if is_note(stored_struct) {
                 // Private notes always occupy a single slot, since the slot is only used as a state variable
                 // identifier.
-                // Someone could store a Note in PublicMutable for whatever reason though.
-                // TODO(#8659): remove the PublicMutable exception above
                 1
             } else {
-                get_packed_size(stored_struct)
+                get_storage_size(typ)
             }
         };
 
         // We assume below that all state variables implement `fn new<Context>(context: Context, slot: Field) -> Self`.
-        (quote { $struct_name::new(context, $slot)}, serialized_size)
+        (quote { $struct_name::new(context, $slot)}, storage_size)
     }
 }
 

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

@@ -1,4 +1,4 @@
-use std::meta::{typ::fresh_type_variable, unquote};
+use std::meta::unquote;
 
 pub(crate) comptime fn get_fn_visibility(f: FunctionDefinition) -> Quoted {
     if f.has_named_attribute("private") {
@@ -219,17 +219,18 @@ pub(crate) comptime fn compute_event_selector(s: StructDefinition) -> Field {
     unquote!(computation_quote)
 }
 
-pub(crate) comptime fn get_packed_size(typ: Type) -> u32 {
-    let any = fresh_type_variable();
-    let maybe_packed_impl =
-        typ.get_trait_impl(quote { protocol_types::traits::Packable<$any> }.as_trait_constraint());
+/// Returns how many storage slots a type needs to reserve for itself. State variables must implement the Storage trait
+/// for slots to be allocated for them.
+pub(crate) comptime fn get_storage_size(typ: Type) -> u32 {
+    // We create a type variable for the storage size. We can't simply read the value used in the implementation because
+    // it may not be a constant (e.g. N + 1). We then bind it to the implementation of the Storage trait.
+    let storage_size = std::meta::typ::fresh_type_variable();
+    assert(
+        typ.implements(quote { crate::state_vars::Storage<$storage_size> }.as_trait_constraint()),
+        f"Attempted to fetch storage size, but {typ} does not implement the Storage trait",
+    );
 
-    maybe_packed_impl
-        .expect(f"Attempted to fetch packed length, but {typ} does not implement the Packable trait"
-            )
-            .trait_generic_args()[0]
-        .as_constant()
-        .unwrap()
+    storage_size.as_constant().unwrap()
 }
 
 pub(crate) comptime fn module_has_storage(m: Module) -> bool {

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

@@ -9,7 +9,7 @@ pub struct Map<K, V, Context> {
 }
 // docs:end:map
 
-impl<K, T, Context, let N: u32> Storage<T, N> for Map<K, T, Context>
+impl<K, T, Context, let N: u32> Storage<N> for Map<K, T, Context>
 where
     T: Packable<N>,
 {

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

@@ -21,7 +21,7 @@ pub struct PrivateImmutable<Note, Context> {
 }
 // docs:end:struct
 
-impl<T, Context, let N: u32> Storage<T, N> for PrivateImmutable<T, Context>
+impl<T, Context, let N: u32> Storage<N> for PrivateImmutable<T, Context>
 where
     T: Packable<N>,
 {

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

@@ -23,7 +23,7 @@ pub struct PrivateMutable<Note, Context> {
 
 mod test;
 
-impl<T, Context, let N: u32> Storage<T, N> for PrivateMutable<T, Context>
+impl<T, Context, let N: u32> Storage<N> for PrivateMutable<T, Context>
 where
     T: Packable<N>,
 {

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

@@ -22,7 +22,7 @@ pub struct PrivateSet<Note, Context> {
 }
 // docs:end:struct
 
-impl<T, Context, let N: u32> Storage<T, N> for PrivateSet<T, Context>
+impl<T, Context, let N: u32> Storage<N> for PrivateSet<T, Context>
 where
     T: Packable<N>,
 {

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

@@ -14,7 +14,7 @@ pub struct PublicImmutable<T, Context> {
 }
 // docs:end:public_immutable_struct
 
-impl<T, Context, let N: u32> Storage<T, N> for PublicImmutable<T, Context>
+impl<T, Context, let N: u32> Storage<N> for PublicImmutable<T, Context>
 where
     T: Packable<N>,
 {

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

@@ -9,7 +9,7 @@ pub struct PublicMutable<T, Context> {
 }
 // docs:end:public_mutable_struct
 
-impl<T, Context, let N: u32> Storage<T, N> for PublicMutable<T, Context>
+impl<T, Context, let N: u32> Storage<N> for PublicMutable<T, Context>
 where
     T: Packable<N>,
 {

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

@@ -36,7 +36,7 @@ global HASH_SEPARATOR: u32 = 2;
 //
 // TODO https://github.com/AztecProtocol/aztec-packages/issues/5736: change the storage allocation scheme so that we
 // can actually use it here
-impl<T, let INITIAL_DELAY: u32, Context, let N: u32> Storage<T, N> for SharedMutable<T, INITIAL_DELAY, Context>
+impl<T, let INITIAL_DELAY: u32, Context, let N: u32> Storage<N> for SharedMutable<T, INITIAL_DELAY, Context>
 where
     T: Packable<N>,
 {

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

@@ -1,9 +1,6 @@
-use dep::protocol_types::traits::Packable;
-
-pub trait Storage<T, let N: u32>
-where
-    T: Packable<N>,
-{
+/// State variables must implement this trait in order to placed in the storage struct (i.e. the one marked with the
+/// #[storage] attribute). The `N` value determines how many storage slots will be reserved for the state variable.
+pub trait Storage<let N: u32> {
     fn get_storage_slot(self) -> Field;
 }
 

noir-projects/aztec-nr/aztec/src/test/mocks/mock_struct.nr

@@ -1,8 +1,8 @@
 use dep::protocol_types::traits::{Deserialize, Packable, Serialize};
 
-pub(crate) struct MockStruct {
-    pub(crate) a: Field,
-    pub(crate) b: Field,
+pub struct MockStruct {
+    pub a: Field,
+    pub b: Field,
 }
 
 impl MockStruct {

noir-projects/aztec-nr/aztec/src/test/mocks/mod.nr

@@ -1,2 +1,2 @@
-pub(crate) mod mock_note;
-pub(crate) mod mock_struct;
+mod mock_note;
+mod mock_struct;

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

@@ -13,14 +13,14 @@ contract Test {
         event::encode_and_encrypt_event_unconstrained, note::encode_and_encrypt_note,
     };
     use dep::aztec::prelude::{
-        AztecAddress, EthAddress, FunctionSelector, NoteGetterOptions, NoteViewerOptions,
+        AztecAddress, EthAddress, FunctionSelector, Map, NoteGetterOptions, NoteViewerOptions,
         PrivateImmutable, PrivateSet,
     };
 
     use dep::aztec::protocol_types::{
         constants::{MAX_NOTE_HASH_READ_REQUESTS_PER_CALL, PRIVATE_LOG_SIZE_IN_FIELDS},
         point::Point,
-        traits::Serialize,
+        traits::{Hash, Packable, Serialize},
         utils::arrays::array_concat,
     };
 
@@ -61,10 +61,24 @@ contract Test {
         value4: Field,
     }
 
+    #[derive(Packable)]
+    struct ExampleStruct {
+        value0: Field,
+        value1: Field,
+        value2: Field,
+        value3: Field,
+        value4: Field,
+    }
+
+    // This struct is used to test the storage slot allocation mechanism - if modified the test_storage_slot_allocation
+    // test function must also be updated accordingly.
     #[storage]
     struct Storage<Context> {
         example_constant: PrivateImmutable<TestNote, Context>,
         example_set: PrivateSet<TestNote, Context>,
+        example_struct: PrivateImmutable<ExampleStruct, Context>,
+        example_struct_in_map: Map<AztecAddress, PrivateImmutable<ExampleStruct, Context>, Context>,
+        another_example_struct: PrivateImmutable<ExampleStruct, Context>,
     }
 
     #[private]
@@ -482,6 +496,7 @@ contract Test {
         assert_eq(get_public_keys(address).npk_m.inner, public_nullifying_key);
     }
 
+    #[derive(Serialize)]
     pub struct DummyNote {
         amount: Field,
         secret_hash: Field,
@@ -497,12 +512,6 @@ contract Test {
         }
     }
 
-    impl Serialize<2> for DummyNote {
-        fn serialize(self) -> [Field; 2] {
-            [self.amount, self.secret_hash]
-        }
-    }
-
     pub struct DeepStruct {
         a_field: Field,
         a_bool: bool,

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

@@ -1,3 +1,4 @@
+use crate::Test;
 use crate::test_note::TestNote;
 use dep::uint_note::uint_note::UintNote;
 use dep::value_note::value_note::ValueNote;
@@ -10,3 +11,47 @@ unconstrained fn test_note_type_id() {
     assert_eq(ValueNote::get_note_type_id(), 1, "ValueNote type id should be 1");
     assert_eq(TestNote::get_note_type_id(), 2, "TestNote type id should be 2");
 }
+
+#[test]
+unconstrained fn test_storage_slot_allocation() {
+    // This tests that sufficient storage slots are assigned to each state variable so that they do not interfere with
+    // one another. The space a state variable needs is determined by the N value in its implementation of the Storage
+    // trait. Most state variables bind N to the packed length of the type they hold.
+    //
+    // This is the storage declaration:
+    //
+    // #[storage]
+    // struct Storage<Context> {
+    //     example_constant: PrivateImmutable<TestNote, Context>,
+    //     example_set: PrivateSet<TestNote, Context>,
+    //     example_struct: PrivateImmutable<ExampleStruct, Context>,
+    //     example_struct_in_map: Map<AztecAddress, PrivateImmutable<ExampleStruct, Context>, Context>,
+    //     another_example_struct: PrivateImmutable<ExampleStruct, Context>,
+    // }
+
+    // We can't directly see how many slots are allocated to each variable, but we can look at the slot increments for
+    // each and deduct the allocation size based off of that. In other words, given a struct with two members a and b,
+    // the number of slots allocated to a will be b.storage_slot - a.storage_slot.
+
+    // The first slot is always 1.
+    let mut expected_slot = 1;
+    assert_eq(Test::storage_layout().example_constant.slot, expected_slot);
+
+    // Even though example_constant holds TestNote, which packs to a length larger than 1, notes always reserve a
+    // single slot.
+    expected_slot += 1;
+    assert_eq(Test::storage_layout().example_set.slot, expected_slot);
+
+    // example_set also held a note, so it should have only allocated a single slot.
+    expected_slot += 1;
+    assert_eq(Test::storage_layout().example_struct.slot, expected_slot);
+
+    // example_struct allocates 5 slots because it is not a note and it's packed length is 5.
+    expected_slot += 5;
+    assert_eq(Test::storage_layout().example_struct_in_map.slot, expected_slot);
+
+    // example_struct_in_map should allocate a single note because it's a map, regardless of whatever it holds. The Map
+    // type is going to deal with its own dynamic allocation based on keys
+    expected_slot += 1;
+    assert_eq(Test::storage_layout().another_example_struct.slot, expected_slot);
+}

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

@@ -1,6 +1,5 @@
 use crate::test::utils;
 use crate::Token;
-use dep::authwit::cheatcodes as authwit_cheatcodes;
 use dep::aztec::test::helpers::cheatcodes;
 
 #[test]