[Feature] Introduce storage_alias for CountedStorageMap (#13366)

* [Feature] Introduce storagage_alias for CountedStorageMap

* bit more dry

* bit more dry

* address review comments

* some tests and fixes

* fix ui tests

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <[email protected]>

* compare metadata

---------

Co-authored-by: parity-processbot <>
Co-authored-by: Bastian Köcher <[email protected]>

Author: ruseinov

frame/support/procedural/src/lib.rs

@@ -69,6 +69,12 @@ fn get_cargo_env_var<T: FromStr>(version_env: &str) -> std::result::Result<T, ()
 	T::from_str(&version).map_err(drop)
 }
 
+/// Generate the counter_prefix related to the storage.
+/// counter_prefix is used by counted storage map.
+fn counter_prefix(prefix: &str) -> String {
+	format!("CounterFor{}", prefix)
+}
+
 /// Declares strongly-typed wrappers around codec-compatible types in storage.
 ///
 /// ## Example

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

@@ -15,9 +15,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::pallet::{
-	parse::storage::{Metadata, QueryKind, StorageDef, StorageGenerics},
-	Def,
+use crate::{
+	counter_prefix,
+	pallet::{
+		parse::storage::{Metadata, QueryKind, StorageDef, StorageGenerics},
+		Def,
+	},
 };
 use quote::ToTokens;
 use std::{collections::HashMap, ops::IndexMut};
@@ -39,12 +42,6 @@ fn counter_prefix_ident(storage_ident: &syn::Ident) -> syn::Ident {
 	)
 }
 
-/// Generate the counter_prefix related to the storage.
-/// counter_prefix is used by counted storage map.
-fn counter_prefix(prefix: &str) -> String {
-	format!("CounterFor{}", prefix)
-}
-
 /// Check for duplicated storage prefixes. This step is necessary since users can specify an
 /// alternative storage prefix using the #[pallet::storage_prefix] syntax, and we need to ensure
 /// that the prefix specified by the user is not a duplicate of an existing one.

frame/support/procedural/src/storage_alias.rs

@@ -17,6 +17,7 @@
 
 //! Implementation of the `storage_alias` attribute macro.
 
+use crate::counter_prefix;
 use frame_support_procedural_tools::generate_crate_access_2018;
 use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens};
