Implement a CountedStorageMap

frame/example/src/lib.rs

@@ -598,6 +598,9 @@ pub mod pallet {
 	#[pallet::getter(fn foo)]
 	pub(super) type Foo<T: Config> = StorageValue<_, T::Balance, ValueQuery>;
 
+	#[pallet::storage]
+	pub type CountedMap<T> = CountedStorageMap<_, Blake2_128Concat, u8, u16>;
+
 	// The genesis config type.
 	#[pallet::genesis_config]
 	pub struct GenesisConfig<T: Config> {

frame/example/src/tests.rs

@@ -175,6 +175,15 @@ fn signed_ext_watch_dummy_works() {
 	})
 }
 
+#[test]
+fn counted_map_works() {
+	new_test_ext().execute_with(|| {
+		assert_eq!(CountedMap::<Test>::count(), 0);
+		CountedMap::<Test>::insert(3, 3);
+		assert_eq!(CountedMap::<Test>::count(), 1);
+	})
+}
+
 #[test]
 fn weights_work() {
 	// must have a defined weight.

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

@@ -18,30 +18,66 @@
 use crate::pallet::{Def, parse::storage::StorageDef};
 use crate::pallet::parse::storage::{Metadata, QueryKind, StorageGenerics};
 use frame_support_procedural_tools::clean_type_string;
-use std::collections::HashSet;
+use std::collections::HashMap;
 
-/// Generate the prefix_ident related the the storage.
+/// Generate the prefix_ident related to the storage.
 /// prefix_ident is used for the prefix struct to be given to storage as first generic param.
 fn prefix_ident(storage: &StorageDef) -> syn::Ident {
 	let storage_ident = &storage.ident;
 	syn::Ident::new(&format!("_GeneratedPrefixForStorage{}", storage_ident), storage_ident.span())
 }
 
+/// Generate the counter_prefix_ident related to the storage.
+/// counter_prefix_ident is used for the prefix struct to be given to counted storage map.
+fn counter_prefix_ident(storage_ident: &syn::Ident) -> syn::Ident {
+	syn::Ident::new(
+		&format!("_GeneratedCounterPrefixForStorage{}", storage_ident),
+		storage_ident.span(),
+	)
+}
+
+/// 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.
 fn check_prefix_duplicates(
 	storage_def: &StorageDef,
-	set: &mut HashSet<String>,
+	// A hashmap of all already used prefix and their associated error if duplication
+	set: &mut HashMap<String, syn::Error>,
 ) -> syn::Result<()> {
 	let prefix = storage_def.prefix();
+	let dup_err = syn::Error::new(
+		storage_def.prefix_span(),
+		format!("Duplicate storage prefixes found for `{}`", prefix),
+	);
+
+	if let Some(other_dup_err) = set.insert(prefix.clone(), dup_err.clone()) {
+		let mut err = dup_err;
+		err.combine(other_dup_err);
+		return Err(err);
+	}
 
-	if !set.insert(prefix.clone()) {
-		let err = syn::Error::new(
+	if let Metadata::CountedMap { .. } = storage_def.metadata {
+		let counter_prefix = counter_prefix(&prefix);
+		let counter_dup_err = syn::Error::new(
 			storage_def.prefix_span(),
-			format!("Duplicate storage prefixes found for `{}`", prefix),
+			format!(
+				"Duplicate storage prefixes found for `{}`, used for counter associated to \
+				counted storage map",
+				counter_prefix,
+			),
 		);
-		return Err(err);
+
+		if let Some(other_dup_err) = set.insert(counter_prefix.clone(), counter_dup_err.clone()) {
+			let mut err = counter_dup_err;
+			err.combine(other_dup_err);
+			return Err(err);
+		}
 	}
 
 	Ok(())
@@ -52,11 +88,8 @@ fn check_prefix_duplicates(
 /// * Add `#[allow(type_alias_bounds)]`
 pub fn process_generics(def: &mut Def) -> syn::Result<()> {
 	let frame_support = &def.frame_support;
-	let mut prefix_set = HashSet::new();
 
 	for storage_def in def.storages.iter_mut() {
-		check_prefix_duplicates(storage_def, &mut prefix_set)?;
-
 		let item = &mut def.item.content.as_mut().expect("Checked by def").1[storage_def.index];
 
 		let typ_item = match item {
@@ -112,6 +145,19 @@ pub fn process_generics(def: &mut Def) -> syn::Result<()> {
 					let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
 					args.args.push(syn::GenericArgument::Type(max_values));
 				}
+				StorageGenerics::CountedMap {
+					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));
+					let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
+					args.args.push(syn::GenericArgument::Type(query_kind));
+					let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone());
+					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::DoubleMap {
 					hasher1, key1, hasher2, key2, value, query_kind, on_empty, max_values,
 				} => {
@@ -157,11 +203,19 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		return e.into_compile_error().into();
 	}
 
+	// Check for duplicate prefixes
+	let mut prefix_set = HashMap::new();
+	let mut errors = def.storages.iter()
+		.filter_map(|storage_def| check_prefix_duplicates(storage_def, &mut prefix_set).err());
+	if let Some(mut final_error) = errors.next() {
+		errors.for_each(|error| final_error.combine(error));
+		return final_error.into_compile_error()
+	}
+
 	let frame_support = &def.frame_support;
 	let frame_system = &def.frame_system;
 	let pallet_ident = &def.pallet_struct.pallet;
 
-
 	let entries = def.storages.iter()
 		.map(|storage| {
 			let docs = &storage.docs;
@@ -179,6 +233,9 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 				Metadata::Map { .. } => quote::quote_spanned!(storage.attr_span =>
 					#frame_support::storage::types::StorageMapMetadata
 				),
+				Metadata::CountedMap { .. } => quote::quote_spanned!(storage.attr_span =>
+					#frame_support::storage::types::CountedStorageMapMetadata
+				),
 				Metadata::DoubleMap { .. } => quote::quote_spanned!(storage.attr_span =>
 					#frame_support::storage::types::StorageDoubleMapMetadata
 				),
@@ -208,6 +265,18 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						}
 					)
 				},
+				Metadata::CountedMap { key, value } => {
+					let value = clean_type_string(&quote::quote!(#value).to_string());
+					let key = clean_type_string(&quote::quote!(#key).to_string());
+					quote::quote_spanned!(storage.attr_span =>
+						#frame_support::metadata::StorageEntryType::Map {
+							hasher: <#full_ident as #metadata_trait>::HASHER,
+							key: #frame_support::metadata::DecodeDifferent::Encode(#key),
+							value: #frame_support::metadata::DecodeDifferent::Encode(#value),
+							unused: false,
+						}
+					)
+				},
 				Metadata::DoubleMap { key1, key2, value } => {
 					let value = clean_type_string(&quote::quote!(#value).to_string());
 					let key1 = clean_type_string(&quote::quote!(#key1).to_string());
@@ -242,7 +311,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 				}
 			};
 
-			quote::quote_spanned!(storage.attr_span =>
+			let mut metadata = quote::quote_spanned!(storage.attr_span =>
 				#(#cfg_attrs)* #frame_support::metadata::StorageEntryMetadata {
 					name: #frame_support::metadata::DecodeDifferent::Encode(
 						<#full_ident as #metadata_trait>::NAME
@@ -256,7 +325,32 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						#( #docs, )*
 					]),
 				}
-			)
+			);
+
+			// Additional metadata for some storages:
+			if let Metadata::CountedMap { .. } = storage.metadata {
+				metadata.extend(quote::quote_spanned!(storage.attr_span =>
+					, #(#cfg_attrs)* #frame_support::metadata::StorageEntryMetadata {
+						name: #frame_support::metadata::DecodeDifferent::Encode(
+							<#full_ident as #metadata_trait>::COUNTER_NAME
+						),
+						modifier: <#full_ident as #metadata_trait>::COUNTER_MODIFIER,
+						ty: #frame_support::metadata::StorageEntryType::Plain(
+							#frame_support::metadata::DecodeDifferent::Encode(
+								<#full_ident as #metadata_trait>::COUNTER_TY
+							)
+						),
+						default: #frame_support::metadata::DecodeDifferent::Encode(
+							<#full_ident as #metadata_trait>::COUNTER_DEFAULT
+						),
+						documentation: #frame_support::metadata::DecodeDifferent::Encode(&[
+							<#full_ident as #metadata_trait>::COUNTER_DOC
+						]),
+					}
+				));
+			}
+
+			metadata
 		});
 
 	let getters = def.storages.iter()
