Implement a CountedStorageMap

frame/example/src/lib.rs

@@ -605,6 +605,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

@@ -180,6 +180,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

@@ -19,27 +19,68 @@ use crate::pallet::{
 	parse::storage::{Metadata, QueryKind, StorageDef, StorageGenerics},
 	Def,
 };
-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>) -> syn::Result<()> {
+fn check_prefix_duplicates(
+	storage_def: &StorageDef,
+	// A hashmap of all already used prefix and their associated error if duplication
+	used_prefixes: &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) = used_prefixes.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) =
+			used_prefixes.insert(counter_prefix.clone(), counter_dup_err.clone())
+		{
+			let mut err = counter_dup_err;
+			err.combine(other_dup_err);
+			return Err(err)
+		}
 	}
 
 	Ok(())
@@ -51,11 +92,8 @@ fn check_prefix_duplicates(storage_def: &StorageDef, set: &mut HashSet<String>)
 /// * 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 {
@@ -109,6 +147,24 @@ 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,
@@ -162,11 +218,22 @@ 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 entries_builder = def.storages.iter().map(|storage| {
 		let docs = &storage.docs;
 
 		let ident = &storage.ident;
@@ -176,14 +243,14 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		let cfg_attrs = &storage.cfg_attrs;
 
 		quote::quote_spanned!(storage.attr_span =>
-			#(#cfg_attrs)* #frame_support::metadata::StorageEntryMetadata {
-				name: <#full_ident as #frame_support::storage::StorageEntryMetadata>::NAME,
-				modifier: <#full_ident as #frame_support::storage::StorageEntryMetadata>::MODIFIER,
-				ty: <#full_ident as #frame_support::storage::StorageEntryMetadata>::ty(),
-				default: <#full_ident as #frame_support::storage::StorageEntryMetadata>::default(),
-				docs: #frame_support::sp_std::vec![
-					#( #docs, )*
-				],
+			#(#cfg_attrs)*
+			{
+				<#full_ident as #frame_support::storage::StorageEntryMetadataBuilder>::build_metadata(
+					#frame_support::sp_std::vec![
+						#( #docs, )*
+					],
+					&mut entries,
+				);
 			}
 		)
 	});
@@ -246,6 +313,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 =>
@@ -311,7 +399,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,)>
@@ -351,9 +476,12 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						#frame_support::traits::PalletInfo
 					>::name::<#pallet_ident<#type_use_gen>>()
 						.expect("Every active pallet has a name in the runtime; qed"),
-					entries: #frame_support::sp_std::vec![
-						#( #entries, )*
-					],
+					entries: {
+						#[allow(unused_mut)]
+						let mut entries = #frame_support::sp_std::vec![];
+						#( #entries_builder )*
+						entries
+					},
 				}
 			}
 		}

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, key2: syn::Type },
 	NMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
 }
@@ -153,6 +154,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>,
@@ -173,6 +182,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)?, keygen, value },
@@ -186,6 +196,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(),
 		}
@@ -195,6 +206,7 @@ impl StorageGenerics {
 enum StorageKind {
 	Value,
 	Map,
+	CountedMap,
 	DoubleMap,
 	NMap,
 }
@@ -324,6 +336,33 @@ 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,
@@ -425,6 +464,11 @@ fn process_unnamed_generics(
 			Metadata::Map { key: retrieve_arg(2)?, value: retrieve_arg(3)? },
 			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 {
@@ -451,6 +495,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 => {