@@ -136,6 +137,7 @@ impl ToTokens for SimpleGenerics {
 mod storage_types {
 	syn::custom_keyword!(StorageValue);
 	syn::custom_keyword!(StorageMap);
+	syn::custom_keyword!(CountedStorageMap);
 	syn::custom_keyword!(StorageDoubleMap);
 	syn::custom_keyword!(StorageNMap);
 }
@@ -168,6 +170,21 @@ enum StorageType {
 		_trailing_comma: Option<Token![,]>,
 		_gt_token: Token![>],
 	},
+	CountedMap {
+		_kw: storage_types::CountedStorageMap,
+		_lt_token: Token![<],
+		prefix: SimplePath,
+		prefix_generics: Option<TypeGenerics>,
+		_hasher_comma: Token![,],
+		hasher_ty: Type,
+		_key_comma: Token![,],
+		key_ty: Type,
+		_value_comma: Token![,],
+		value_ty: Type,
+		query_type: Option<(Token![,], Type)>,
+		_trailing_comma: Option<Token![,]>,
+		_gt_token: Token![>],
+	},
 	DoubleMap {
 		_kw: storage_types::StorageDoubleMap,
 		_lt_token: Token![<],
@@ -235,12 +252,22 @@ impl StorageType {
 					>;
 				}
 			},
+			Self::CountedMap {
+				value_ty, query_type, hasher_ty, key_ty, prefix_generics, ..
+			} |
 			Self::Map { value_ty, query_type, hasher_ty, key_ty, prefix_generics, .. } => {
 				let query_type = query_type.as_ref().map(|(c, t)| quote!(#c #t));
+				let map_type = Ident::new(
+					match self {
+						Self::Map { .. } => "StorageMap",
+						_ => "CountedStorageMap",
+					},
+					Span::call_site(),
+				);
 
 				quote! {
 					#( #attributes )*
-					#visibility type #storage_name #storage_generics = #crate_::storage::types::StorageMap<
+					#visibility type #storage_name #storage_generics = #crate_::storage::types::#map_type<
 						#storage_instance #prefix_generics,
 						#hasher_ty,
 						#key_ty,
@@ -296,6 +323,7 @@ impl StorageType {
 		match self {
 			Self::Value { prefix, .. } |
 			Self::Map { prefix, .. } |
+			Self::CountedMap { prefix, .. } |
 			Self::NMap { prefix, .. } |
 			Self::DoubleMap { prefix, .. } => prefix,
 		}
@@ -306,6 +334,7 @@ impl StorageType {
 		match self {
 			Self::Value { prefix_generics, .. } |
 			Self::Map { prefix_generics, .. } |
+			Self::CountedMap { prefix_generics, .. } |
 			Self::NMap { prefix_generics, .. } |
 			Self::DoubleMap { prefix_generics, .. } => prefix_generics.as_ref(),
 		}
@@ -363,6 +392,22 @@ impl Parse for StorageType {
 				_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
 				_gt_token: input.parse()?,
 			})
+		} else if lookahead.peek(storage_types::CountedStorageMap) {
+			Ok(Self::CountedMap {
+				_kw: input.parse()?,
+				_lt_token: input.parse()?,
+				prefix: input.parse()?,
+				prefix_generics: parse_pallet_generics(input)?,
+				_hasher_comma: input.parse()?,
+				hasher_ty: input.parse()?,
+				_key_comma: input.parse()?,
+				key_ty: input.parse()?,
+				_value_comma: input.parse()?,
+				value_ty: input.parse()?,
+				query_type: parse_query_type(input)?,
+				_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
+				_gt_token: input.parse()?,
+			})
 		} else if lookahead.peek(storage_types::StorageDoubleMap) {
 			Ok(Self::DoubleMap {
 				_kw: input.parse()?,
@@ -476,6 +521,7 @@ pub fn storage_alias(input: TokenStream) -> Result<TokenStream> {
 		input.storage_type.prefix(),
 		input.storage_type.prefix_generics(),
 		&input.visibility,
+		matches!(input.storage_type, StorageType::CountedMap { .. }),
 	)?;
 
 	let definition = input.storage_type.generate_type_declaration(
@@ -511,6 +557,7 @@ fn generate_storage_instance(
 	prefix: &SimplePath,
 	prefix_generics: Option<&TypeGenerics>,
 	visibility: &Visibility,
+	is_counted_map: bool,
 ) -> Result<StorageInstance> {
 	if let Some(ident) = prefix.get_ident().filter(|i| *i == "_") {
 		return Err(Error::new(ident.span(), "`_` is not allowed as prefix by `storage_alias`."))
@@ -546,9 +593,37 @@ fn generate_storage_instance(
 
 	let where_clause = storage_where_clause.map(|w| quote!(#w)).unwrap_or_default();
 
-	let name = Ident::new(&format!("{}_Storage_Instance", storage_name), Span::call_site());
+	let name_str = format!("{}_Storage_Instance", storage_name);
+	let name = Ident::new(&name_str, Span::call_site());
 	let storage_name_str = storage_name.to_string();
 
+	let counter_code = is_counted_map.then(|| {
+		let counter_name = Ident::new(&counter_prefix(&name_str), Span::call_site());
+		let counter_storage_name_str = counter_prefix(&storage_name_str);
+
+		quote! {
+			#visibility struct #counter_name< #impl_generics >(
+				#crate_::sp_std::marker::PhantomData<(#type_generics)>
+			) #where_clause;
+
+			impl<#impl_generics> #crate_::traits::StorageInstance
+				for #counter_name< #type_generics > #where_clause
+			{
+				fn pallet_prefix() -> &'static str {
+					#pallet_prefix
+				}
+
+				const STORAGE_PREFIX: &'static str = #counter_storage_name_str;
+			}
+
+			impl<#impl_generics> #crate_::storage::types::CountedStorageMapInstance
+				for #name< #type_generics > #where_clause
+			{
+				type CounterPrefix = #counter_name < #type_generics >;
+			}
+		}
+	});
+
 	// Implement `StorageInstance` trait.
 	let code = quote! {
 		#visibility struct #name< #impl_generics >(
@@ -564,6 +639,8 @@ fn generate_storage_instance(
 
 			const STORAGE_PREFIX: &'static str = #storage_name_str;
 		}
+
+		#counter_code
 	};
 
 	Ok(StorageInstance { name, code })

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

@@ -543,6 +543,16 @@ mod test {
 			97
 		}
 	}
+	#[crate::storage_alias]
+	type ExampleCountedMap = CountedStorageMap<Prefix, Twox64Concat, u16, u32>;
+
+	#[test]
+	fn storage_alias_works() {
+		TestExternalities::default().execute_with(|| {
+			assert_eq!(ExampleCountedMap::count(), 0);
+			ExampleCountedMap::insert(3, 10);
+		})
+	}
 
 	#[test]
 	fn test_value_query() {

frame/support/test/tests/pallet.rs

@@ -19,7 +19,7 @@ use frame_support::{
 	dispatch::{
 		DispatchClass, DispatchInfo, GetDispatchInfo, Parameter, Pays, UnfilteredDispatchable,
 	},
-	pallet_prelude::ValueQuery,
+	pallet_prelude::{StorageInfoTrait, ValueQuery},
 	storage::unhashed,
 	traits::{
 		ConstU32, GetCallName, GetStorageVersion, OnFinalize, OnGenesis, OnInitialize,
@@ -1906,14 +1906,30 @@ fn assert_type_all_pallets_without_system_reversed_is_correct() {
 
 #[test]
 fn test_storage_alias() {
+	use frame_support::Twox64Concat;
+
 	#[frame_support::storage_alias]
 	type Value<T: pallet::Config>
 	where
 		<T as frame_system::Config>::AccountId: From<SomeType1> + SomeAssociation1,
 	= StorageValue<pallet::Pallet<T>, u32, ValueQuery>;
 
+	#[frame_support::storage_alias]
+	type SomeCountedStorageMap<T: pallet2::Config>
+	where
+		<T as frame_system::Config>::AccountId: From<SomeType1> + SomeAssociation1,
+	= CountedStorageMap<pallet2::Pallet<T>, Twox64Concat, u8, u32>;
+
 	TestExternalities::default().execute_with(|| {
 		pallet::Value::<Runtime>::put(10);
 		assert_eq!(10, Value::<Runtime>::get());
+
+		pallet2::SomeCountedStorageMap::<Runtime>::insert(10, 100);
+		assert_eq!(Some(100), SomeCountedStorageMap::<Runtime>::get(10));
+		assert_eq!(1, SomeCountedStorageMap::<Runtime>::count());
+		assert_eq!(
+			SomeCountedStorageMap::<Runtime>::storage_info(),
+			pallet2::SomeCountedStorageMap::<Runtime>::storage_info()
+		);
 	})
 }

frame/support/test/tests/storage_alias_ui/forbid_underscore_as_prefix.stderr

@@ -1,4 +1,4 @@
-error: expected one of: `StorageValue`, `StorageMap`, `StorageDoubleMap`, `StorageNMap`
+error: expected one of: `StorageValue`, `StorageMap`, `CountedStorageMap`, `StorageDoubleMap`, `StorageNMap`
  --> tests/storage_alias_ui/forbid_underscore_as_prefix.rs:2:14
   |
 2 | type Ident = CustomStorage<Hello, u32>;