CountedNMap implementation (#10621)

* add initial CountedDoubleMap implementation

* extend CountedDoubleMap functionality

* add some traits implementation for CountedStorageDoubleMap

* add basic tests for CountedStorageDoubleMap

* add mutate functions implementation

* add additional tests

* add test_option_query test

* add try_append_decode_len_works, append_decode_len_works tests

* add migrate_keys_works, translate_values tests

* add test_iter_drain_translate test

* add test_metadata test

* add remove_prefix implementation,  add test_iter_drain_prefix test

* update

* refactor PrefixIterator usage

* Fix CI build

* fix storage_ensure_span_are_ok_wrong_gen.rs storage_ensure_span_are_ok_wrong_gen_unnamed.rs

* add counted_nmap implementation

* add tests, fixes

* remove counted double map impl

* fix metadata checks

* update clear func

* fix clear, clear with prefix

* fix set function

* update

* final fix

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Anton <[email protected]>

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Anton <[email protected]>

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Anton <[email protected]>

* fix comments

* fix suggestion

* cargo update

* Relocate impl of Sealed for Ref to module root

* fix StorageEntryMetadata type

* Update frame/support/src/storage/types/nmap.rs

Co-authored-by: Guillaume Yu Thiolliere <[email protected]>

* removed StorageNMap and StoragePrefixedMap traits impl

* fix tests

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Guillaume Yu Thiolliere <[email protected]>

* extend pallet::storage macro with CountedStorageNMap usage

* fix

* add tests

* fix

* fix

* Add counter_storage_final_key(), map_storage_final_prefix() functions

* update tests

* fix

* fix

* fix

* update tests

* fix fmt

* fix fmt

---------

Co-authored-by: Anton <[email protected]>
Co-authored-by: Keith Yeung <[email protected]>
Co-authored-by: Guillaume Yu Thiolliere <[email protected]>
Co-authored-by: parity-processbot <>

Author: 4 people

frame/support/procedural/src/pallet/expand/storage.rs

@@ -194,18 +194,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
 					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
 					args.args.push(syn::GenericArgument::Type(on_empty));
 				},
-				StorageGenerics::Map { hasher, key, value, query_kind, on_empty, max_values } => {
-					args.args.push(syn::GenericArgument::Type(hasher));
-					args.args.push(syn::GenericArgument::Type(key));
-					args.args.push(syn::GenericArgument::Type(value.clone()));
-					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
-					set_result_query_type_parameter(&mut query_kind)?;
-					args.args.push(syn::GenericArgument::Type(query_kind));
-					let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
-					args.args.push(syn::GenericArgument::Type(on_empty));
-					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
-					args.args.push(syn::GenericArgument::Type(max_values));
-				},
+				StorageGenerics::Map { hasher, key, value, query_kind, on_empty, max_values } |
 				StorageGenerics::CountedMap {
 					hasher,
 					key,
@@ -248,7 +237,14 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
 					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
 					args.args.push(syn::GenericArgument::Type(max_values));
 				},
-				StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } => {
+				StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } |
+				StorageGenerics::CountedNMap {
+					keygen,
+					value,
+					query_kind,
+					on_empty,
+					max_values,
+				} => {
 					args.args.push(syn::GenericArgument::Type(keygen));
 					args.args.push(syn::GenericArgument::Type(value.clone()));
 					let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
@@ -265,7 +261,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
 
 			let (value_idx, query_idx, on_empty_idx) = match storage_def.metadata {
 				Metadata::Value { .. } => (1, 2, 3),
-				Metadata::NMap { .. } => (2, 3, 4),
+				Metadata::NMap { .. } | Metadata::CountedNMap { .. } => (2, 3, 4),
 				Metadata::Map { .. } | Metadata::CountedMap { .. } => (3, 4, 5),
 				Metadata::DoubleMap { .. } => (5, 6, 7),
 			};
@@ -359,6 +355,17 @@ fn augment_final_docs(def: &mut Def) {
 			);
 			push_string_literal(&doc_line, storage);
 		},
