Allow to specify some max number of values for storages in pallet macro.

frame/support/procedural/src/lib.rs

@@ -156,6 +156,9 @@ use proc_macro::TokenStream;
 /// * \[optional\] `config(#field_name)`: `field_name` is optional if get is set.
 /// Will include the item in `GenesisConfig`.
 /// * \[optional\] `build(#closure)`: Closure called with storage overlays.
+/// * \[optional\] `max_values(#expr)`: `expr` is an expression returning a `u32`. It is used to
+/// implement `StorageInfoTrait`. Note this attribute is not available for storage value as the maximum
+/// number of values is 1.
 /// * `#type`: Storage type.
 /// * \[optional\] `#default`: Value returned when none.
 ///
@@ -234,11 +237,20 @@ use proc_macro::TokenStream;
 /// add_extra_genesis {
 /// 	config(phantom): std::marker::PhantomData<I>,
 /// }
-/// ...
+/// ```
 ///
 /// This adds a field to your `GenesisConfig` with the name `phantom` that you can initialize with
 /// `Default::default()`.
 ///
+/// ## PoV information
+///
+/// To implement the trait `StorageInfoTrait` for storages an additional attribute can be used
+/// `generate_storage_info`:
+/// ```nocompile
+/// decl_storage! { generate_storage_info
+/// 	trait Store for ...
+/// }
+/// ```
 #[proc_macro]
 pub fn decl_storage(input: TokenStream) -> TokenStream {
 	storage::decl_storage_impl(input)

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

@@ -15,7 +15,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::pallet::{Def, parse::helper::get_doc_literals};
+use crate::pallet::{Def, expand::merge_where_clauses, parse::helper::get_doc_literals};
 
 /// * Add derive trait on Pallet
 /// * Implement GetPalletVersion on Pallet
@@ -24,6 +24,7 @@ use crate::pallet::{Def, parse::helper::get_doc_literals};
 /// * declare Module type alias for construct_runtime
 /// * replace the first field type of `struct Pallet` with `PhantomData` if it is `_`
 /// * implementation of `PalletInfoAccess` information
+/// * implementation of `StorageInfoTrait` on Pallet
 pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 	let frame_support = &def.frame_support;
 	let frame_system = &def.frame_system;
@@ -33,6 +34,10 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 	let pallet_ident = &def.pallet_struct.pallet;
 	let config_where_clause = &def.config.where_clause;
 
+	let mut storages_where_clauses = vec![&def.config.where_clause];
+	storages_where_clauses.extend(def.storages.iter().map(|storage| &storage.where_clause));
+	let storages_where_clauses = merge_where_clauses(&storages_where_clauses);
+
 	let pallet_item = {
 		let pallet_module_items = &mut def.item.content.as_mut().expect("Checked by def").1;
 		let item = &mut pallet_module_items[def.pallet_struct.index];
@@ -97,6 +102,41 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 		)
 	};
 
+	let storage_info = if let Some(storage_info_span) = def.pallet_struct.generate_storage_info {
+		let storage_names = &def.storages.iter().map(|storage| &storage.ident).collect::<Vec<_>>();
+		let storage_cfg_attrs = &def.storages.iter()
+			.map(|storage| &storage.cfg_attrs)
+			.collect::<Vec<_>>();
+
+		quote::quote_spanned!(storage_info_span =>
+			impl<#type_impl_gen> #frame_support::traits::StorageInfoTrait
+				for #pallet_ident<#type_use_gen>
+				#storages_where_clauses
+			{
+				fn storage_info()
+					-> #frame_support::sp_std::vec::Vec<#frame_support::traits::StorageInfo>
+				{
+					let mut res = #frame_support::sp_std::vec![];
+
+					#(
+						#(#storage_cfg_attrs)*
+						{
+							let mut storage_info = <
+								#storage_names<#type_use_gen>
+								as #frame_support::traits::StorageInfoTrait
+							>::storage_info();
+							res.append(&mut storage_info);
+						}
+					)*
+
+					res
+				}
+			}
+		)
+	} else {
+		Default::default()
+	};
+
 	quote::quote_spanned!(def.pallet_struct.attr_span =>
 		#module_error_metadata
 
@@ -157,5 +197,7 @@ pub fn expand_pallet_struct(def: &mut Def) -> proc_macro2::TokenStream {
 						implemented by the runtime")
 			}
 		}
