feat!: storage_layout and #[aztec(storage)] (#5387)

Requires #5386

Closes #5079

Leveraging the new `#[abi(tag)]` export attribute, the
`#[aztec(storage]` and `#[aztec(note)]` decorators are now used by the
aztec macros to generate exportable global structs that include storage
and notes info. These are then interpreted by `noir-compiler` in the ts
codegen and exposed through the contract classes.

```typescript

MyContractClass.storage.my_storage_variable.slot: Fr;
MyContractClass.notes.MyNote.id: Fr;

```

In the process I realized we didn't really need the
`process_def_collector` pass, which has been removed (one less thing to
worry about before moving to metaprogramming!) in favor of using
`#[aztec(note)]` to identify notes that need to get included in the
autogenerated `compute_note_hash_and_nullifier` implementation.

---------

Co-authored-by: esau <152162806+sklppy88@users.noreply.github.com>

docs/docs/developers/contracts/references/storage/main.md

@@ -7,7 +7,7 @@ Smart contracts rely on storage, acting as the persistent memory on the blockcha
 To learn how to define a storage struct, read [this guide](../../writing_contracts/storage/define_storage.md).
 To learn more about storage slots, read [this explainer](../../writing_contracts/storage/storage_slots.md).
 
-You control this storage in Aztec using the `Storage` struct. This struct serves as the housing unit for all your smart contract's state variables - the data it needs to keep track of and maintain.
+You control this storage in Aztec using a struct annotated with `#[aztec(storage)]`. This struct serves as the housing unit for all your smart contract's state variables - the data it needs to keep track of and maintain.
 
 These state variables come in two forms: public and private. Public variables are visible to anyone, and private variables remain hidden within the contract.
 

docs/docs/developers/contracts/writing_contracts/storage/define_storage.md

@@ -16,4 +16,4 @@ struct Storage {
 }
 ```
 
-If you have defined a `Storage` struct following this naming scheme, then it will be made available to you through the reserved `storage` keyword within your contract functions.
+If you have defined a struct and annotated it as `#[aztec(storage)]`, then it will be made available to you through the reserved `storage` keyword within your contract functions.

docs/docs/misc/migration_notes.md

@@ -8,6 +8,49 @@ Aztec is in full-speed development. Literally every version breaks compatibility
 
 ## TBD
 
+### [Aztec.nr] Storage struct annotation
+
+The storage struct now identified by the annotation `#[aztec(storage)]`, instead of having to rely on it being called `Storage`.
+
+```diff
+- struct Storage {
+-    ...
+- }
++ #[aztec(storage)]
++ struct MyStorageStruct {
++    ...
++ }
+```
+
+### [Aztec.js] Storage layout and note info
+
+Storage layout and note information are now exposed in the TS contract artifact
+
+```diff
+- const note = new Note([new Fr(mintAmount), secretHash]);
+- const pendingShieldStorageSlot = new Fr(5n); // storage slot for pending_shields
+- const noteTypeId = new Fr(84114971101151129711410111011678111116101n); // note type id for TransparentNote
+- const extendedNote = new ExtendedNote(
+-   note,
+-   admin.address,
+-   token.address,
+-   pendingShieldStorageSlot,
+-   noteTypeId,
+-   receipt.txHash,
+- );
+- await pxe.addNote(extendedNote);
++ const note = new Note([new Fr(mintAmount), secretHash]);
++ const extendedNote = new ExtendedNote(
++   note,
++   admin.address,
++   token.address,
++   TokenContract.storage.pending_shields.slot,
++   TokenContract.notes.TransparentNote.id,
++   receipt.txHash,
++ );
++ await pxe.addNote(extendedNote);
+```
+
 ### [Aztec.nr] rand oracle is now called unsafe_rand
 `oracle::rand::rand` has been renamed to `oracle::unsafe_rand::unsafe_rand`.
 This change was made to communicate that we do not constrain the value in circuit and instead we just trust our PXE.

noir-projects/aztec-nr/aztec/src/prelude.nr