@@ -317,6 +411,27 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						}
 					)
 				},
+				Metadata::CountedMap { key, value } => {
+					let query = match storage.query_kind.as_ref().expect("Checked by def") {
+						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
+							Option<#value>
+						),
+						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 {
+							#( #docs )*
+							pub fn #getter<KArg>(k: KArg) -> #query where
+								KArg: #frame_support::codec::EncodeLike<#key>,
+							{
+								// NOTE: we can't use any trait here because CountedStorageMap
+								// doesn't implement any.
+								<#full_ident>::get(k)
+							}
+						}
+					)
+				},
 				Metadata::DoubleMap { key1, key2, value } => {
 					let query = match storage.query_kind.as_ref().expect("Checked by def") {
 						QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
@@ -381,7 +496,44 @@ 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)*
+				#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("Every active pallet has a name in the runtime; qed")
+					}
+					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()
+		};
+
 		quote::quote_spanned!(storage_def.attr_span =>
+			#maybe_counter
+
 			#(#cfg_attrs)*
 			#prefix_struct_vis struct #prefix_struct_ident<#type_use_gen>(
 				core::marker::PhantomData<(#type_use_gen,)>

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

@@ -86,6 +86,7 @@ impl syn::parse::Parse for PalletStorageAttr {
 pub enum Metadata {
 	Value { value: syn::Type },
 	Map { value: syn::Type, key: syn::Type },
+	CountedMap { value: syn::Type, key: syn::Type },
 	DoubleMap {
 		value: syn::Type,
 		key1: syn::Type,
@@ -160,6 +161,14 @@ pub enum StorageGenerics {
 		on_empty: Option<syn::Type>,
 		max_values: Option<syn::Type>,
 	},
+	CountedMap {
+		hasher: syn::Type,
+		key: syn::Type,
+		value: syn::Type,
+		query_kind: Option<syn::Type>,
+		on_empty: Option<syn::Type>,
+		max_values: Option<syn::Type>,
+	},
 	Value {
 		value: syn::Type,
 		query_kind: Option<syn::Type>,
@@ -180,6 +189,7 @@ impl StorageGenerics {
 		let res = match self.clone() {
 			Self::DoubleMap { value, key1, key2, .. } => Metadata::DoubleMap { value, key1, key2 },
 			Self::Map { value, key, .. } => Metadata::Map { value, key },
+			Self::CountedMap { value, key, .. } => Metadata::CountedMap { value, key },
 			Self::Value { value, .. } => Metadata::Value { value },
 			Self::NMap { keygen, value, .. } => Metadata::NMap {
 				keys: collect_keys(&keygen)?,
@@ -196,6 +206,7 @@ impl StorageGenerics {
 		match &self {
 			Self::DoubleMap { query_kind, .. }
 			| Self::Map { query_kind, .. }
+			| Self::CountedMap { query_kind, .. }
 			| Self::Value { query_kind, .. }
 			| Self::NMap { query_kind, .. }
 				=> query_kind.clone(),
@@ -206,6 +217,7 @@ impl StorageGenerics {
 enum StorageKind {
 	Value,
 	Map,
+	CountedMap,
 	DoubleMap,
 	NMap,
 }
@@ -333,6 +345,30 @@ fn process_named_generics(
 				max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
 			}
 		}
+		StorageKind::CountedMap => {
+			check_generics(
+				&parsed,
+				&["Hasher", "Key", "Value"],
+				&["QueryKind", "OnEmpty", "MaxValues"],
+				"CountedStorageMap",
+				args_span,
+			)?;
+
+			StorageGenerics::CountedMap {
+				hasher: parsed.remove("Hasher")
+					.map(|binding| binding.ty)
+					.expect("checked above as mandatory generic"),
+				key: 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),
+			}
+		}
 		StorageKind::DoubleMap => {
 			check_generics(
 				&parsed,
@@ -436,6 +472,14 @@ fn process_unnamed_generics(
 			},
 			retrieve_arg(4).ok(),
 		),
+		StorageKind::CountedMap => (
+			None,
+			Metadata::CountedMap {
+				key: retrieve_arg(2)?,
+				value: retrieve_arg(3)?,
+			},
+			retrieve_arg(4).ok(),
+		),
 		StorageKind::DoubleMap => (
 			None,
 			Metadata::DoubleMap {
@@ -470,6 +514,7 @@ fn process_generics(
 	let storage_kind = match &*segment.ident.to_string() {
 		"StorageValue" => StorageKind::Value,
 		"StorageMap" => StorageKind::Map,
+		"CountedStorageMap" => StorageKind::CountedMap,
 		"StorageDoubleMap" => StorageKind::DoubleMap,
 		"StorageNMap" => StorageKind::NMap,
 		found => {