Allow renaming storage item prefixes (#9016)

* Implement parsing for #[pallet::storage_name] on storage items

* Rename storage prefix when a #[pallet::storage_name] is supplied

* Fix test_storage_info

* Rename storage_name to storage_prefix

* Check for duplicates when renaming storage prefixes

* Allow only string literals for storage_prefix renames

* Use proper spans for attribute errors

* Check for valid identifiers when parsing storage prefix renames

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

@@ -15,22 +15,48 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::pallet::Def;
+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;
 
 /// Generate the prefix_ident related the the storage.
 /// prefix_ident is used for the prefix struct to be given to storage as first generic param.
-fn prefix_ident(storage_ident: &syn::Ident) -> syn::Ident {
+fn prefix_ident(storage: &StorageDef) -> syn::Ident {
+	let storage_ident = &storage.ident;
 	syn::Ident::new(&format!("_GeneratedPrefixForStorage{}", storage_ident), storage_ident.span())
 }
 
+/// 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<()> {
+	let prefix = storage_def.prefix();
+
+	if !set.insert(prefix.clone()) {
+		let err = syn::Error::new(
+			storage_def.prefix_span(),
+			format!("Duplicate storage prefixes found for `{}`", prefix),
+		);
+		return Err(err);
+	}
+
+	Ok(())
+}
+
 /// * if generics are unnamed: replace the first generic `_` by the generated prefix structure
 /// * if generics are named: reorder the generic, remove their name, and add the missing ones.
 /// * Add `#[allow(type_alias_bounds)]`
-pub fn process_generics(def: &mut Def) {
+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 {
@@ -50,7 +76,7 @@ pub fn process_generics(def: &mut Def) {
 			_ => unreachable!("Checked by def"),
 		};
 
-		let prefix_ident = prefix_ident(&storage_def.ident);
+		let prefix_ident = prefix_ident(&storage_def);
 		let type_use_gen = if def.config.has_instance {
 			quote::quote_spanned!(storage_def.attr_span => T, I)
 		} else {
@@ -116,6 +142,8 @@ pub fn process_generics(def: &mut Def) {
 			args.args[0] = syn::parse_quote!( #prefix_ident<#type_use_gen> );
 		}
 	}
+
+	Ok(())
 }
 
 /// * generate StoragePrefix structs (e.g. for a storage `MyStorage` a struct with the name
@@ -125,7 +153,9 @@ pub fn process_generics(def: &mut Def) {
 /// * Add `#[allow(type_alias_bounds)]` on storages type alias
 /// * generate metadatas
 pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
-	process_generics(def);
+	if let Err(e) = process_generics(def) {
+		return e.into_compile_error().into();
+	}
 
 	let frame_support = &def.frame_support;
 	let frame_system = &def.frame_system;
@@ -344,9 +374,9 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
 	let prefix_structs = def.storages.iter().map(|storage_def| {
 		let type_impl_gen = &def.type_impl_generics(storage_def.attr_span);
 		let type_use_gen = &def.type_use_generics(storage_def.attr_span);
-		let prefix_struct_ident = prefix_ident(&storage_def.ident);
+		let prefix_struct_ident = prefix_ident(&storage_def);
 		let prefix_struct_vis = &storage_def.vis;
-		let prefix_struct_const = storage_def.ident.to_string();
+		let prefix_struct_const = storage_def.prefix();
 		let config_where_clause = &def.config.where_clause;
 
 		let cfg_attrs = &storage_def.cfg_attrs;

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

@@ -25,28 +25,60 @@ mod keyword {
 	syn::custom_keyword!(Error);
 	syn::custom_keyword!(pallet);
 	syn::custom_keyword!(getter);
+	syn::custom_keyword!(storage_prefix);
 	syn::custom_keyword!(OptionQuery);
 	syn::custom_keyword!(ValueQuery);
 }
 
-/// Parse for `#[pallet::getter(fn dummy)]`
-pub struct PalletStorageAttr {
-	getter: syn::Ident,
+/// Parse for one of the following:
+/// * `#[pallet::getter(fn dummy)]`
+/// * `#[pallet::storage_prefix = "CustomName"]`
+pub enum PalletStorageAttr {
+	Getter(syn::Ident, proc_macro2::Span),
+	StorageName(syn::LitStr, proc_macro2::Span),
+}
+
+impl PalletStorageAttr {
+	fn attr_span(&self) -> proc_macro2::Span {
+		match self {
+			Self::Getter(_, span) | Self::StorageName(_, span) => *span,
+		}
+	}
 }
 
 impl syn::parse::Parse for PalletStorageAttr {
 	fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
 		input.parse::<syn::Token![#]>()?;
+		let attr_span = input.span();
 		let content;
 		syn::bracketed!(content in input);
 		content.parse::<keyword::pallet>()?;
 		content.parse::<syn::Token![::]>()?;
-		content.parse::<keyword::getter>()?;
 
-		let generate_content;
-		syn::parenthesized!(generate_content in content);
-		generate_content.parse::<syn::Token![fn]>()?;
-		Ok(Self { getter: generate_content.parse::<syn::Ident>()? })
+		let lookahead = content.lookahead1();
+		if lookahead.peek(keyword::getter) {
+			content.parse::<keyword::getter>()?;
+
+			let generate_content;
+			syn::parenthesized!(generate_content in content);
+			generate_content.parse::<syn::Token![fn]>()?;
+			Ok(Self::Getter(generate_content.parse::<syn::Ident>()?, attr_span))
+		} else if lookahead.peek(keyword::storage_prefix) {
+			content.parse::<keyword::storage_prefix>()?;
+			content.parse::<syn::Token![=]>()?;
+
+			let renamed_prefix = content.parse::<syn::LitStr>()?;
+			// Ensure the renamed prefix is a proper Rust identifier
+			syn::parse_str::<syn::Ident>(&renamed_prefix.value())
+				.map_err(|_| {
+					let msg = format!("`{}` is not a valid identifier", renamed_prefix.value());
+					syn::Error::new(renamed_prefix.span(), msg)
+				})?;
+
+			Ok(Self::StorageName(renamed_prefix, attr_span))
+		} else {
+			Err(lookahead.error())
+		}
 	}
 }
 
@@ -89,6 +121,8 @@ pub struct StorageDef {
 	pub instances: Vec<helper::InstanceUsage>,
 	/// Optional getter to generate. If some then query_kind is ensured to be some as well.
 	pub getter: Option<syn::Ident>,
+	/// Optional expression that evaluates to a type that can be used as StoragePrefix instead of ident.
+	pub rename_as: Option<syn::LitStr>,
 	/// Whereas the querytype of the storage is OptionQuery or ValueQuery.
 	/// Note that this is best effort as it can't be determined when QueryKind is generic, and
 	/// result can be false if user do some unexpected type alias.
@@ -105,7 +139,6 @@ pub struct StorageDef {
 	pub named_generics: Option<StorageGenerics>,
 }
 
-
 /// The parsed generic from the
 #[derive(Clone)]
 pub enum StorageGenerics {
@@ -541,6 +574,25 @@ fn extract_key(ty: &syn::Type) -> syn::Result<syn::Type> {
 }
 
 impl StorageDef {
+	/// Return the storage prefix for this storage item
+	pub fn prefix(&self) -> String {
+		self
+			.rename_as
+			.as_ref()
+			.map(syn::LitStr::value)
+			.unwrap_or(self.ident.to_string())
+	}
+
+	/// Return either the span of the ident or the span of the literal in the
+	/// #[storage_prefix] attribute
+	pub fn prefix_span(&self) -> proc_macro2::Span {
+		self
+			.rename_as
+			.as_ref()
+			.map(syn::LitStr::span)
+			.unwrap_or(self.ident.span())
+	}
+
 	pub fn try_from(
 		attr_span: proc_macro2::Span,
 		index: usize,
@@ -552,12 +604,30 @@ impl StorageDef {
 			return Err(syn::Error::new(item.span(), "Invalid pallet::storage, expect item type."));
 		};
 
-		let mut attrs: Vec<PalletStorageAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
-		if attrs.len() > 1 {
+		let attrs: Vec<PalletStorageAttr> = helper::take_item_pallet_attrs(&mut item.attrs)?;
+		let (mut getters, mut names) = attrs
+			.into_iter()
+			.partition::<Vec<_>, _>(|attr| matches!(attr, PalletStorageAttr::Getter(..)));
+		if getters.len() > 1 {
 			let msg = "Invalid pallet::storage, multiple argument pallet::getter found";
-			return Err(syn::Error::new(attrs[1].getter.span(), msg));
+			return Err(syn::Error::new(getters[1].attr_span(), msg));
 		}
-		let getter = attrs.pop().map(|attr| attr.getter);
+		if names.len() > 1 {
+			let msg = "Invalid pallet::storage, multiple argument pallet::storage_prefix found";
+			return Err(syn::Error::new(names[1].attr_span(), msg));
+		}
+		let getter = getters.pop().map(|attr| {
+			match attr {
+				PalletStorageAttr::Getter(ident, _) => ident,
+				_ => unreachable!(),
+			}
+		});
+		let rename_as = names.pop().map(|attr| {
+			match attr {
+				PalletStorageAttr::StorageName(lit, _) => lit,
+				_ => unreachable!(),
+			}
+		});
 
 		let cfg_attrs = helper::get_item_cfg_attrs(&item.attrs);
 
@@ -609,6 +679,7 @@ impl StorageDef {
 			metadata,
 			docs,
 			getter,
+			rename_as,
 			query_kind,
 			where_clause,
 			cfg_attrs,