use crate::{counter_prefix, pallet::parse::helper};
use frame_support_procedural_tools::generate_access_from_frame_or_crate;
use proc_macro2::{Span, TokenStream};
use quote::{quote, ToTokens};
use syn::{
parenthesized,
parse::{Parse, ParseStream},
punctuated::Punctuated,
spanned::Spanned,
token,
visit::Visit,
Attribute, Error, Ident, Result, Token, Type, TypeParam, Visibility, WhereClause,
};
trait TypeExt {
fn get_ident(&self) -> Option<&Ident>;
fn contains_ident(&self, ident: &Ident) -> bool;
}
impl TypeExt for Type {
fn get_ident(&self) -> Option<&Ident> {
match self {
Type::Path(p) => match &p.qself {
Some(qself) => qself.ty.get_ident(),
None => p.path.get_ident(),
},
_ => None,
}
}
fn contains_ident(&self, ident: &Ident) -> bool {
struct ContainsIdent<'a> {
ident: &'a Ident,
found: bool,
}
impl<'a, 'ast> Visit<'ast> for ContainsIdent<'a> {
fn visit_ident(&mut self, i: &'ast Ident) {
if i == self.ident {
self.found = true;
}
}
}
let mut visitor = ContainsIdent { ident, found: false };
syn::visit::visit_type(&mut visitor, self);
visitor.found
}
}
struct SimpleGenerics {
lt_token: Token![<],
params: Punctuated<TypeParam, token::Comma>,
gt_token: Token![>],
}
impl SimpleGenerics {
fn type_generics(&self) -> impl Iterator<Item = &Ident> {
self.params.iter().map(|p| &p.ident)
}
fn impl_generics(&self) -> impl Iterator<Item = &TypeParam> {
self.params.iter()
}
}
impl Parse for SimpleGenerics {
fn parse(input: ParseStream<'_>) -> Result<Self> {
Ok(Self {
lt_token: input.parse()?,
params: Punctuated::parse_separated_nonempty(input)?,
gt_token: input.parse()?,
})
}
}
impl ToTokens for SimpleGenerics {
fn to_tokens(&self, tokens: &mut TokenStream) {
self.lt_token.to_tokens(tokens);
self.params.to_tokens(tokens);
self.gt_token.to_tokens(tokens);
}
}
mod storage_types {
syn::custom_keyword!(StorageValue);
syn::custom_keyword!(StorageMap);
syn::custom_keyword!(CountedStorageMap);
syn::custom_keyword!(StorageDoubleMap);
syn::custom_keyword!(StorageNMap);
}
mod prefix_types {
syn::custom_keyword!(verbatim);
syn::custom_keyword!(pallet_name);
syn::custom_keyword!(dynamic);
}
enum StorageType {
Value {
_kw: storage_types::StorageValue,
_lt_token: Token![<],
prefix: Type,
_value_comma: Token![,],
value_ty: Type,
query_type: Option<(Token![,], Type)>,
_trailing_comma: Option<Token![,]>,
_gt_token: Token![>],
},
Map {
_kw: storage_types::StorageMap,
_lt_token: Token![<],
prefix: Type,
_hasher_comma: Token![,],
hasher_ty: Type,
_key_comma: Token![,],
key_ty: Type,
_value_comma: Token![,],
value_ty: Type,
query_type: Option<(Token![,], Type)>,
_trailing_comma: Option<Token![,]>,
_gt_token: Token![>],
},
CountedMap {
_kw: storage_types::CountedStorageMap,
_lt_token: Token![<],
prefix: Type,
_hasher_comma: Token![,],
hasher_ty: Type,
_key_comma: Token![,],
key_ty: Type,
_value_comma: Token![,],
value_ty: Type,
query_type: Option<(Token![,], Type)>,
_trailing_comma: Option<Token![,]>,
_gt_token: Token![>],
},
DoubleMap {
_kw: storage_types::StorageDoubleMap,
_lt_token: Token![<],
prefix: Type,
_hasher1_comma: Token![,],
hasher1_ty: Type,
_key1_comma: Token![,],
key1_ty: Type,
_hasher2_comma: Token![,],
hasher2_ty: Type,
_key2_comma: Token![,],
key2_ty: Type,
_value_comma: Token![,],
value_ty: Type,
query_type: Option<(Token![,], Type)>,
_trailing_comma: Option<Token![,]>,
_gt_token: Token![>],
},
NMap {
_kw: storage_types::StorageNMap,
_lt_token: Token![<],
prefix: Type,
_paren_comma: Token![,],
_paren_token: token::Paren,
key_types: Punctuated<Type, Token![,]>,
_value_comma: Token![,],
value_ty: Type,
query_type: Option<(Token![,], Type)>,
_trailing_comma: Option<Token![,]>,
_gt_token: Token![>],
},
}
impl StorageType {
fn generate_type_declaration(
&self,
crate_: &syn::Path,
storage_instance: &StorageInstance,
storage_name: &Ident,
storage_generics: Option<&SimpleGenerics>,
visibility: &Visibility,
attributes: &[Attribute],
) -> TokenStream {
let storage_instance_generics = &storage_instance.generics;
let storage_instance = &storage_instance.name;
let attributes = attributes.iter();
let storage_generics = storage_generics.map(|g| {
let generics = g.type_generics();
quote!( < #( #generics ),* > )
});
match self {
Self::Value { value_ty, query_type, .. } => {
let query_type = query_type.as_ref().map(|(c, t)| quote!(#c #t));
quote! {
#( #attributes )*
#visibility type #storage_name #storage_generics = #crate_::storage::types::StorageValue<
#storage_instance #storage_instance_generics,
#value_ty
#query_type
>;
}
},
Self::CountedMap { value_ty, query_type, hasher_ty, key_ty, .. } |
Self::Map { value_ty, query_type, hasher_ty, key_ty, .. } => {
let query_type = query_type.as_ref().map(|(c, t)| quote!(#c #t));
let map_type = Ident::new(
match self {
Self::Map { .. } => "StorageMap",
_ => "CountedStorageMap",
},
Span::call_site(),
);
quote! {
#( #attributes )*
#visibility type #storage_name #storage_generics = #crate_::storage::types::#map_type<
#storage_instance #storage_instance_generics,
#hasher_ty,
#key_ty,
#value_ty
#query_type
>;
}
},
Self::DoubleMap {
value_ty,
query_type,
hasher1_ty,
key1_ty,
hasher2_ty,
key2_ty,
..
} => {
let query_type = query_type.as_ref().map(|(c, t)| quote!(#c #t));
quote! {
#( #attributes )*
#visibility type #storage_name #storage_generics = #crate_::storage::types::StorageDoubleMap<
#storage_instance #storage_instance_generics,
#hasher1_ty,
#key1_ty,
#hasher2_ty,
#key2_ty,
#value_ty
#query_type
>;
}
},
Self::NMap { value_ty, query_type, key_types, .. } => {
let query_type = query_type.as_ref().map(|(c, t)| quote!(#c #t));
let key_types = key_types.iter();
quote! {
#( #attributes )*
#visibility type #storage_name #storage_generics = #crate_::storage::types::StorageNMap<
#storage_instance #storage_instance_generics,
( #( #key_types ),* ),
#value_ty
#query_type
>;
}
},
}
}
fn prefix(&self) -> &Type {
match self {
Self::Value { prefix, .. } |
Self::Map { prefix, .. } |
Self::CountedMap { prefix, .. } |
Self::NMap { prefix, .. } |
Self::DoubleMap { prefix, .. } => prefix,
}
}
}
impl Parse for StorageType {
fn parse(input: ParseStream<'_>) -> Result<Self> {
let lookahead = input.lookahead1();
let parse_query_type = |input: ParseStream<'_>| -> Result<Option<(Token![,], Type)>> {
if input.peek(Token![,]) && !input.peek2(Token![>]) {
Ok(Some((input.parse()?, input.parse()?)))
} else {
Ok(None)
}
};
if lookahead.peek(storage_types::StorageValue) {
Ok(Self::Value {
_kw: input.parse()?,
_lt_token: input.parse()?,
prefix: input.parse()?,
_value_comma: input.parse()?,
value_ty: input.parse()?,
query_type: parse_query_type(input)?,
_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
_gt_token: input.parse()?,
})
} else if lookahead.peek(storage_types::StorageMap) {
Ok(Self::Map {
_kw: input.parse()?,
_lt_token: input.parse()?,
prefix: input.parse()?,
_hasher_comma: input.parse()?,
hasher_ty: input.parse()?,
_key_comma: input.parse()?,
key_ty: input.parse()?,
_value_comma: input.parse()?,
value_ty: input.parse()?,
query_type: parse_query_type(input)?,
_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
_gt_token: input.parse()?,
})
} else if lookahead.peek(storage_types::CountedStorageMap) {
Ok(Self::CountedMap {
_kw: input.parse()?,
_lt_token: input.parse()?,
prefix: input.parse()?,
_hasher_comma: input.parse()?,
hasher_ty: input.parse()?,
_key_comma: input.parse()?,
key_ty: input.parse()?,
_value_comma: input.parse()?,
value_ty: input.parse()?,
query_type: parse_query_type(input)?,
_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
_gt_token: input.parse()?,
})
} else if lookahead.peek(storage_types::StorageDoubleMap) {
Ok(Self::DoubleMap {
_kw: input.parse()?,
_lt_token: input.parse()?,
prefix: input.parse()?,
_hasher1_comma: input.parse()?,
hasher1_ty: input.parse()?,
_key1_comma: input.parse()?,
key1_ty: input.parse()?,
_hasher2_comma: input.parse()?,
hasher2_ty: input.parse()?,
_key2_comma: input.parse()?,
key2_ty: input.parse()?,
_value_comma: input.parse()?,
value_ty: input.parse()?,
query_type: parse_query_type(input)?,
_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
_gt_token: input.parse()?,
})
} else if lookahead.peek(storage_types::StorageNMap) {
let content;
Ok(Self::NMap {
_kw: input.parse()?,
_lt_token: input.parse()?,
prefix: input.parse()?,
_paren_comma: input.parse()?,
_paren_token: parenthesized!(content in input),
key_types: Punctuated::parse_terminated(&content)?,
_value_comma: input.parse()?,
value_ty: input.parse()?,
query_type: parse_query_type(input)?,
_trailing_comma: input.peek(Token![,]).then(|| input.parse()).transpose()?,
_gt_token: input.parse()?,
})
} else {
Err(lookahead.error())
}
}
}
struct Input {
attributes: Vec<Attribute>,
visibility: Visibility,
_type: Token![type],
storage_name: Ident,
storage_generics: Option<SimpleGenerics>,
where_clause: Option<WhereClause>,
_equal: Token![=],
storage_type: StorageType,
_semicolon: Token![;],
}
impl Parse for Input {
fn parse(input: ParseStream<'_>) -> Result<Self> {
let attributes = input.call(Attribute::parse_outer)?;
let visibility = input.parse()?;
let _type = input.parse()?;
let storage_name = input.parse()?;
let lookahead = input.lookahead1();
let storage_generics = if lookahead.peek(Token![<]) {
Some(input.parse()?)
} else if lookahead.peek(Token![=]) {
None
} else {
return Err(lookahead.error())
};
let lookahead = input.lookahead1();
let where_clause = if lookahead.peek(Token![where]) {
Some(input.parse()?)
} else if lookahead.peek(Token![=]) {
None
} else {
return Err(lookahead.error())
};
let _equal = input.parse()?;
let storage_type = input.parse()?;
let _semicolon = input.parse()?;
Ok(Self {
attributes,
visibility,
_type,
storage_name,
storage_generics,
_equal,
storage_type,
where_clause,
_semicolon,
})
}
}
#[derive(Clone, Copy)]
enum PrefixType {
Compatibility,
Verbatim,
PalletName,
Dynamic,
}
pub fn storage_alias(attributes: TokenStream, input: TokenStream) -> Result<TokenStream> {
let input = syn::parse2::<Input>(input)?;
let crate_ = generate_access_from_frame_or_crate("frame-support")?;
let prefix_type = if attributes.is_empty() {
PrefixType::Compatibility
} else if syn::parse2::<prefix_types::verbatim>(attributes.clone()).is_ok() {
PrefixType::Verbatim
} else if syn::parse2::<prefix_types::pallet_name>(attributes.clone()).is_ok() {
PrefixType::PalletName
} else if syn::parse2::<prefix_types::dynamic>(attributes.clone()).is_ok() {
PrefixType::Dynamic
} else {
return Err(Error::new(attributes.span(), "Unknown attributes"))
};
let storage_instance = generate_storage_instance(
&crate_,
&input.storage_name,
input.storage_generics.as_ref(),
input.where_clause.as_ref(),
input.storage_type.prefix(),
&input.visibility,
matches!(input.storage_type, StorageType::CountedMap { .. }),
prefix_type,
)?;
let definition = input.storage_type.generate_type_declaration(
&crate_,
&storage_instance,
&input.storage_name,
input.storage_generics.as_ref(),
&input.visibility,
&input.attributes,
);
let storage_instance_code = storage_instance.code;
Ok(quote! {
#storage_instance_code
#definition
})
}
struct StorageInstance {
name: Ident,
generics: TokenStream,
code: TokenStream,
}
fn generate_storage_instance(
crate_: &syn::Path,
storage_name: &Ident,
storage_generics: Option<&SimpleGenerics>,
storage_where_clause: Option<&WhereClause>,
prefix: &Type,
visibility: &Visibility,
is_counted_map: bool,
prefix_type: PrefixType,
) -> Result<StorageInstance> {
if let Type::Infer(_) = prefix {
return Err(Error::new(prefix.span(), "`_` is not allowed as prefix by `storage_alias`."))
}
let impl_generics_used_by_prefix = storage_generics
.as_ref()
.map(|g| {
g.impl_generics()
.filter(|g| prefix.contains_ident(&g.ident))
.collect::<Vec<_>>()
})
.unwrap_or_default();
let (pallet_prefix, impl_generics, type_generics) = match prefix_type {
PrefixType::Compatibility =>
if !impl_generics_used_by_prefix.is_empty() {
let type_generics = impl_generics_used_by_prefix.iter().map(|g| &g.ident);
let impl_generics = impl_generics_used_by_prefix.iter();
(
quote! {
< #prefix as #crate_::traits::PalletInfoAccess>::name()
},
quote!( #( #impl_generics ),* ),
quote!( #( #type_generics ),* ),
)
} else if let Some(prefix) = prefix.get_ident() {
let prefix_str = prefix.to_string();
(quote!(#prefix_str), quote!(), quote!())
} else {
return Err(Error::new_spanned(
prefix,
"If there are no generics, the prefix is only allowed to be an identifier.",
))
},
PrefixType::Verbatim => {
let prefix_str = match prefix.get_ident() {
Some(p) => p.to_string(),
None =>
return Err(Error::new_spanned(
prefix,
"Prefix type `verbatim` requires that the prefix is an ident.",
)),
};
(quote!(#prefix_str), quote!(), quote!())
},
PrefixType::PalletName => {
let type_generics = impl_generics_used_by_prefix.iter().map(|g| &g.ident);
let impl_generics = impl_generics_used_by_prefix.iter();
(
quote! {
<#prefix as #crate_::traits::PalletInfoAccess>::name()
},
quote!( #( #impl_generics ),* ),
quote!( #( #type_generics ),* ),
)
},
PrefixType::Dynamic => {
let type_generics = impl_generics_used_by_prefix.iter().map(|g| &g.ident);
let impl_generics = impl_generics_used_by_prefix.iter();
(
quote! {
<#prefix as #crate_::traits::Get<_>>::get()
},
quote!( #( #impl_generics ),* ),
quote!( #( #type_generics ),* ),
)
},
};
let where_clause = storage_where_clause.map(|w| quote!(#w)).unwrap_or_default();
let name_str = format!("{}_Storage_Instance", storage_name);
let name = Ident::new(&name_str, Span::call_site());
let storage_name_str = storage_name.to_string();
let counter_code = is_counted_map.then(|| {
let counter_name = Ident::new(&counter_prefix(&name_str), Span::call_site());
let counter_storage_name_str = counter_prefix(&storage_name_str);
let storage_prefix_hash = helper::two128_str(&counter_storage_name_str);
quote! {
#visibility struct #counter_name< #impl_generics >(
::core::marker::PhantomData<(#type_generics)>
) #where_clause;
impl<#impl_generics> #crate_::traits::StorageInstance
for #counter_name< #type_generics > #where_clause
{
fn pallet_prefix() -> &'static str {
#pallet_prefix
}
const STORAGE_PREFIX: &'static str = #counter_storage_name_str;
fn storage_prefix_hash() -> [u8; 16] {
#storage_prefix_hash
}
}
impl<#impl_generics> #crate_::storage::types::CountedStorageMapInstance
for #name< #type_generics > #where_clause
{
type CounterPrefix = #counter_name < #type_generics >;
}
}
});
let storage_prefix_hash = helper::two128_str(&storage_name_str);
let code = quote! {
#[allow(non_camel_case_types)]
#visibility struct #name< #impl_generics >(
::core::marker::PhantomData<(#type_generics)>
) #where_clause;
impl<#impl_generics> #crate_::traits::StorageInstance
for #name< #type_generics > #where_clause
{
fn pallet_prefix() -> &'static str {
#pallet_prefix
}
const STORAGE_PREFIX: &'static str = #storage_name_str;
fn storage_prefix_hash() -> [u8; 16] {
#storage_prefix_hash
}
}
#counter_code
};
Ok(StorageInstance { name, code, generics: quote!( < #type_generics > ) })
}