fix!: fix storage layout export (#8880)

Closes: #8532 Includes the contract name in the exported storage layout, so that the artifact generator can distinguish between main and dependent contracts.
AztecProtocol · Sep 30, 2024 · c8f43b3 · c8f43b3
1 parent ed1e23e
commit c8f43b3
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 8 deletions.
diff --git a/noir-projects/aztec-nr/aztec/src/macros/mod.nr b/noir-projects/aztec-nr/aztec/src/macros/mod.nr
@@ -47,8 +47,8 @@ comptime fn generate_contract_interface(m: Module) -> Quoted {
     let storage_layout_getter = if has_storage_layout {
         let storage_layout_name = STORAGE_LAYOUT_NAME.get(m).unwrap();
         quote {
-            pub fn storage_layout() -> StorageLayout {
-                $storage_layout_name
+            pub fn storage_layout() -> StorageLayoutFields {
+                $storage_layout_name.fields
             }
         }
     } else {

diff --git a/noir-projects/aztec-nr/aztec/src/macros/storage/mod.nr b/noir-projects/aztec-nr/aztec/src/macros/storage/mod.nr
@@ -1,4 +1,4 @@
-use std::{collections::umap::UHashMap, hash::{BuildHasherDefault, poseidon2::Poseidon2Hasher}};
+use std::{collections::umap::UHashMap, hash::{BuildHasherDefault, poseidon2::Poseidon2Hasher}, meta::unquote};
 
 use super::utils::get_serialized_size;
 use super::utils::is_note;
@@ -57,18 +57,25 @@ pub comptime fn storage(s: StructDefinition) -> Quoted {
     let module = s.module();
     let module_name = module.name();
     let storage_layout_name = f"STORAGE_LAYOUT_{module_name}".quoted_contents();
+    let module_name_str = module_name.as_str_quote();
     STORAGE_LAYOUT_NAME.insert(module, storage_layout_name);
 
     quote {
         $storage_impl
 
-        struct StorageLayout {
+        struct StorageLayoutFields {
             $storage_layout_fields
         }
 
+        struct StorageLayout<let N: u32> {
+            contract_name: str<N>,
+            fields: StorageLayoutFields
+        }
+
         #[abi(storage)]
         global $storage_layout_name = StorageLayout {
-            $storage_layout_constructors
+            contract_name: $module_name_str,
+            fields: StorageLayoutFields { $storage_layout_constructors }
         };
     }
 }

diff --git a/yarn-project/types/src/abi/contract_artifact.ts b/yarn-project/types/src/abi/contract_artifact.ts
@@ -2,6 +2,7 @@ import {
   type ABIParameter,
   type ABIParameterVisibility,
   type AbiType,
+  type BasicValue,
   type ContractArtifact,
   type ContractNote,
   type FieldLayout,
@@ -217,10 +218,23 @@ function hasKernelFunctionInputs(params: ABIParameter[]): boolean {
  * @returns A storage layout for the contract.
  */
 function getStorageLayout(input: NoirCompiledContract) {
-  const storage = input.outputs.globals.storage ? (input.outputs.globals.storage[0] as StructValue) : { fields: [] };
-  const storageFields = storage.fields as TypedStructFieldValue<StructValue>[];
+  // If another contract is imported by the main contract, its storage layout its going to also show up here.
+  // The layout export includes the contract name, so here we can find the one that belongs to the current one and
+  // ignore the rest.
+  const storageExports = input.outputs.globals.storage ? (input.outputs.globals.storage as StructValue[]) : [];
+  const storageForContract = storageExports.find(storageExport => {
+    const contractNameField = storageExport.fields.find(field => field.name === 'contract_name')?.value as BasicValue<
+      'string',
+      string
+    >;
+    return contractNameField.value === input.name;
+  });
+  const storageFields = storageForContract
+    ? ((storageForContract.fields.find(field => field.name == 'fields') as TypedStructFieldValue<StructValue>).value
+        .fields as TypedStructFieldValue<StructValue>[])
+    : [];
 
-  if (!storageFields) {
+  if (storageFields.length === 0) {
     return {};
   }