+		Metadata::CountedNMap { keys, value, .. } => {
+			let doc_line = format!(
+				"Storage type is [`CountedStorageNMap`] with keys type ({}) and value type {}.",
+				keys.iter()
+					.map(|k| k.to_token_stream().to_string())
+					.collect::<Vec<_>>()
+					.join(", "),
+				value.to_token_stream()
+			);
+			push_string_literal(&doc_line, storage);
+		},
 		Metadata::CountedMap { key, value } => {
 			let doc_line = format!(
 				"Storage type is [`CountedStorageMap`] with key type {} and value type {}.",
@@ -579,6 +586,36 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						}
 					)
 				},
+				Metadata::CountedNMap { keygen, value, .. } => {
+					let query = match storage.query_kind.as_ref().expect("Checked by def") {
+						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
+							Option<#value>
+						),
+						QueryKind::ResultQuery(error_path, _) => {
+							quote::quote_spanned!(storage.attr_span =>
+								Result<#value, #error_path>
+							)
+						},
+						QueryKind::ValueQuery => quote::quote!(#value),
+					};
+					quote::quote_spanned!(storage.attr_span =>
+						#(#cfg_attrs)*
+						impl<#type_impl_gen> #pallet_ident<#type_use_gen> #completed_where_clause {
+							#[doc = #getter_doc_line]
+							pub fn #getter<KArg>(key: KArg) -> #query
+							where
+								KArg: #frame_support::storage::types::EncodeLikeTuple<
+									<#keygen as #frame_support::storage::types::KeyGenerator>::KArg
+								>
+									+ #frame_support::storage::types::TupleToEncodedIter,
+							{
+								// NOTE: we can't use any trait here because CountedStorageNMap
+								// doesn't implement any.
+								<#full_ident>::get(key)
+							}
+						}
+					)
+				},
 			}
 		} else {
 			Default::default()
@@ -595,40 +632,72 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 
 		let cfg_attrs = &storage_def.cfg_attrs;
 
-		let maybe_counter = if let Metadata::CountedMap { .. } = storage_def.metadata {
-			let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident);
-			let counter_prefix_struct_const = counter_prefix(&prefix_struct_const);
-
-			quote::quote_spanned!(storage_def.attr_span =>
-				#(#cfg_attrs)*
-				#[doc(hidden)]
-				#prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>(
-					core::marker::PhantomData<(#type_use_gen,)>
-				);
-				#(#cfg_attrs)*
-				impl<#type_impl_gen> #frame_support::traits::StorageInstance
-					for #counter_prefix_struct_ident<#type_use_gen>
-					#config_where_clause
-				{
-					fn pallet_prefix() -> &'static str {
-						<
-							<T as #frame_system::Config>::PalletInfo
-							as #frame_support::traits::PalletInfo
-						>::name::<Pallet<#type_use_gen>>()
-							.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
+		let maybe_counter = match storage_def.metadata {
+			Metadata::CountedMap { .. } => {
+				let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident);
+				let counter_prefix_struct_const = counter_prefix(&prefix_struct_const);
+				quote::quote_spanned!(storage_def.attr_span =>
+					#(#cfg_attrs)*
+					#[doc(hidden)]
+					#prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>(
+						core::marker::PhantomData<(#type_use_gen,)>
+					);
+					#(#cfg_attrs)*
+					impl<#type_impl_gen> #frame_support::traits::StorageInstance
+						for #counter_prefix_struct_ident<#type_use_gen>
+						#config_where_clause
+					{
+						fn pallet_prefix() -> &'static str {
+							<
+								<T as #frame_system::Config>::PalletInfo
+								as #frame_support::traits::PalletInfo
+							>::name::<Pallet<#type_use_gen>>()
+								.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
+						}
+						const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
 					}
-					const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
-				}
-				#(#cfg_attrs)*
-				impl<#type_impl_gen> #frame_support::storage::types::CountedStorageMapInstance
-					for #prefix_struct_ident<#type_use_gen>
-					#config_where_clause
-				{
-					type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
-				}
-			)
-		} else {
-			proc_macro2::TokenStream::default()
+					#(#cfg_attrs)*
+					impl<#type_impl_gen> #frame_support::storage::types::CountedStorageMapInstance
+						for #prefix_struct_ident<#type_use_gen>
+						#config_where_clause
+					{
+						type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
+					}
+				)
+			},
+			Metadata::CountedNMap { .. } => {
+				let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident);
+				let counter_prefix_struct_const = counter_prefix(&prefix_struct_const);
+				quote::quote_spanned!(storage_def.attr_span =>
+					#(#cfg_attrs)*
+					#[doc(hidden)]
+					#prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>(
+						core::marker::PhantomData<(#type_use_gen,)>
+					);
+					#(#cfg_attrs)*
+					impl<#type_impl_gen> #frame_support::traits::StorageInstance
+						for #counter_prefix_struct_ident<#type_use_gen>
+						#config_where_clause
+					{
+						fn pallet_prefix() -> &'static str {
+							<
+								<T as #frame_system::Config>::PalletInfo
+								as #frame_support::traits::PalletInfo
+							>::name::<Pallet<#type_use_gen>>()
+								.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
+						}
+						const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
+					}
+					#(#cfg_attrs)*
+					impl<#type_impl_gen> #frame_support::storage::types::CountedStorageNMapInstance
+						for #prefix_struct_ident<#type_use_gen>
+						#config_where_clause
+					{
+						type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
+					}
+				)
+			},
+			_ => proc_macro2::TokenStream::default(),
 		};
 
 		quote::quote_spanned!(storage_def.attr_span =>

frame/support/procedural/src/pallet/parse/storage.rs

@@ -138,6 +138,7 @@ pub enum Metadata {
 	CountedMap { value: syn::Type, key: syn::Type },
 	DoubleMap { value: syn::Type, key1: syn::Type, key2: syn::Type },
 	NMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
+	CountedNMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
 }
 
 pub enum QueryKind {
@@ -230,6 +231,13 @@ pub enum StorageGenerics {
 		on_empty: Option<syn::Type>,
 		max_values: Option<syn::Type>,
 	},
+	CountedNMap {
+		keygen: syn::Type,
+		value: syn::Type,
+		query_kind: Option<syn::Type>,
+		on_empty: Option<syn::Type>,
+		max_values: Option<syn::Type>,
+	},
 }
 
 impl StorageGenerics {
@@ -242,6 +250,8 @@ impl StorageGenerics {
 			Self::Value { value, .. } => Metadata::Value { value },
 			Self::NMap { keygen, value, .. } =>
 				Metadata::NMap { keys: collect_keys(&keygen)?, keygen, value },
+			Self::CountedNMap { keygen, value, .. } =>
+				Metadata::CountedNMap { keys: collect_keys(&keygen)?, keygen, value },
 		};
 
 		Ok(res)
@@ -254,7 +264,8 @@ impl StorageGenerics {
 			Self::Map { query_kind, .. } |
 			Self::CountedMap { query_kind, .. } |
 			Self::Value { query_kind, .. } |
-			Self::NMap { query_kind, .. } => query_kind.clone(),
+			Self::NMap { query_kind, .. } |
+			Self::CountedNMap { query_kind, .. } => query_kind.clone(),
 		}
 	}
 }