+
+		#storage_info
 	)
 }

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

@@ -24,6 +24,7 @@ mod keyword {
 	syn::custom_keyword!(pallet);
 	syn::custom_keyword!(Pallet);
 	syn::custom_keyword!(generate_store);
+	syn::custom_keyword!(generate_storage_info);
 	syn::custom_keyword!(Store);
 }
 
@@ -39,12 +40,30 @@ pub struct PalletStructDef {
 	pub store: Option<(syn::Visibility, keyword::Store)>,
 	/// The span of the pallet::pallet attribute.
 	pub attr_span: proc_macro2::Span,
+	/// Whether to specify the storages max encoded len when implementing `StorageInfoTrait`.
+	/// Contains the span of the attribute.
+	pub generate_storage_info: Option<proc_macro2::Span>,
 }
 
-/// Parse for `#[pallet::generate_store($vis trait Store)]`
-pub struct PalletStructAttr {
-	vis: syn::Visibility,
-	keyword: keyword::Store,
+/// Parse for one variant of:
+/// * `#[pallet::generate_store($vis trait Store)]`
+/// * `#[pallet::generate_storage_info]`
+pub enum PalletStructAttr {
+	GenerateStore {
+		span: proc_macro2::Span,
+		vis: syn::Visibility,
+		keyword: keyword::Store,
+	},
+	GenerateStorageInfoTrait(proc_macro2::Span),
+}
+
+impl PalletStructAttr {
+	fn span(&self) -> proc_macro2::Span {
+		match self {
+			Self::GenerateStore { span, .. } => *span,
+			Self::GenerateStorageInfoTrait(span) => *span,
+		}
+	}
 }
 
 impl syn::parse::Parse for PalletStructAttr {
@@ -54,14 +73,23 @@ impl syn::parse::Parse for PalletStructAttr {
 		syn::bracketed!(content in input);
 		content.parse::<keyword::pallet>()?;
 		content.parse::<syn::Token![::]>()?;
-		content.parse::<keyword::generate_store>()?;
-
-		let generate_content;
-		syn::parenthesized!(generate_content in content);
-		let vis = generate_content.parse::<syn::Visibility>()?;
-		generate_content.parse::<syn::Token![trait]>()?;
-		let keyword = generate_content.parse::<keyword::Store>()?;
-		Ok(Self { vis, keyword })
+
+		let lookahead = content.lookahead1();
+		if lookahead.peek(keyword::generate_store) {
+			let span = content.parse::<keyword::generate_store>()?.span();
+
+			let generate_content;
+			syn::parenthesized!(generate_content in content);
+			let vis = generate_content.parse::<syn::Visibility>()?;
+			generate_content.parse::<syn::Token![trait]>()?;
+			let keyword = generate_content.parse::<keyword::Store>()?;
+			Ok(Self::GenerateStore { vis, keyword, span })
+		} else if lookahead.peek(keyword::generate_storage_info) {
+			let span = content.parse::<keyword::generate_storage_info>()?.span();
+			Ok(Self::GenerateStorageInfoTrait(span))
+		} else {
+			Err(lookahead.error())
+		}
 	}
 }
 
@@ -78,12 +106,24 @@ impl PalletStructDef {
 			return Err(syn::Error::new(item.span(), msg));
 		};
 
-		let mut event_attrs: Vec<PalletStructAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
-		if event_attrs.len() > 1 {
-			let msg = "Invalid pallet::pallet, multiple argument pallet::generate_store found";
-			return Err(syn::Error::new(event_attrs[1].keyword.span(), msg));
+		let mut store = None;
+		let mut generate_storage_info = None;
+
+		let struct_attrs: Vec<PalletStructAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
+		for attr in struct_attrs {
+			match attr {
+				PalletStructAttr::GenerateStore { vis, keyword, .. } if store.is_none() => {
+					store = Some((vis, keyword));
+				},
+				PalletStructAttr::GenerateStorageInfoTrait(span) if generate_storage_info.is_none() => {
+					generate_storage_info = Some(span);
+				},
+				attr => {
+					let msg = "Unexpected duplicated attribute";
+					return Err(syn::Error::new(attr.span(), msg));
+				},
+			}
 		}
-		let store = event_attrs.pop().map(|attr| (attr.vis, attr.keyword));
 
 		let pallet = syn::parse2::<keyword::Pallet>(item.ident.to_token_stream())?;
 
@@ -100,6 +140,6 @@ impl PalletStructDef {
 		let mut instances = vec![];
 		instances.push(helper::check_type_def_gen_no_bounds(&item.generics, item.ident.span())?);
 
-		Ok(Self { index, instances, pallet, store, attr_span })
+		Ok(Self { index, instances, pallet, store, attr_span, generate_storage_info })
 	}
 }

frame/support/procedural/src/storage/mod.rs

@@ -18,6 +18,7 @@
 //! `decl_storage` input definition and expansion.
 
 mod storage_struct;
+mod storage_info;
 mod parse;
 mod store_trait;
 mod getters;
@@ -35,6 +36,8 @@ use frame_support_procedural_tools::{
 
 /// All information contained in input of decl_storage
 pub struct DeclStorageDef {
+	/// Whether to generate the storage info
+	generate_storage_info: bool,
 	/// Name of the module used to import hidden imports.
 	hidden_crate: Option<syn::Ident>,
 	/// Visibility of store trait.
@@ -69,6 +72,8 @@ impl syn::parse::Parse for DeclStorageDef {
 
 /// Extended version of `DeclStorageDef` with useful precomputed value.
 pub struct DeclStorageDefExt {
+	/// Whether to generate the storage info
+	generate_storage_info: bool,
 	/// Name of the module used to import hidden imports.
 	hidden_crate: proc_macro2::TokenStream,
 	/// Hidden imports used by the module.
@@ -154,6 +159,7 @@ impl From<DeclStorageDef> for DeclStorageDefExt {
 		Self {
 			hidden_crate,
 			hidden_imports,
+			generate_storage_info: def.generate_storage_info,
 			visibility: def.visibility,
 			store_trait: def.store_trait,
 			module_name: def.module_name,
@@ -193,6 +199,8 @@ pub struct StorageLineDef {
 	getter: Option<syn::Ident>,
 	/// The name of the field to be used in genesis config if any.
 	config: Option<syn::Ident>,
+	/// The given max values with `max_values` attribute, or a none if not specified.
+	max_values: Option<syn::Expr>,
 	/// The build function of the storage if any.
 	build: Option<syn::Expr>,
 	/// Default value of genesis config field and also for storage when no value available.
@@ -210,6 +218,8 @@ pub struct StorageLineDefExt {
 	getter: Option<syn::Ident>,
 	/// The name of the field to be used in genesis config if any.
 	config: Option<syn::Ident>,
+	/// The given max values with `max_values` attribute, or a none if not specified.
+	max_values: Option<syn::Expr>,
 	/// The build function of the storage if any.
 	build: Option<syn::Expr>,
 	/// Default value of genesis config field and also for storage when no value available.
@@ -333,6 +343,7 @@ impl StorageLineDefExt {
 			name: storage_def.name,
 			getter: storage_def.getter,
 			config: storage_def.config,
+			max_values: storage_def.max_values,
 			build: storage_def.build,
 			default_value: storage_def.default_value,
 			storage_type: storage_def.storage_type,
@@ -469,6 +480,7 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr
 	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);
+	let storage_info = storage_info::impl_storage_info(&def_ext);
 
 	quote!(
 		use #scrate::{
@@ -489,5 +501,6 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr
 		#instance_trait
 		#genesis_config
 		#storage_struct
+		#storage_info
 	).into()
 }

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

@@ -21,10 +21,12 @@ use frame_support_procedural_tools::{ToTokens, Parse, syn_ext as ext};
 use syn::{Ident, Token, spanned::Spanned};
 
 mod keyword {
+	syn::custom_keyword!(generate_storage_info);
 	syn::custom_keyword!(hiddencrate);
 	syn::custom_keyword!(add_extra_genesis);
 	syn::custom_keyword!(extra_genesis_skip_phantom_data_field);
 	syn::custom_keyword!(config);
+	syn::custom_keyword!(max_values);
 	syn::custom_keyword!(build);
 	syn::custom_keyword!(get);
 	syn::custom_keyword!(map);
@@ -73,6 +75,7 @@ macro_rules! impl_parse_for_opt {
 /// Parsing usage only
 #[derive(Parse, ToTokens, Debug)]
 struct StorageDefinition {
+	pub generate_storage_info: Opt<GenerateStorageInfo>,
 	pub hidden_crate: Opt<SpecificHiddenCrate>,
 	pub visibility: syn::Visibility,
 	pub trait_token: Token![trait],
@@ -97,6 +100,12 @@ struct StorageDefinition {
 	pub extra_genesis: Opt<AddExtraGenesis>,
 }
 
+#[derive(Parse, ToTokens, Debug)]
+struct GenerateStorageInfo {
+	pub keyword: keyword::generate_storage_info,
+}
+impl_parse_for_opt!(GenerateStorageInfo => keyword::generate_storage_info);
+
 #[derive(Parse, ToTokens, Debug)]
 struct SpecificHiddenCrate {
 	pub keyword: keyword::hiddencrate,
@@ -160,6 +169,7 @@ struct DeclStorageLine {
 	pub name: Ident,
 	pub getter: Opt<DeclStorageGetter>,
 	pub config: Opt<DeclStorageConfig>,
+	pub max_values: Opt<DeclStorageMaxValues>,
 	pub build: Opt<DeclStorageBuild>,
 	pub coldot_token: Token![:],
 	pub storage_type: DeclStorageType,
@@ -188,6 +198,13 @@ struct DeclStorageConfig {
 
 impl_parse_for_opt!(DeclStorageConfig => keyword::config);
 
+#[derive(Parse, ToTokens, Debug)]
+struct DeclStorageMaxValues {
+	pub max_values_keyword: keyword::max_values,
+	pub expr: ext::Parens<syn::Expr>,
+}
+impl_parse_for_opt!(DeclStorageMaxValues => keyword::max_values);
+
 #[derive(Parse, ToTokens, Debug)]
 struct DeclStorageBuild {
 	pub build_keyword: keyword::build,
@@ -437,6 +454,7 @@ pub fn parse(input: syn::parse::ParseStream) -> syn::Result<super::DeclStorageDe
 	let storage_lines = parse_storage_line_defs(def.content.content.inner.into_iter())?;
 
 	Ok(super::DeclStorageDef {
+		generate_storage_info: def.generate_storage_info.inner.is_some(),
 		hidden_crate: def.hidden_crate.inner.map(|i| i.ident.content),
 		visibility: def.visibility,
 		module_name: def.module_ident,
@@ -490,6 +508,21 @@ fn parse_storage_line_defs(
 			})?;
 		}
 
+		let max_values = match &line.storage_type {
+			DeclStorageType::Map(_) | DeclStorageType::DoubleMap(_) | DeclStorageType::NMap(_) => {
+				line.max_values.inner.map(|i| i.expr.content)
+			},
+			DeclStorageType::Simple(_) => {
+				if let Some(max_values) = line.max_values.inner {
+					let msg = "unexpected max_values attribute for storage value.";
+					let span = max_values.max_values_keyword.span();
+					return Err(syn::Error::new(span, msg));
+				} else {
+					Some(syn::parse_quote!(1u32))
+				}
+			},
+		};
+
 		let span = line.storage_type.span();
 		let no_hasher_error = || syn::Error::new(
 			span,
@@ -534,6 +567,7 @@ fn parse_storage_line_defs(
 			name: line.name,
 			getter,
 			config,
+			max_values,
 			build: line.build.inner.map(|o| o.expr.content),
 			default_value: line.default_value.inner.map(|o| o.expr),
 			storage_type,