Implement StorageNMap (#8635)

* Implement StorageNMap * Change copyright date to 2021 * Rewrite keys to use impl_for_tuples instead of recursion * Implement prefix iteration on StorageNMap * Implement EncodeLike for key arguments * Rename KeyGenerator::Arg to KeyGenerator::KArg * Support StorageNMap in decl_storage and #[pallet::storage] macros * Use StorageNMap in assets pallet * Support migrate_keys in StorageNMap * Reduce line characters on select files * Refactor crate imports in decl_storage macros * Some more line char reductions and doc comment update * Update UI test expectations * Revert whitespace changes to untouched files * Generate Key struct instead of a 1-tuple when only 1 pair of key and hasher is provided * Revert formatting changes to unrelated files * Introduce KeyGeneratorInner * Add tests for StorageNMap in FRAMEv2 pallet macro * Small fixes to unit tests for StorageNMap * Bump runtime metadata version * Remove unused import * Update tests to use runtime metadata v13 * Introduce and use EncodeLikeTuple as a trait bound for KArg * Add some rustdocs * Revert usage of StorageNMap in assets pallet * Make use of ext::PunctuatedTrailing * Add rustdoc for final_hash * Fix StorageNMap proc macro expansions for single key cases * Create associated const in KeyGenerator for hasher metadata * Refactor code according to comments from Basti * Add module docs for generator/nmap.rs * Re-export storage::Key as NMapKey in pallet prelude * Seal the EncodeLikeTuple trait * Extract sealing code out of key.rs Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>
2026-07-07 23:37:22 +00:00 · 2021-05-14 02:44:29 -07:00
parent c6b1240e51
commit 033d8289f0
26 changed files with 3210 additions and 50 deletions
@@ -70,7 +70,9 @@ impl syn::parse::Parse for DeclStorageDef {
 /// Extended version of `DeclStorageDef` with useful precomputed value.
 pub struct DeclStorageDefExt {
 	/// Name of the module used to import hidden imports.
-	hidden_crate: Option<syn::Ident>,
+	hidden_crate: proc_macro2::TokenStream,
+	/// Hidden imports used by the module.
+	hidden_imports: proc_macro2::TokenStream,
 	/// Visibility of store trait.
 	visibility: syn::Visibility,
 	/// Name of store trait: usually `Store`.
@@ -108,9 +110,15 @@ pub struct DeclStorageDefExt {

 impl From<DeclStorageDef> for DeclStorageDefExt {
 	fn from(mut def: DeclStorageDef) -> Self {
+		let hidden_crate_name = def.hidden_crate.as_ref().map(|i| i.to_string())
+			.unwrap_or_else(|| "decl_storage".to_string());
+	
+		let hidden_crate = generate_crate_access(&hidden_crate_name, "frame-support");
+		let hidden_imports = generate_hidden_includes(&hidden_crate_name, "frame-support");
+
 		let storage_lines = def.storage_lines.drain(..).collect::<Vec<_>>();
 		let storage_lines = storage_lines.into_iter()
-			.map(|line| StorageLineDefExt::from_def(line, &def))
+			.map(|line| StorageLineDefExt::from_def(line, &def, &hidden_crate))
 			.collect();

 		let (
@@ -144,7 +152,8 @@ impl From<DeclStorageDef> for DeclStorageDefExt {
 		);

 		Self {
-			hidden_crate: def.hidden_crate,
+			hidden_crate,
+			hidden_imports,
 			visibility: def.visibility,
 			store_trait: def.store_trait,
 			module_name: def.module_name,
@@ -230,7 +239,11 @@ pub struct StorageLineDefExt {
 }

 impl StorageLineDefExt {
-	fn from_def(storage_def: StorageLineDef, def: &DeclStorageDef) -> Self {
+	fn from_def(
+		storage_def: StorageLineDef,
+		def: &DeclStorageDef,
+		hidden_crate: &proc_macro2::TokenStream,
+	) -> Self {
 		let is_generic = match &storage_def.storage_type {
 			StorageLineTypeDef::Simple(value) => {
 				ext::type_contains_ident(&value, &def.module_runtime_generic)
@@ -244,12 +257,17 @@ impl StorageLineDefExt {
 					|| ext::type_contains_ident(&map.key2, &def.module_runtime_generic)
 					|| ext::type_contains_ident(&map.value, &def.module_runtime_generic)
 			}
+			StorageLineTypeDef::NMap(map) => {
+				map.keys.iter().any(|key| ext::type_contains_ident(key, &def.module_runtime_generic))
+					|| ext::type_contains_ident(&map.value, &def.module_runtime_generic)
+			}
 		};

 		let query_type = match &storage_def.storage_type {
 			StorageLineTypeDef::Simple(value) => value.clone(),
 			StorageLineTypeDef::Map(map) => map.value.clone(),
 			StorageLineTypeDef::DoubleMap(map) => map.value.clone(),
+			StorageLineTypeDef::NMap(map) => map.value.clone(),
 		};
 		let is_option = ext::extract_type_option(&query_type).is_some();
 		let value_type = ext::extract_type_option(&query_type).unwrap_or_else(|| query_type.clone());
@@ -295,6 +313,10 @@ impl StorageLineDefExt {
 				let key2 = &map.key2;
 				quote!( StorageDoubleMap<#key1, #key2, #value_type> )
 			},
+			StorageLineTypeDef::NMap(map) => {
+				let keygen = map.to_keygen_struct(hidden_crate);
+				quote!( StorageNMap<#keygen, #value_type> )
+			}
 		};

 		let storage_trait = quote!( storage::#storage_trait_truncated );
@@ -332,6 +354,7 @@ impl StorageLineDefExt {
 pub enum StorageLineTypeDef {
 	Map(MapDef),
 	DoubleMap(Box<DoubleMapDef>),
+	NMap(NMapDef),
 	Simple(syn::Type),
 }

@@ -351,6 +374,42 @@ pub struct DoubleMapDef {
 	pub value: syn::Type,
 }

+pub struct NMapDef {
+	pub hashers: Vec<HasherKind>,
+	pub keys: Vec<syn::Type>,
+	pub value: syn::Type,
+}
+
+impl NMapDef {
+	fn to_keygen_struct(&self, scrate: &proc_macro2::TokenStream) -> proc_macro2::TokenStream {
+		if self.keys.len() == 1 {
+			let hasher = &self.hashers[0].to_storage_hasher_struct();
+			let key = &self.keys[0];
+			return quote!( #scrate::storage::types::Key<#scrate::#hasher, #key> );
+		}
+
+		let key_hasher = self.keys.iter().zip(&self.hashers).map(|(key, hasher)| {
+			let hasher = hasher.to_storage_hasher_struct();
+			quote!( #scrate::storage::types::Key<#scrate::#hasher, #key> )
+		})
+		.collect::<Vec<_>>();
+		quote!(( #(#key_hasher,)* ))
+	}
+
+	fn to_key_tuple(&self) -> proc_macro2::TokenStream {
+		if self.keys.len() == 1 {
+			let key = &self.keys[0];
+			return quote!(#key);
+		}
+
+		let tuple = self.keys.iter().map(|key| {
+			quote!(#key)
+		})
+		.collect::<Vec<_>>();
+		quote!(( #(#tuple,)* ))
+	}
+}
+
 pub struct ExtraGenesisLineDef {
 	attrs: Vec<syn::Attribute>,
 	name: syn::Ident,
@@ -402,26 +461,24 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr

 	print_pallet_upgrade::maybe_print_pallet_upgrade(&def_ext);

-	let hidden_crate_name = def_ext.hidden_crate.as_ref().map(|i| i.to_string())
-		.unwrap_or_else(|| "decl_storage".to_string());
-
-	let scrate = generate_crate_access(&hidden_crate_name, "frame-support");
-	let scrate_decl = generate_hidden_includes(&hidden_crate_name, "frame-support");
-
+	let scrate = &def_ext.hidden_crate;
+	let scrate_decl = &def_ext.hidden_imports;
 	let store_trait = store_trait::decl_and_impl(&def_ext);
-	let getters = getters::impl_getters(&scrate, &def_ext);
-	let metadata = metadata::impl_metadata(&scrate, &def_ext);
-	let instance_trait = instance_trait::decl_and_impl(&scrate, &def_ext);
-	let genesis_config = genesis_config::genesis_config_and_build_storage(&scrate, &def_ext);
-	let storage_struct = storage_struct::decl_and_impl(&scrate, &def_ext);
+	let getters = getters::impl_getters(&def_ext);
+	let metadata = metadata::impl_metadata(&def_ext);
+	let instance_trait = instance_trait::decl_and_impl(&def_ext);
+	let genesis_config = genesis_config::genesis_config_and_build_storage(&def_ext);
+	let storage_struct = storage_struct::decl_and_impl(&def_ext);

 	quote!(
 		use #scrate::{
 			StorageValue as _,
 			StorageMap as _,
 			StorageDoubleMap as _,
+			StorageNMap as _,
 			StoragePrefixedMap as _,
 			IterableStorageMap as _,
+			IterableStorageNMap as _,
 			IterableStorageDoubleMap as _,
 		};