@@ -265,6 +276,7 @@ enum StorageKind {
 	CountedMap,
 	DoubleMap,
 	NMap,
+	CountedNMap,
 }
 
 /// Check the generics in the `map` contains the generics in `gen` may contains generics in
@@ -493,6 +505,29 @@ fn process_named_generics(
 				max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
 			}
 		},
+		StorageKind::CountedNMap => {
+			check_generics(
+				&parsed,
+				&["Key", "Value"],
+				&["QueryKind", "OnEmpty", "MaxValues"],
+				"CountedStorageNMap",
+				args_span,
+			)?;
+
+			StorageGenerics::CountedNMap {
+				keygen: parsed
+					.remove("Key")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				value: parsed
+					.remove("Value")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				query_kind: parsed.remove("QueryKind").map(|binding| binding.ty),
+				on_empty: parsed.remove("OnEmpty").map(|binding| binding.ty),
+				max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
+			}
+		},
 	};
 
 	let metadata = generics.metadata()?;
@@ -578,6 +613,16 @@ fn process_unnamed_generics(
 				false,
 			)
 		},
+		StorageKind::CountedNMap => {
+			let keygen = retrieve_arg(1)?;
+			let keys = collect_keys(&keygen)?;
+			(
+				None,
+				Metadata::CountedNMap { keys, keygen, value: retrieve_arg(2)? },
+				retrieve_arg(3).ok(),
+				false,
+			)
+		},
 	};
 
 	Ok(res)
