Support pallet::storage conditional compilation (#8324)

* Support pallet::storage conditional compilation.

* Add docs for cfg attributes.

* Keep strong types for get cfg attrs return.

* Update frame/support/procedural/src/pallet/parse/helper.rs

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

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/.gitlab-ci.yml

@@ -277,6 +277,7 @@ test-linux-stable:                 &test-linux
   script:
     # this job runs all tests in former runtime-benchmarks, frame-staking and wasmtime tests
     - time cargo test --workspace --locked --release --verbose --features runtime-benchmarks --manifest-path bin/node/cli/Cargo.toml
+    - time cargo test -p frame-support-test --features=conditional-storage --manifest-path frame/support/test/Cargo.toml
     - SUBSTRATE_TEST_TIMEOUT=1 time cargo test -p substrate-test-utils --release --verbose --locked -- --ignored timeout
     - sccache -s
 

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

@@ -78,6 +78,8 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 			let gen = &def.type_use_generics(storage.attr_span);
 			let full_ident = quote::quote_spanned!(storage.attr_span => #ident<#gen> );
 
+			let cfg_attrs = &storage.cfg_attrs;
+
 			let metadata_trait = match &storage.metadata {
 				Metadata::Value { .. } => quote::quote_spanned!(storage.attr_span =>
 					#frame_support::storage::types::StorageValueMetadata
@@ -128,7 +130,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 			};
 
 			quote::quote_spanned!(storage.attr_span =>
-				#frame_support::metadata::StorageEntryMetadata {
+				#(#cfg_attrs)* #frame_support::metadata::StorageEntryMetadata {
 					name: #frame_support::metadata::DecodeDifferent::Encode(
 						<#full_ident as #metadata_trait>::NAME
 					),
@@ -159,6 +161,8 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 			let type_use_gen = &def.type_use_generics(storage.attr_span);
 			let full_ident = quote::quote_spanned!(storage.attr_span => #ident<#gen> );
 
+			let cfg_attrs = &storage.cfg_attrs;
+
 			match &storage.metadata {
 				Metadata::Value { value } => {
 					let query = match storage.query_kind.as_ref().expect("Checked by def") {
@@ -168,6 +172,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						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() -> #query {
@@ -186,6 +191,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						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
@@ -206,6 +212,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 						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<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> #query where
@@ -233,10 +240,14 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 		let prefix_struct_const = storage_def.ident.to_string();
 		let config_where_clause = &def.config.where_clause;
 
+		let cfg_attrs = &storage_def.cfg_attrs;
+
 		quote::quote_spanned!(storage_def.attr_span =>
+			#(#cfg_attrs)*
 			#prefix_struct_vis struct #prefix_struct_ident<#type_use_gen>(
 				core::marker::PhantomData<(#type_use_gen,)>
 			);
+			#(#cfg_attrs)*
 			impl<#type_impl_gen> #frame_support::traits::StorageInstance
 				for #prefix_struct_ident<#type_use_gen>
 				#config_where_clause

substrate/frame/support/procedural/src/pallet/expand/store_trait.rs

@@ -37,10 +37,12 @@ pub fn expand_store_trait(def: &mut Def) -> proc_macro2::TokenStream {
 	let completed_where_clause = super::merge_where_clauses(&where_clauses);
 
 	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!(trait_store.span() =>
 		#trait_vis trait #trait_store {
 			#(
+				#(#storage_cfg_attrs)*
 				type #storage_names;
 			)*
 		}
@@ -48,6 +50,7 @@ pub fn expand_store_trait(def: &mut Def) -> proc_macro2::TokenStream {
 			#completed_where_clause
 		{
 			#(
+				#(#storage_cfg_attrs)*
 				type #storage_names = #storage_names<#type_use_gen>;
 			)*
 		}

substrate/frame/support/procedural/src/pallet/parse/call.rs

@@ -172,7 +172,7 @@ impl CallDef {
 				}
 
 				let mut call_var_attrs: Vec<FunctionAttr> =
-					helper::take_item_attrs(&mut method.attrs)?;
+					helper::take_item_pallet_attrs(&mut method.attrs)?;
 
 				if call_var_attrs.len() != 1 {
 					let msg = if call_var_attrs.is_empty() {
@@ -193,7 +193,7 @@ impl CallDef {
 					};
 
 					let arg_attrs: Vec<ArgAttrIsCompact> =
-						helper::take_item_attrs(&mut arg.attrs)?;
+						helper::take_item_pallet_attrs(&mut arg.attrs)?;
 
 					if arg_attrs.len() > 1 {
 						let msg = "Invalid pallet::call, argument has too many attributes";

substrate/frame/support/procedural/src/pallet/parse/config.rs

@@ -309,7 +309,7 @@ impl ConfigDef {
 				|| check_event_type(frame_system, trait_item, has_instance)?;
 
 			// Parse for constant
-			let type_attrs_const: Vec<TypeAttrConst> = helper::take_item_attrs(trait_item)?;
+			let type_attrs_const: Vec<TypeAttrConst> = helper::take_item_pallet_attrs(trait_item)?;
 
 			if type_attrs_const.len() > 1 {
 				let msg = "Invalid attribute in pallet::config, only one attribute is expected";
@@ -339,7 +339,7 @@ impl ConfigDef {
 			}
 		}
 
-		let attr: Option<DisableFrameSystemSupertraitCheck> = helper::take_first_item_attr(
+		let attr: Option<DisableFrameSystemSupertraitCheck> = helper::take_first_item_pallet_attr(
 			&mut item.attrs
 		)?;
 

substrate/frame/support/procedural/src/pallet/parse/event.rs

@@ -163,7 +163,7 @@ impl EventDef {
 			return Err(syn::Error::new(item.span(), "Invalid pallet::event, expected item enum"))
 		};
 
-		let event_attrs: Vec<PalletEventAttr> = helper::take_item_attrs(&mut item.attrs)?;
+		let event_attrs: Vec<PalletEventAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
 		let attr_info = PalletEventAttrInfo::from_attrs(event_attrs)?;
 		let metadata = attr_info.metadata.unwrap_or_else(Vec::new);
 		let deposit_event = attr_info.deposit_event;

substrate/frame/support/procedural/src/pallet/parse/helper.rs

@@ -47,7 +47,7 @@ pub trait MutItemAttrs {
 }
 
 /// Take the first pallet attribute (e.g. attribute like `#[pallet..]`) and decode it to `Attr`
-pub fn take_first_item_attr<Attr>(item: &mut impl MutItemAttrs) -> syn::Result<Option<Attr>> where
+pub fn take_first_item_pallet_attr<Attr>(item: &mut impl MutItemAttrs) -> syn::Result<Option<Attr>> where
 	Attr: syn::parse::Parse,
 {
 	let attrs = if let Some(attrs) = item.mut_item_attrs() {
@@ -69,18 +69,29 @@ pub fn take_first_item_attr<Attr>(item: &mut impl MutItemAttrs) -> syn::Result<O
 }
 
 /// Take all the pallet attributes (e.g. attribute like `#[pallet..]`) and decode them to `Attr`
-pub fn take_item_attrs<Attr>(item: &mut impl MutItemAttrs) -> syn::Result<Vec<Attr>> where
+pub fn take_item_pallet_attrs<Attr>(item: &mut impl MutItemAttrs) -> syn::Result<Vec<Attr>> where
 	Attr: syn::parse::Parse,
 {
 	let mut pallet_attrs = Vec::new();
 
-	while let Some(attr) = take_first_item_attr(item)? {
+	while let Some(attr) = take_first_item_pallet_attr(item)? {
 		pallet_attrs.push(attr)
 	}
 
 	Ok(pallet_attrs)
 }
 
+/// Get all the cfg attributes (e.g. attribute like `#[cfg..]`) and decode them to `Attr`
+pub fn get_item_cfg_attrs(attrs: &[syn::Attribute]) -> Vec<syn::Attribute> {
+	attrs.iter().filter_map(|attr|  {
+		if attr.path.segments.first().map_or(false, |segment| segment.ident == "cfg") {
+			Some(attr.clone())
+		} else {
+			None
+		}
+	}).collect::<Vec<_>>()
+}
+
 impl MutItemAttrs for syn::Item {
 	fn mut_item_attrs(&mut self) -> Option<&mut Vec<syn::Attribute>> {
 		match self {

substrate/frame/support/procedural/src/pallet/parse/mod.rs

@@ -89,7 +89,7 @@ impl Def {
 		let mut type_values = vec![];
 
 		for (index, item) in items.iter_mut().enumerate() {
-			let pallet_attr: Option<PalletAttr> = helper::take_first_item_attr(item)?;
+			let pallet_attr: Option<PalletAttr> = helper::take_first_item_pallet_attr(item)?;
 
 			match pallet_attr {
 				Some(PalletAttr::Config(span)) if config.is_none() =>

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

@@ -78,7 +78,7 @@ impl PalletStructDef {
 			return Err(syn::Error::new(item.span(), msg));
 		};
 
-		let mut event_attrs: Vec<PalletStructAttr> = helper::take_item_attrs(&mut item.attrs)?;
+		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));

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

@@ -91,6 +91,8 @@ pub struct StorageDef {
 	pub where_clause: Option<syn::WhereClause>,
 	/// The span of the pallet::storage attribute.
 	pub attr_span: proc_macro2::Span,
+	/// The `cfg` attributes.
+	pub cfg_attrs: Vec<syn::Attribute>,
 }
 
 /// In `Foo<A, B, C>` retrieve the argument at given position, i.e. A is argument at position 0.
@@ -125,13 +127,15 @@ impl StorageDef {
 			return Err(syn::Error::new(item.span(), "Invalid pallet::storage, expected item type"));
 		};
 
-		let mut attrs: Vec<PalletStorageAttr> = helper::take_item_attrs(&mut item.attrs)?;
+		let mut attrs: Vec<PalletStorageAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
 		if attrs.len() > 1 {
 			let msg = "Invalid pallet::storage, multiple argument pallet::getter found";
 			return Err(syn::Error::new(attrs[1].getter.span(), msg));
 		}
 		let getter = attrs.pop().map(|attr| attr.getter);
 
+		let cfg_attrs = helper::get_item_cfg_attrs(&item.attrs);
+
 		let mut instances = vec![];
 		instances.push(helper::check_type_def_gen(&item.generics, item.ident.span())?);
 
@@ -223,6 +227,7 @@ impl StorageDef {
 			getter,
 			query_kind,
 			where_clause,
+			cfg_attrs,
 		})
 	}
 }

substrate/frame/support/src/lib.rs

@@ -1418,6 +1418,18 @@ pub mod pallet_prelude {
 /// pub(super) type MyStorage<T> = StorageMap<_, Blake2_128Concat, u32, u32>;
 /// ```
 ///
+/// The optional attributes `#[cfg(..)]` allow conditional compilation for the storage.
+///
+/// E.g:
+/// ```ignore
+/// #[cfg(feature = "my-feature")]
+/// #[pallet::storage]
+/// pub(super) type MyStorage<T> = StorageValue<_, u32>;
+/// ```
+///
+/// All the `cfg` attributes are automatically copied to the items generated for the storage, i.e. the
+/// getter, storage prefix, and the metadata element etc.
+///
 /// NOTE: If the `QueryKind` generic parameter is still generic at this stage or is using some type
 /// alias then the generation of the getter might fail. In this case the getter can be implemented
 /// manually.

substrate/frame/support/test/Cargo.toml

@@ -42,3 +42,4 @@ std = [
 	"sp-state-machine",
 ]
 try-runtime = ["frame-support/try-runtime"]
+conditional-storage = []

substrate/frame/support/test/tests/pallet.rs

@@ -217,6 +217,28 @@ pub mod pallet {
 	#[pallet::storage]
 	pub type DoubleMap2<T> = StorageDoubleMap<_, Twox64Concat, u16, Blake2_128Concat, u32, u64>;
 
+	#[pallet::storage]
+	#[pallet::getter(fn conditional_value)]
+	#[cfg(feature = "conditional-storage")]
+	pub type ConditionalValue<T> = StorageValue<_, u32>;
+
+	#[cfg(feature = "conditional-storage")]
+	#[pallet::storage]
+	#[pallet::getter(fn conditional_map)]
+	pub type ConditionalMap<T> = StorageMap<_, Twox64Concat, u16, u32>;
+
+	#[cfg(feature = "conditional-storage")]
+	#[pallet::storage]
+	#[pallet::getter(fn conditional_double_map)]
+	pub type ConditionalDoubleMap<T> = StorageDoubleMap<
+		_,
+		Blake2_128Concat,
+		u8,
+		Twox64Concat,
+		u16,
+		u32,
+	>;
+
 	#[pallet::genesis_config]
 	#[derive(Default)]
 	pub struct GenesisConfig {
@@ -522,6 +544,13 @@ fn storage_expand() {
 		k.extend(2u32.using_encoded(blake2_128_concat));
 		assert_eq!(unhashed::get::<u64>(&k), Some(3u64));
 		assert_eq!(&k[..32], &<pallet::DoubleMap2<Runtime>>::final_prefix());
+
+		#[cfg(feature = "conditional-storage")]
+		{
+			pallet::ConditionalValue::<Runtime>::put(1);
+			pallet::ConditionalMap::<Runtime>::insert(1, 2);
+			pallet::ConditionalDoubleMap::<Runtime>::insert(1, 2, 3);
+		}
 	})
 }
 
@@ -646,6 +675,38 @@ fn metadata() {
 					default: DecodeDifferent::Decoded(vec![0]),
 					documentation: DecodeDifferent::Decoded(vec![]),
 				},
+				#[cfg(feature = "conditional-storage")] StorageEntryMetadata {
+					name: DecodeDifferent::Decoded("ConditionalValue".to_string()),
+					modifier: StorageEntryModifier::Optional,
+					ty: StorageEntryType::Plain(DecodeDifferent::Decoded("u32".to_string())),
+					default: DecodeDifferent::Decoded(vec![0]),
+					documentation: DecodeDifferent::Decoded(vec![]),
+				},
+				#[cfg(feature = "conditional-storage")] StorageEntryMetadata {
+					name: DecodeDifferent::Decoded("ConditionalMap".to_string()),
+					modifier: StorageEntryModifier::Optional,
+					ty: StorageEntryType::Map {
+						key: DecodeDifferent::Decoded("u16".to_string()),
+						value: DecodeDifferent::Decoded("u32".to_string()),
+						hasher: StorageHasher::Twox64Concat,
+						unused: false,
+					},
+					default: DecodeDifferent::Decoded(vec![0]),
+					documentation: DecodeDifferent::Decoded(vec![]),
+				},
+				#[cfg(feature = "conditional-storage")] StorageEntryMetadata {
+					name: DecodeDifferent::Decoded("ConditionalDoubleMap".to_string()),
+					modifier: StorageEntryModifier::Optional,
+					ty: StorageEntryType::DoubleMap {
+						value: DecodeDifferent::Decoded("u32".to_string()),
+						key1: DecodeDifferent::Decoded("u8".to_string()),
+						key2: DecodeDifferent::Decoded("u16".to_string()),
+						hasher: StorageHasher::Blake2_128Concat,
+						key2_hasher: StorageHasher::Twox64Concat,
+					},
+					default: DecodeDifferent::Decoded(vec![0]),
+					documentation: DecodeDifferent::Decoded(vec![]),
+				},
 			]),
 		})),
 		calls: Some(DecodeDifferent::Decoded(vec![