@@ -4,7 +4,7 @@ use crate::{
     state_vars::{
     map::Map, private_immutable::PrivateImmutable, private_mutable::PrivateMutable,
     public_immutable::PublicImmutable, public_mutable::PublicMutable, private_set::PrivateSet,
-    shared_immutable::SharedImmutable
+    shared_immutable::SharedImmutable, storage::Storable
 },
     log::{emit_unencrypted_log, emit_encrypted_log}, context::PrivateContext,
     note::{

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

@@ -5,3 +5,12 @@ trait Storage<T> where T: Serialize<N> + Deserialize<N> {
         self.storage_slot
     }
 }
+
+// Struct representing an exportable storage variable in the contract
+// Every entry in the storage struct will be exported in the compilation artifact as a 
+// Storable entity, containing the storage slot and the type of the variable
+struct Storable<N> {
+        slot: Field,
+        typ: str<N>
+    }
+

noir/noir-repo/aztec_macros/src/lib.rs

@@ -5,16 +5,15 @@ use transforms::{
     compute_note_hash_and_nullifier::inject_compute_note_hash_and_nullifier,
     events::{generate_selector_impl, transform_events},
     functions::{transform_function, transform_unconstrained},
-    note_interface::generate_note_interface_impl,
+    note_interface::{generate_note_interface_impl, inject_note_exports},
     storage::{
         assign_storage_slots, check_for_storage_definition, check_for_storage_implementation,
-        generate_storage_implementation,
+        generate_storage_implementation, generate_storage_layout,
     },
 };
 
-use noirc_frontend::{
-    hir::def_collector::dc_crate::{UnresolvedFunctions, UnresolvedTraitImpl},
-    macros_api::{CrateId, FileId, HirContext, MacroError, MacroProcessor, SortedModule, Span},
+use noirc_frontend::macros_api::{
+    CrateId, FileId, HirContext, MacroError, MacroProcessor, SortedModule, Span,
 };
 
 use utils::{
@@ -36,16 +35,6 @@ impl MacroProcessor for AztecMacro {
         transform(ast, crate_id, file_id, context)
     }
 
-    fn process_collected_defs(
-        &self,
-        crate_id: &CrateId,
-        context: &mut HirContext,
-        collected_trait_impls: &[UnresolvedTraitImpl],
-        collected_functions: &mut [UnresolvedFunctions],
-    ) -> Result<(), (MacroError, FileId)> {
-        transform_collected_defs(crate_id, context, collected_trait_impls, collected_functions)
-    }
-
     fn process_typed_ast(
         &self,
         crate_id: &CrateId,
@@ -95,6 +84,7 @@ fn transform_module(module: &mut SortedModule) -> Result<bool, AztecMacroError>
         if !check_for_storage_implementation(module, &storage_struct_name) {
             generate_storage_implementation(module, &storage_struct_name)?;
         }
+        generate_storage_layout(module, storage_struct_name)?;
     }
 
     for structure in module.types.iter_mut() {
@@ -185,24 +175,6 @@ fn transform_module(module: &mut SortedModule) -> Result<bool, AztecMacroError>
     Ok(has_transformed_module)
 }
 
-fn transform_collected_defs(
-    crate_id: &CrateId,
-    context: &mut HirContext,
-    collected_trait_impls: &[UnresolvedTraitImpl],
-    collected_functions: &mut [UnresolvedFunctions],
-) -> Result<(), (MacroError, FileId)> {
-    if has_aztec_dependency(crate_id, context) {
-        inject_compute_note_hash_and_nullifier(
-            crate_id,
-            context,
-            collected_trait_impls,
-            collected_functions,
-        )
-    } else {
-        Ok(())
-    }
-}
-
 //
 //                    Transform Hir Nodes for Aztec
 //
@@ -212,6 +184,12 @@ fn transform_hir(
     crate_id: &CrateId,
     context: &mut HirContext,
 ) -> Result<(), (AztecMacroError, FileId)> {
-    transform_events(crate_id, context)?;
-    assign_storage_slots(crate_id, context)
+    if has_aztec_dependency(crate_id, context) {
+        transform_events(crate_id, context)?;
+        inject_compute_note_hash_and_nullifier(crate_id, context)?;
+        assign_storage_slots(crate_id, context)?;
+        inject_note_exports(crate_id, context)
+    } else {
+        Ok(())
+    }
 }

noir/noir-repo/aztec_macros/src/transforms/compute_note_hash_and_nullifier.rs

@@ -1,48 +1,43 @@
 use noirc_errors::{Location, Span};
 use noirc_frontend::{
     graph::CrateId,
-    hir::{
-        def_collector::dc_crate::{UnresolvedFunctions, UnresolvedTraitImpl},
-        def_map::{LocalModuleId, ModuleId},
-    },
-    macros_api::{FileId, HirContext, MacroError},
-    node_interner::FuncId,
-    parse_program, FunctionReturnType, ItemVisibility, NoirFunction, UnresolvedTypeData,
+    macros_api::{FileId, HirContext},
+    parse_program, FunctionReturnType, NoirFunction, Type, UnresolvedTypeData,
 };
 
-use crate::utils::hir_utils::fetch_struct_trait_impls;
+use crate::utils::{
+    errors::AztecMacroError,
+    hir_utils::{collect_crate_functions, fetch_notes, get_contract_module_data, inject_fn},
+};
 
 // Check if "compute_note_hash_and_nullifier(AztecAddress,Field,Field,Field,[Field; N]) -> [Field; 4]" is defined
 fn check_for_compute_note_hash_and_nullifier_definition(
-    functions_data: &[(LocalModuleId, FuncId, NoirFunction)],
-    module_id: LocalModuleId,
+    crate_id: &CrateId,
+    context: &HirContext,
 ) -> bool {
-    functions_data.iter().filter(|func_data| func_data.0 == module_id).any(|func_data| {
-        func_data.2.def.name.0.contents == "compute_note_hash_and_nullifier"
-                && func_data.2.def.parameters.len() == 5
-                && match &func_data.2.def.parameters[0].typ.typ {
-                    UnresolvedTypeData::Named(path, _, _) => path.segments.last().unwrap().0.contents == "AztecAddress",
-                    _ => false,
-                }
-                && func_data.2.def.parameters[1].typ.typ == UnresolvedTypeData::FieldElement
-                && func_data.2.def.parameters[2].typ.typ == UnresolvedTypeData::FieldElement
-                && func_data.2.def.parameters[3].typ.typ == UnresolvedTypeData::FieldElement
-                // checks if the 5th parameter is an array and the Box<UnresolvedType> in
-                // Array(Option<UnresolvedTypeExpression>, Box<UnresolvedType>) contains only fields
-                && match &func_data.2.def.parameters[4].typ.typ {
-                    UnresolvedTypeData::Array(_, inner_type) => {
-                        matches!(inner_type.typ, UnresolvedTypeData::FieldElement)
-                    },
-                    _ => false,
-                }
+    collect_crate_functions(crate_id, context).iter().any(|funct_id| {
+        let func_data = context.def_interner.function_meta(funct_id);
+        let func_name = context.def_interner.function_name(funct_id);
+        func_name == "compute_note_hash_and_nullifier"
+                && func_data.parameters.len() == 5
+                && func_data.parameters.0.first().is_some_and(| (_, typ, _) | match typ {
+                    Type::Struct(struct_typ, _) => struct_typ.borrow().name.0.contents == "AztecAddress",
+                    _ => false
+                })
+                && func_data.parameters.0.get(1).is_some_and(|(_, typ, _)| typ.is_field())
+                && func_data.parameters.0.get(2).is_some_and(|(_, typ, _)| typ.is_field())
+                && func_data.parameters.0.get(3).is_some_and(|(_, typ, _)|  typ.is_field())
+                // checks if the 5th parameter is an array and contains only fields
+                && func_data.parameters.0.get(4).is_some_and(|(_, typ, _)|  match typ {
+                    Type::Array(_, inner_type) => inner_type.to_owned().is_field(),
+                    _ => false
+                })
                 // We check the return type the same way as we did the 5th parameter
-                && match &func_data.2.def.return_type {
+                && match &func_data.return_type {
                     FunctionReturnType::Default(_) => false,
                     FunctionReturnType::Ty(unresolved_type) => {
                         match &unresolved_type.typ {
-                            UnresolvedTypeData::Array(_, inner_type) => {
-                                matches!(inner_type.typ, UnresolvedTypeData::FieldElement)
-                            },
+                            UnresolvedTypeData::Array(_, inner_type) => matches!(inner_type.typ, UnresolvedTypeData::FieldElement),
                             _ => false,
                         }
                     }
@@ -53,77 +48,33 @@ fn check_for_compute_note_hash_and_nullifier_definition(
 pub fn inject_compute_note_hash_and_nullifier(
     crate_id: &CrateId,
     context: &mut HirContext,
-    unresolved_traits_impls: &[UnresolvedTraitImpl],
-    collected_functions: &mut [UnresolvedFunctions],
-) -> Result<(), (MacroError, FileId)> {
-    // We first fetch modules in this crate which correspond to contracts, along with their file id.
-    let contract_module_file_ids: Vec<(LocalModuleId, FileId)> = context
-        .def_map(crate_id)
-        .expect("ICE: Missing crate in def_map")
-        .modules()
-        .iter()
-        .filter(|(_, module)| module.is_contract)
-        .map(|(idx, module)| (LocalModuleId(idx), module.location.file))
-        .collect();
-
-    // If the current crate does not contain a contract module we simply skip it.
-    if contract_module_file_ids.is_empty() {
-        return Ok(());
-    } else if contract_module_file_ids.len() != 1 {
-        panic!("Found multiple contracts in the same crate");
+) -> Result<(), (AztecMacroError, FileId)> {
+    if let Some((module_id, file_id)) = get_contract_module_data(context, crate_id) {
+        // If compute_note_hash_and_nullifier is already defined by the user, we skip auto-generation in order to provide an
+        // escape hatch for this mechanism.
+        // TODO(#4647): improve this diagnosis and error messaging.
+        if check_for_compute_note_hash_and_nullifier_definition(crate_id, context) {
+            return Ok(());
+        }
+
+        // In order to implement compute_note_hash_and_nullifier, we need to know all of the different note types the
+        // contract might use. These are the types that are marked as #[aztec(note)].
+        let note_types = fetch_notes(context)
+            .iter()
+            .map(|(_, note)| note.borrow().name.0.contents.clone())
+            .collect::<Vec<_>>();
+
+        // We can now generate a version of compute_note_hash_and_nullifier tailored for the contract in this crate.
+        let func = generate_compute_note_hash_and_nullifier(&note_types);
+
+        // And inject the newly created function into the contract.
+
+        // TODO(#4373): We don't have a reasonable location for the source code of this autogenerated function, so we simply
+        // pass an empty span. This function should not produce errors anyway so this should not matter.
+        let location = Location::new(Span::empty(0), file_id);
+
+        inject_fn(crate_id, context, func, location, module_id, file_id);
     }
-
-    let (module_id, file_id) = contract_module_file_ids[0];
-
-    // If compute_note_hash_and_nullifier is already defined by the user, we skip auto-generation in order to provide an
-    // escape hatch for this mechanism.
-    // TODO(#4647): improve this diagnosis and error messaging.
-    if collected_functions.iter().any(|coll_funcs_data| {
-        check_for_compute_note_hash_and_nullifier_definition(&coll_funcs_data.functions, module_id)
-    }) {
-        return Ok(());
-    }
-
-    // In order to implement compute_note_hash_and_nullifier, we need to know all of the different note types the
-    // contract might use. These are the types that implement the NoteInterface trait, which provides the
-    // get_note_type_id function.
-    let note_types = fetch_struct_trait_impls(context, unresolved_traits_impls, "NoteInterface");
-
-    // We can now generate a version of compute_note_hash_and_nullifier tailored for the contract in this crate.
-    let func = generate_compute_note_hash_and_nullifier(&note_types);
-
-    // And inject the newly created function into the contract.
-
-    // TODO(#4373): We don't have a reasonable location for the source code of this autogenerated function, so we simply
-    // pass an empty span. This function should not produce errors anyway so this should not matter.
-    let location = Location::new(Span::empty(0), file_id);
-
-    // These are the same things the ModCollector does when collecting functions: we push the function to the
-    // NodeInterner, declare it in the module (which checks for duplicate definitions), and finally add it to the list
-    // on collected but unresolved functions.
-
-    let func_id = context.def_interner.push_empty_fn();
-    context.def_interner.push_function(
-        func_id,
-        &func.def,
-        ModuleId { krate: *crate_id, local_id: module_id },
-        location,
-    );
-
-    context.def_map_mut(crate_id).unwrap()
-        .modules_mut()[module_id.0]
-        .declare_function(
-            func.name_ident().clone(), ItemVisibility::Public, func_id
-        ).expect(
-            "Failed to declare the autogenerated compute_note_hash_and_nullifier function, likely due to a duplicate definition. See https://github.com/AztecProtocol/aztec-packages/issues/4647."
-        );
-
-    collected_functions
-        .iter_mut()
-        .find(|fns| fns.file_id == file_id)
-        .expect("ICE: no functions found in contract file")
-        .push_fn(module_id, func_id, func.clone());
-
     Ok(())
 }
 

noir/noir-repo/aztec_macros/src/transforms/events.rs

@@ -174,7 +174,7 @@ pub fn transform_events(
     crate_id: &CrateId,
     context: &mut HirContext,
 ) -> Result<(), (AztecMacroError, FileId)> {
-    for struct_id in collect_crate_structs(crate_id, context) {
+    for (_, struct_id) in collect_crate_structs(crate_id, context) {
         let attributes = context.def_interner.struct_attributes(&struct_id);
         if attributes.iter().any(|attr| is_custom_attribute(attr, "aztec(event)")) {
             transform_event(struct_id, &mut context.def_interner)?;