@@ -594,10 +639,11 @@ fn process_generics(
 		"CountedStorageMap" => StorageKind::CountedMap,
 		"StorageDoubleMap" => StorageKind::DoubleMap,
 		"StorageNMap" => StorageKind::NMap,
+		"CountedStorageNMap" => StorageKind::CountedNMap,
 		found => {
 			let msg = format!(
 				"Invalid pallet::storage, expected ident: `StorageValue` or \
-				`StorageMap` or `CountedStorageMap` or `StorageDoubleMap` or `StorageNMap` \
+				`StorageMap` or `CountedStorageMap` or `StorageDoubleMap` or `StorageNMap` or `CountedStorageNMap` \
 				in order to expand metadata, found `{}`.",
 				found,
 			);

frame/support/src/lib.rs

@@ -1545,8 +1545,8 @@ pub mod pallet_prelude {
 		storage::{
 			bounded_vec::BoundedVec,
 			types::{
-				CountedStorageMap, Key as NMapKey, OptionQuery, ResultQuery, StorageDoubleMap,
-				StorageMap, StorageNMap, StorageValue, ValueQuery,
+				CountedStorageMap, CountedStorageNMap, Key as NMapKey, OptionQuery, ResultQuery,
+				StorageDoubleMap, StorageMap, StorageNMap, StorageValue, ValueQuery,
 			},
 			StorageList,
 		},

frame/support/src/storage/mod.rs

@@ -1428,6 +1428,7 @@ mod private {
 	impl<K, V, S> Sealed for bounded_btree_map::BoundedBTreeMap<K, V, S> {}
 	impl<T, S> Sealed for bounded_btree_set::BoundedBTreeSet<T, S> {}
 	impl<T: Encode> Sealed for BTreeSet<T> {}
+	impl<'a, T: EncodeLike<U>, U: Encode> Sealed for codec::Ref<'a, T, U> {}
 
 	macro_rules! impl_sealed_for_tuple {
 		($($elem:ident),+) => {

frame/support/src/storage/types/counted_map.rs

@@ -612,8 +612,9 @@ mod test {
 			assert_eq!(A::count(), 2);
 
 			// Insert an existing key, shouldn't increment counted values.
-			A::insert(3, 11);
+			A::insert(3, 12);
 
+			assert_eq!(A::try_get(3), Ok(12));
 			assert_eq!(A::count(), 2);
 
 			// Remove non-existing.
@@ -706,17 +707,17 @@ mod test {
 			// Try succeed mutate existing to existing.
 			A::try_mutate_exists(1, |query| {
 				assert_eq!(*query, Some(43));
-				*query = Some(43);
+				*query = Some(45);
 				Result::<(), ()>::Ok(())
 			})
 			.unwrap();
 
-			assert_eq!(A::try_get(1), Ok(43));
+			assert_eq!(A::try_get(1), Ok(45));
 			assert_eq!(A::count(), 4);
 
 			// Try succeed mutate existing to non-existing.
 			A::try_mutate_exists(1, |query| {
-				assert_eq!(*query, Some(43));
+				assert_eq!(*query, Some(45));
 				*query = None;
 				Result::<(), ()>::Ok(())
 			})