From a1d42aaa9f52531b27f882832cca9efe370ffcb6 Mon Sep 17 00:00:00 2001 From: Roman Useinov Date: Mon, 13 Feb 2023 23:22:51 +0100 Subject: [PATCH] [Feature] Introduce storage_alias for CountedStorageMap (#13366) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [Feature] Introduce storagage_alias for CountedStorageMap * bit more dry * bit more dry * address review comments * some tests and fixes * fix ui tests * Apply suggestions from code review Co-authored-by: Bastian Köcher * compare metadata --------- Co-authored-by: parity-processbot <> Co-authored-by: Bastian Köcher --- frame/support/procedural/src/lib.rs | 6 ++ .../procedural/src/pallet/expand/storage.rs | 15 ++-- frame/support/procedural/src/storage_alias.rs | 81 ++++++++++++++++++- .../support/src/storage/types/counted_map.rs | 10 +++ frame/support/test/tests/pallet.rs | 18 ++++- .../forbid_underscore_as_prefix.stderr | 2 +- 6 files changed, 119 insertions(+), 13 deletions(-) diff --git a/frame/support/procedural/src/lib.rs b/frame/support/procedural/src/lib.rs index 8fa46b48d005f..67e592f4ca2db 100644 --- a/frame/support/procedural/src/lib.rs +++ b/frame/support/procedural/src/lib.rs @@ -69,6 +69,12 @@ fn get_cargo_env_var(version_env: &str) -> std::result::Result String { + format!("CounterFor{}", prefix) +} + /// Declares strongly-typed wrappers around codec-compatible types in storage. /// /// ## Example diff --git a/frame/support/procedural/src/pallet/expand/storage.rs b/frame/support/procedural/src/pallet/expand/storage.rs index 195a62431f279..0bb19ad6c29cd 100644 --- a/frame/support/procedural/src/pallet/expand/storage.rs +++ b/frame/support/procedural/src/pallet/expand/storage.rs @@ -15,9 +15,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::pallet::{ - parse::storage::{Metadata, QueryKind, StorageDef, StorageGenerics}, - Def, +use crate::{ + counter_prefix, + pallet::{ + parse::storage::{Metadata, QueryKind, StorageDef, StorageGenerics}, + Def, + }, }; use quote::ToTokens; use std::{collections::HashMap, ops::IndexMut}; @@ -39,12 +42,6 @@ fn counter_prefix_ident(storage_ident: &syn::Ident) -> syn::Ident { ) } -/// 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. diff --git a/frame/support/procedural/src/storage_alias.rs b/frame/support/procedural/src/storage_alias.rs index e0df0123595b9..eaf8591a1d999 100644 --- a/frame/support/procedural/src/storage_alias.rs +++ b/frame/support/procedural/src/storage_alias.rs @@ -17,6 +17,7 @@ //! Implementation of the `storage_alias` attribute macro. +use crate::counter_prefix; use frame_support_procedural_tools::generate_crate_access_2018; use proc_macro2::{Span, TokenStream}; use quote::{quote, ToTokens}; @@ -136,6 +137,7 @@ impl ToTokens for SimpleGenerics { mod storage_types { syn::custom_keyword!(StorageValue); syn::custom_keyword!(StorageMap); + syn::custom_keyword!(CountedStorageMap); syn::custom_keyword!(StorageDoubleMap); syn::custom_keyword!(StorageNMap); } @@ -168,6 +170,21 @@ enum StorageType { _trailing_comma: Option, _gt_token: Token![>], }, + CountedMap { + _kw: storage_types::CountedStorageMap, + _lt_token: Token![<], + prefix: SimplePath, + prefix_generics: Option, + _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, + _gt_token: Token![>], + }, DoubleMap { _kw: storage_types::StorageDoubleMap, _lt_token: Token![<], @@ -235,12 +252,22 @@ impl StorageType { >; } }, + Self::CountedMap { + value_ty, query_type, hasher_ty, key_ty, prefix_generics, .. + } | Self::Map { value_ty, query_type, hasher_ty, key_ty, prefix_generics, .. } => { 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::StorageMap< + #visibility type #storage_name #storage_generics = #crate_::storage::types::#map_type< #storage_instance #prefix_generics, #hasher_ty, #key_ty, @@ -296,6 +323,7 @@ impl StorageType { match self { Self::Value { prefix, .. } | Self::Map { prefix, .. } | + Self::CountedMap { prefix, .. } | Self::NMap { prefix, .. } | Self::DoubleMap { prefix, .. } => prefix, } @@ -306,6 +334,7 @@ impl StorageType { match self { Self::Value { prefix_generics, .. } | Self::Map { prefix_generics, .. } | + Self::CountedMap { prefix_generics, .. } | Self::NMap { prefix_generics, .. } | Self::DoubleMap { prefix_generics, .. } => prefix_generics.as_ref(), } @@ -363,6 +392,22 @@ impl Parse for StorageType { _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()?, + prefix_generics: parse_pallet_generics(input)?, + _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()?, @@ -476,6 +521,7 @@ pub fn storage_alias(input: TokenStream) -> Result { input.storage_type.prefix(), input.storage_type.prefix_generics(), &input.visibility, + matches!(input.storage_type, StorageType::CountedMap { .. }), )?; let definition = input.storage_type.generate_type_declaration( @@ -511,6 +557,7 @@ fn generate_storage_instance( prefix: &SimplePath, prefix_generics: Option<&TypeGenerics>, visibility: &Visibility, + is_counted_map: bool, ) -> Result { if let Some(ident) = prefix.get_ident().filter(|i| *i == "_") { return Err(Error::new(ident.span(), "`_` is not allowed as prefix by `storage_alias`.")) @@ -546,9 +593,37 @@ fn generate_storage_instance( let where_clause = storage_where_clause.map(|w| quote!(#w)).unwrap_or_default(); - let name = Ident::new(&format!("{}_Storage_Instance", storage_name), Span::call_site()); + 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); + + quote! { + #visibility struct #counter_name< #impl_generics >( + #crate_::sp_std::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; + } + + impl<#impl_generics> #crate_::storage::types::CountedStorageMapInstance + for #name< #type_generics > #where_clause + { + type CounterPrefix = #counter_name < #type_generics >; + } + } + }); + // Implement `StorageInstance` trait. let code = quote! { #visibility struct #name< #impl_generics >( @@ -564,6 +639,8 @@ fn generate_storage_instance( const STORAGE_PREFIX: &'static str = #storage_name_str; } + + #counter_code }; Ok(StorageInstance { name, code }) diff --git a/frame/support/src/storage/types/counted_map.rs b/frame/support/src/storage/types/counted_map.rs index 3361c4093dce0..7d7d9ebcf70cf 100644 --- a/frame/support/src/storage/types/counted_map.rs +++ b/frame/support/src/storage/types/counted_map.rs @@ -543,6 +543,16 @@ mod test { 97 } } + #[crate::storage_alias] + type ExampleCountedMap = CountedStorageMap; + + #[test] + fn storage_alias_works() { + TestExternalities::default().execute_with(|| { + assert_eq!(ExampleCountedMap::count(), 0); + ExampleCountedMap::insert(3, 10); + }) + } #[test] fn test_value_query() { diff --git a/frame/support/test/tests/pallet.rs b/frame/support/test/tests/pallet.rs index c0376d5aa450f..36150b28058a7 100644 --- a/frame/support/test/tests/pallet.rs +++ b/frame/support/test/tests/pallet.rs @@ -19,7 +19,7 @@ use frame_support::{ dispatch::{ DispatchClass, DispatchInfo, GetDispatchInfo, Parameter, Pays, UnfilteredDispatchable, }, - pallet_prelude::ValueQuery, + pallet_prelude::{StorageInfoTrait, ValueQuery}, storage::unhashed, traits::{ ConstU32, GetCallName, GetStorageVersion, OnFinalize, OnGenesis, OnInitialize, @@ -1906,14 +1906,30 @@ fn assert_type_all_pallets_without_system_reversed_is_correct() { #[test] fn test_storage_alias() { + use frame_support::Twox64Concat; + #[frame_support::storage_alias] type Value where ::AccountId: From + SomeAssociation1, = StorageValue, u32, ValueQuery>; + #[frame_support::storage_alias] + type SomeCountedStorageMap + where + ::AccountId: From + SomeAssociation1, + = CountedStorageMap, Twox64Concat, u8, u32>; + TestExternalities::default().execute_with(|| { pallet::Value::::put(10); assert_eq!(10, Value::::get()); + + pallet2::SomeCountedStorageMap::::insert(10, 100); + assert_eq!(Some(100), SomeCountedStorageMap::::get(10)); + assert_eq!(1, SomeCountedStorageMap::::count()); + assert_eq!( + SomeCountedStorageMap::::storage_info(), + pallet2::SomeCountedStorageMap::::storage_info() + ); }) } diff --git a/frame/support/test/tests/storage_alias_ui/forbid_underscore_as_prefix.stderr b/frame/support/test/tests/storage_alias_ui/forbid_underscore_as_prefix.stderr index 3aa517ecfa314..3b5e3e9c23cca 100644 --- a/frame/support/test/tests/storage_alias_ui/forbid_underscore_as_prefix.stderr +++ b/frame/support/test/tests/storage_alias_ui/forbid_underscore_as_prefix.stderr @@ -1,4 +1,4 @@ -error: expected one of: `StorageValue`, `StorageMap`, `StorageDoubleMap`, `StorageNMap` +error: expected one of: `StorageValue`, `StorageMap`, `CountedStorageMap`, `StorageDoubleMap`, `StorageNMap` --> tests/storage_alias_ui/forbid_underscore_as_prefix.rs:2:14 | 2 | type Ident = CustomStorage;