diff --git a/frame/example/src/lib.rs b/frame/example/src/lib.rs index f5014b75640ba..db609ed7b9110 100644 --- a/frame/example/src/lib.rs +++ b/frame/example/src/lib.rs @@ -598,6 +598,9 @@ pub mod pallet { #[pallet::getter(fn foo)] pub(super) type Foo = StorageValue<_, T::Balance, ValueQuery>; + #[pallet::storage] + pub type CountedMap = CountedStorageMap<_, Blake2_128Concat, u8, u16>; + // The genesis config type. #[pallet::genesis_config] pub struct GenesisConfig { diff --git a/frame/example/src/tests.rs b/frame/example/src/tests.rs index a290ea0f6576f..5d39c53d7b149 100644 --- a/frame/example/src/tests.rs +++ b/frame/example/src/tests.rs @@ -175,6 +175,15 @@ fn signed_ext_watch_dummy_works() { }) } +#[test] +fn counted_map_works() { + new_test_ext().execute_with(|| { + assert_eq!(CountedMap::::count(), 0); + CountedMap::::insert(3, 3); + assert_eq!(CountedMap::::count(), 1); + }) +} + #[test] fn weights_work() { // must have a defined weight. diff --git a/frame/support/procedural/src/pallet/expand/storage.rs b/frame/support/procedural/src/pallet/expand/storage.rs index 0000051dd9b94..b3c89f0808d4f 100644 --- a/frame/support/procedural/src/pallet/expand/storage.rs +++ b/frame/support/procedural/src/pallet/expand/storage.rs @@ -18,30 +18,66 @@ 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; +use std::collections::HashMap; -/// Generate the prefix_ident related the the storage. +/// Generate the prefix_ident related to the storage. /// prefix_ident is used for the prefix struct to be given to storage as first generic param. fn prefix_ident(storage: &StorageDef) -> syn::Ident { let storage_ident = &storage.ident; syn::Ident::new(&format!("_GeneratedPrefixForStorage{}", storage_ident), storage_ident.span()) } +/// Generate the counter_prefix_ident related to the storage. +/// counter_prefix_ident is used for the prefix struct to be given to counted storage map. +fn counter_prefix_ident(storage_ident: &syn::Ident) -> syn::Ident { + syn::Ident::new( + &format!("_GeneratedCounterPrefixForStorage{}", storage_ident), + storage_ident.span(), + ) +} + +/// 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. fn check_prefix_duplicates( storage_def: &StorageDef, - set: &mut HashSet, + // A hashmap of all already used prefix and their associated error if duplication + set: &mut HashMap, ) -> syn::Result<()> { let prefix = storage_def.prefix(); + let dup_err = syn::Error::new( + storage_def.prefix_span(), + format!("Duplicate storage prefixes found for `{}`", prefix), + ); + + if let Some(other_dup_err) = set.insert(prefix.clone(), dup_err.clone()) { + let mut err = dup_err; + err.combine(other_dup_err); + return Err(err); + } - if !set.insert(prefix.clone()) { - let err = syn::Error::new( + if let Metadata::CountedMap { .. } = storage_def.metadata { + let counter_prefix = counter_prefix(&prefix); + let counter_dup_err = syn::Error::new( storage_def.prefix_span(), - format!("Duplicate storage prefixes found for `{}`", prefix), + format!( + "Duplicate storage prefixes found for `{}`, used for counter associated to \ + counted storage map", + counter_prefix, + ), ); - return Err(err); + + if let Some(other_dup_err) = set.insert(counter_prefix.clone(), counter_dup_err.clone()) { + let mut err = counter_dup_err; + err.combine(other_dup_err); + return Err(err); + } } Ok(()) @@ -52,11 +88,8 @@ fn check_prefix_duplicates( /// * Add `#[allow(type_alias_bounds)]` 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 { @@ -112,6 +145,19 @@ pub fn process_generics(def: &mut Def) -> syn::Result<()> { let max_values = max_values.unwrap_or_else(|| default_max_values.clone()); args.args.push(syn::GenericArgument::Type(max_values)); } + StorageGenerics::CountedMap { + hasher, key, value, query_kind, on_empty, max_values + } => { + args.args.push(syn::GenericArgument::Type(hasher)); + args.args.push(syn::GenericArgument::Type(key)); + args.args.push(syn::GenericArgument::Type(value)); + let query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone()); + args.args.push(syn::GenericArgument::Type(query_kind)); + let on_empty = on_empty.unwrap_or_else(|| default_on_empty.clone()); + args.args.push(syn::GenericArgument::Type(on_empty)); + let max_values = max_values.unwrap_or_else(|| default_max_values.clone()); + args.args.push(syn::GenericArgument::Type(max_values)); + } StorageGenerics::DoubleMap { hasher1, key1, hasher2, key2, value, query_kind, on_empty, max_values, } => { @@ -157,11 +203,19 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { return e.into_compile_error().into(); } + // Check for duplicate prefixes + let mut prefix_set = HashMap::new(); + let mut errors = def.storages.iter() + .filter_map(|storage_def| check_prefix_duplicates(storage_def, &mut prefix_set).err()); + if let Some(mut final_error) = errors.next() { + errors.for_each(|error| final_error.combine(error)); + return final_error.into_compile_error() + } + let frame_support = &def.frame_support; let frame_system = &def.frame_system; let pallet_ident = &def.pallet_struct.pallet; - let entries = def.storages.iter() .map(|storage| { let docs = &storage.docs; @@ -179,6 +233,9 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { Metadata::Map { .. } => quote::quote_spanned!(storage.attr_span => #frame_support::storage::types::StorageMapMetadata ), + Metadata::CountedMap { .. } => quote::quote_spanned!(storage.attr_span => + #frame_support::storage::types::CountedStorageMapMetadata + ), Metadata::DoubleMap { .. } => quote::quote_spanned!(storage.attr_span => #frame_support::storage::types::StorageDoubleMapMetadata ), @@ -208,6 +265,18 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { } ) }, + Metadata::CountedMap { key, value } => { + let value = clean_type_string("e::quote!(#value).to_string()); + let key = clean_type_string("e::quote!(#key).to_string()); + quote::quote_spanned!(storage.attr_span => + #frame_support::metadata::StorageEntryType::Map { + hasher: <#full_ident as #metadata_trait>::HASHER, + key: #frame_support::metadata::DecodeDifferent::Encode(#key), + value: #frame_support::metadata::DecodeDifferent::Encode(#value), + unused: false, + } + ) + }, Metadata::DoubleMap { key1, key2, value } => { let value = clean_type_string("e::quote!(#value).to_string()); let key1 = clean_type_string("e::quote!(#key1).to_string()); @@ -242,7 +311,7 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { } }; - quote::quote_spanned!(storage.attr_span => + let mut metadata = quote::quote_spanned!(storage.attr_span => #(#cfg_attrs)* #frame_support::metadata::StorageEntryMetadata { name: #frame_support::metadata::DecodeDifferent::Encode( <#full_ident as #metadata_trait>::NAME @@ -256,7 +325,32 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { #( #docs, )* ]), } - ) + ); + + // Additional metadata for some storages: + if let Metadata::CountedMap { .. } = storage.metadata { + metadata.extend(quote::quote_spanned!(storage.attr_span => + , #(#cfg_attrs)* #frame_support::metadata::StorageEntryMetadata { + name: #frame_support::metadata::DecodeDifferent::Encode( + <#full_ident as #metadata_trait>::COUNTER_NAME + ), + modifier: <#full_ident as #metadata_trait>::COUNTER_MODIFIER, + ty: #frame_support::metadata::StorageEntryType::Plain( + #frame_support::metadata::DecodeDifferent::Encode( + <#full_ident as #metadata_trait>::COUNTER_TY + ) + ), + default: #frame_support::metadata::DecodeDifferent::Encode( + <#full_ident as #metadata_trait>::COUNTER_DEFAULT + ), + documentation: #frame_support::metadata::DecodeDifferent::Encode(&[ + <#full_ident as #metadata_trait>::COUNTER_DOC + ]), + } + )); + } + + metadata }); let getters = def.storages.iter() @@ -317,6 +411,27 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { } ) }, + Metadata::CountedMap { key, value } => { + let query = match storage.query_kind.as_ref().expect("Checked by def") { + QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span => + Option<#value> + ), + 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(k: KArg) -> #query where + KArg: #frame_support::codec::EncodeLike<#key>, + { + // NOTE: we can't use any trait here because CountedStorageMap + // doesn't implement any. + <#full_ident>::get(k) + } + } + ) + }, Metadata::DoubleMap { key1, key2, value } => { let query = match storage.query_kind.as_ref().expect("Checked by def") { QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span => @@ -381,7 +496,44 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream { let cfg_attrs = &storage_def.cfg_attrs; + let maybe_counter = if let Metadata::CountedMap { .. } = storage_def.metadata { + let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident); + let counter_prefix_struct_const = counter_prefix(&prefix_struct_const); + + quote::quote_spanned!(storage_def.attr_span => + #(#cfg_attrs)* + #prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>( + core::marker::PhantomData<(#type_use_gen,)> + ); + #(#cfg_attrs)* + impl<#type_impl_gen> #frame_support::traits::StorageInstance + for #counter_prefix_struct_ident<#type_use_gen> + #config_where_clause + { + fn pallet_prefix() -> &'static str { + < + ::PalletInfo + as #frame_support::traits::PalletInfo + >::name::>() + .expect("Every active pallet has a name in the runtime; qed") + } + const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const; + } + #(#cfg_attrs)* + impl<#type_impl_gen> #frame_support::storage::types::CountedStorageMapInstance + for #prefix_struct_ident<#type_use_gen> + #config_where_clause + { + type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>; + } + ) + } else { + proc_macro2::TokenStream::default() + }; + quote::quote_spanned!(storage_def.attr_span => + #maybe_counter + #(#cfg_attrs)* #prefix_struct_vis struct #prefix_struct_ident<#type_use_gen>( core::marker::PhantomData<(#type_use_gen,)> diff --git a/frame/support/procedural/src/pallet/parse/storage.rs b/frame/support/procedural/src/pallet/parse/storage.rs index 9ec890e66e57a..9dc053a3d980d 100644 --- a/frame/support/procedural/src/pallet/parse/storage.rs +++ b/frame/support/procedural/src/pallet/parse/storage.rs @@ -86,6 +86,7 @@ impl syn::parse::Parse for PalletStorageAttr { pub enum Metadata { Value { value: syn::Type }, Map { value: syn::Type, key: syn::Type }, + CountedMap { value: syn::Type, key: syn::Type }, DoubleMap { value: syn::Type, key1: syn::Type, @@ -160,6 +161,14 @@ pub enum StorageGenerics { on_empty: Option, max_values: Option, }, + CountedMap { + hasher: syn::Type, + key: syn::Type, + value: syn::Type, + query_kind: Option, + on_empty: Option, + max_values: Option, + }, Value { value: syn::Type, query_kind: Option, @@ -180,6 +189,7 @@ impl StorageGenerics { let res = match self.clone() { Self::DoubleMap { value, key1, key2, .. } => Metadata::DoubleMap { value, key1, key2 }, Self::Map { value, key, .. } => Metadata::Map { value, key }, + Self::CountedMap { value, key, .. } => Metadata::CountedMap { value, key }, Self::Value { value, .. } => Metadata::Value { value }, Self::NMap { keygen, value, .. } => Metadata::NMap { keys: collect_keys(&keygen)?, @@ -196,6 +206,7 @@ impl StorageGenerics { match &self { Self::DoubleMap { query_kind, .. } | Self::Map { query_kind, .. } + | Self::CountedMap { query_kind, .. } | Self::Value { query_kind, .. } | Self::NMap { query_kind, .. } => query_kind.clone(), @@ -206,6 +217,7 @@ impl StorageGenerics { enum StorageKind { Value, Map, + CountedMap, DoubleMap, NMap, } @@ -333,6 +345,30 @@ fn process_named_generics( max_values: parsed.remove("MaxValues").map(|binding| binding.ty), } } + StorageKind::CountedMap => { + check_generics( + &parsed, + &["Hasher", "Key", "Value"], + &["QueryKind", "OnEmpty", "MaxValues"], + "CountedStorageMap", + args_span, + )?; + + StorageGenerics::CountedMap { + hasher: parsed.remove("Hasher") + .map(|binding| binding.ty) + .expect("checked above as mandatory generic"), + key: parsed.remove("Key") + .map(|binding| binding.ty) + .expect("checked above as mandatory generic"), + value: parsed.remove("Value") + .map(|binding| binding.ty) + .expect("checked above as mandatory generic"), + query_kind: parsed.remove("QueryKind").map(|binding| binding.ty), + on_empty: parsed.remove("OnEmpty").map(|binding| binding.ty), + max_values: parsed.remove("MaxValues").map(|binding| binding.ty), + } + } StorageKind::DoubleMap => { check_generics( &parsed, @@ -436,6 +472,14 @@ fn process_unnamed_generics( }, retrieve_arg(4).ok(), ), + StorageKind::CountedMap => ( + None, + Metadata::CountedMap { + key: retrieve_arg(2)?, + value: retrieve_arg(3)?, + }, + retrieve_arg(4).ok(), + ), StorageKind::DoubleMap => ( None, Metadata::DoubleMap { @@ -470,6 +514,7 @@ fn process_generics( let storage_kind = match &*segment.ident.to_string() { "StorageValue" => StorageKind::Value, "StorageMap" => StorageKind::Map, + "CountedStorageMap" => StorageKind::CountedMap, "StorageDoubleMap" => StorageKind::DoubleMap, "StorageNMap" => StorageKind::NMap, found => { diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs index 45988c1c7372b..01f4f1a8bc741 100644 --- a/frame/support/src/lib.rs +++ b/frame/support/src/lib.rs @@ -1251,7 +1251,7 @@ pub mod pallet_prelude { weights::{DispatchClass, Pays, Weight}, storage::types::{ Key as NMapKey, StorageDoubleMap, StorageMap, StorageNMap, StorageValue, ValueQuery, - OptionQuery, + OptionQuery, CountedStorageMap, }, storage::bounded_vec::BoundedVec, }; @@ -1626,23 +1626,37 @@ pub mod pallet_prelude { /// * [`pallet_prelude::StorageValue`] expect `Value` and optionally `QueryKind` and `OnEmpty`, /// * [`pallet_prelude::StorageMap`] expect `Hasher`, `Key`, `Value` and optionally `QueryKind` and /// `OnEmpty`, +/// * [`pallet_prelude::CountedStorageMap`] expect `Hasher`, `Key`, `Value` and optionally +/// `QueryKind` and `OnEmpty`, /// * [`pallet_prelude::StorageDoubleMap`] expect `Hasher1`, `Key1`, `Hasher2`, `Key2`, `Value` and /// optionally `QueryKind` and `OnEmpty`. /// /// For unnamed generic argument: Their first generic must be `_` as it is replaced by the macro /// and other generic must declared as a normal declaration of type generic in rust. /// -/// The Prefix generic written by the macro is generated using `PalletInfo::name::>()` -/// and the name of the storage type. +/// The `Prefix` generic written by the macro. It implements `StorageInstance` using +/// `PalletInfo::name::>()` for the pallet prefix and the name of the storage type or +/// the rename as in `#[pallet::storage_prefix = "Bar"]` for storage prefix. /// E.g. if runtime names the pallet "MyExample" then the storage `type Foo = ...` use the -/// prefix: `Twox128(b"MyExample") ++ Twox128(b"Foo")`. +/// pallet prefix: "MyExample" and the storage prefix "Foo". +/// +/// NOTE: Storages use those prefixes so that values are stored after the key +/// `Twox128(b"MyExample") ++ Twox128(b"Foo")`. +/// +/// For the `CountedStorageMap` variant, the Prefix also implements `CountedStorageMapInstance`. +/// It associate a `CounterPrefix`, which is implemented same as above, but the storage prefix is +/// prepend with `"CounterFor"`. /// /// The optional attribute `#[pallet::getter(fn $my_getter_fn_name)]` allow to define a /// getter function on `Pallet`. /// +/// The optional attribute `#[pallet::storage_prefix = "SomeName"]` allow to define the storage +/// prefix to use, see how `Prefix` generic is implemented above. +/// /// E.g: /// ```ignore /// #[pallet::storage] +/// #[pallet::storage_prefix = "foo"] /// #[pallet::getter(fn my_storage)] /// pub(super) type MyStorage = StorageMap; /// ``` @@ -1679,6 +1693,8 @@ pub mod pallet_prelude { /// `_GeneratedPrefixForStorage$NameOfStorage`, and implements /// [`StorageInstance`](traits::StorageInstance) on it using the pallet and storage name. It then /// uses it as the first generic of the aliased type. +/// For `CountedStorageMap`, `CountedStorageMapInstance` is implemented, and another similar struct +/// is generated. /// /// For named generic, the macro will reorder the generics, and remove the names. /// @@ -1966,6 +1982,7 @@ pub mod pallet_prelude { /// /// // Another storage declaration /// #[pallet::storage] +/// #[pallet::storage_prefix = "SomeName"] /// #[pallet::getter(fn my_storage)] /// pub(super) type MyStorage = /// StorageMap; diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs index 836ae25bdbbcf..284b32d80aa7b 100644 --- a/frame/support/src/storage/generator/double_map.rs +++ b/frame/support/src/storage/generator/double_map.rs @@ -226,6 +226,7 @@ impl storage::StorageDoubleMap for G where previous_key: prefix, drain: false, closure: |_raw_key, mut raw_value| V::decode(&mut raw_value), + phantom: Default::default(), } } @@ -353,6 +354,7 @@ impl< let mut key_material = G::Hasher2::reverse(raw_key_without_prefix); Ok((K2::decode(&mut key_material)?, V::decode(&mut raw_value)?)) }, + phantom: Default::default(), } } @@ -375,6 +377,7 @@ impl< let k2 = K2::decode(&mut k2_material)?; Ok((k1, k2, V::decode(&mut raw_value)?)) }, + phantom: Default::default(), } } diff --git a/frame/support/src/storage/generator/map.rs b/frame/support/src/storage/generator/map.rs index 9abc7883937dd..3cfc634c40bb2 100644 --- a/frame/support/src/storage/generator/map.rs +++ b/frame/support/src/storage/generator/map.rs @@ -152,6 +152,7 @@ impl< let mut key_material = G::Hasher::reverse(raw_key_without_prefix); Ok((K::decode(&mut key_material)?, V::decode(&mut raw_value)?)) }, + phantom: Default::default(), } } diff --git a/frame/support/src/storage/generator/nmap.rs b/frame/support/src/storage/generator/nmap.rs index 62f188a26db8d..dd701a48d7234 100755 --- a/frame/support/src/storage/generator/nmap.rs +++ b/frame/support/src/storage/generator/nmap.rs @@ -213,6 +213,7 @@ where previous_key: prefix, drain: false, closure: |_raw_key, mut raw_value| V::decode(&mut raw_value), + phantom: Default::default(), } } @@ -325,6 +326,7 @@ impl> let partial_key = K::decode_partial_key(raw_key_without_prefix)?; Ok((partial_key, V::decode(&mut raw_value)?)) }, + phantom: Default::default(), } } @@ -347,6 +349,7 @@ impl> let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?; Ok((final_key, V::decode(&mut raw_value)?)) }, + phantom: Default::default(), } } diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs index 62db2eff839fb..8b85899972061 100644 --- a/frame/support/src/storage/migration.rs +++ b/frame/support/src/storage/migration.rs @@ -175,7 +175,7 @@ pub fn storage_iter_with_suffix( Ok((raw_key_without_prefix.to_vec(), value)) }; - PrefixIterator { prefix, previous_key, drain: false, closure } + PrefixIterator { prefix, previous_key, drain: false, closure, phantom: Default::default() } } /// Construct iterator to iterate over map items in `module` for the map called `item`. @@ -203,7 +203,7 @@ pub fn storage_key_iter_with_suffix { prefix: from_prefix.to_vec(), previous_key: from_prefix.to_vec(), drain: true, closure: |key, value| Ok((key.to_vec(), value.to_vec())), + phantom: Default::default(), }; for (key, value) in iter { diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 6a02c6572c7f2..e272ab4887622 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -668,10 +668,12 @@ pub trait StorageNMap { KArg: EncodeLikeTuple + TupleToEncodedIter; } -/// Iterate over a prefix and decode raw_key and raw_value into `T`. +/// Iterate or drain over a prefix and decode raw_key and raw_value into `T`. /// /// If any decoding fails it skips it and continues to the next key. -pub struct PrefixIterator { +/// +/// If draining, then the hook `OnRemoval::on_removal` is called after each removal. +pub struct PrefixIterator { prefix: Vec, previous_key: Vec, /// If true then value are removed while iterating @@ -679,9 +681,20 @@ pub struct PrefixIterator { /// Function that take `(raw_key_without_prefix, raw_value)` and decode `T`. /// `raw_key_without_prefix` is the raw storage key without the prefix iterated on. closure: fn(&[u8], &[u8]) -> Result, + phantom: core::marker::PhantomData, +} + +/// Trait for specialising on removal logic of [`PrefixIterator`]. +pub trait PrefixIteratorOnRemoval { + fn on_removal(); +} + +/// No-op implement. +impl PrefixIteratorOnRemoval for () { + fn on_removal() {} } -impl PrefixIterator { +impl PrefixIterator { /// Mutate this iterator into a draining iterator; items iterated are removed from storage. pub fn drain(mut self) -> Self { self.drain = true; @@ -689,7 +702,7 @@ impl PrefixIterator { } } -impl Iterator for PrefixIterator { +impl Iterator for PrefixIterator { type Item = T; fn next(&mut self) -> Option { @@ -710,7 +723,8 @@ impl Iterator for PrefixIterator { } }; if self.drain { - unhashed::kill(&self.previous_key) + unhashed::kill(&self.previous_key); + OnRemoval::on_removal(); } let raw_key_without_prefix = &self.previous_key[self.prefix.len()..]; let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) { @@ -896,6 +910,7 @@ pub trait StoragePrefixedMap { previous_key: prefix.to_vec(), drain: false, closure: |_raw_key, mut raw_value| Value::decode(&mut raw_value), + phantom: Default::default(), } } diff --git a/frame/support/src/storage/types/counted_map.rs b/frame/support/src/storage/types/counted_map.rs new file mode 100644 index 0000000000000..217b26ee9df06 --- /dev/null +++ b/frame/support/src/storage/types/counted_map.rs @@ -0,0 +1,969 @@ +use codec::{FullCodec, Decode, EncodeLike, Encode, Ref}; +use crate::{ + Never, + traits::{StorageInstance, GetDefault, Get, StorageInfo}, + storage::{ + StorageAppend, StorageTryAppend, StorageDecodeLength, generator::StorageMap as _, + types::{ + OptionQuery, ValueQuery, StorageMap, StorageValue, QueryKindTrait, + StorageMapMetadata, StorageValueMetadata + }, + }, +}; +use sp_std::prelude::*; +use sp_runtime::traits::Saturating; +use frame_metadata::{ + DefaultByteGetter, StorageEntryModifier, +}; +use max_encoded_len::MaxEncodedLen; + +/// A wrapper around a `StorageMap` and a `StorageValue` to keep track of how many items +/// are in a map, without needing to iterate all the values. +/// +/// This storage item has additional storage read and write overhead when manipulating values +/// compared to a regular storage map. +/// +/// For functions where we only add or remove a value, a single storage read is needed to check if +/// that value already exists. For mutate functions, two storage reads are used to check if the +/// value existed before and after the mutation. +/// +/// Whenever the counter needs to be updated, an additional read and write occurs to update that +/// counter. +pub struct CountedStorageMap< + Prefix, Hasher, Key, Value, QueryKind=OptionQuery, OnEmpty=GetDefault, MaxValues=GetDefault, +>( + core::marker::PhantomData<(Prefix, Hasher, Key, Value, QueryKind, OnEmpty, MaxValues)> +); + +/// The requirement for an instance of [`CountedStorageMap`]. +pub trait CountedStorageMapInstance: StorageInstance { + /// The prefix to use for the counter storage value. + type CounterPrefix: StorageInstance; +} + +// Internal helper trait to access map from counted storage map. +trait Helper { + type Map; +} + +impl + Helper + for CountedStorageMap +{ + type Map = StorageMap; +} + +type CounterFor

= StorageValue<

::CounterPrefix, u32, ValueQuery>; + +/// On removal logic for updating counter while draining upon some prefix with +/// [`crate::storage::PrefixIterator`]. +pub struct OnRemovalCounterUpdate(core::marker::PhantomData); + +impl + crate::storage::PrefixIteratorOnRemoval + for OnRemovalCounterUpdate +{ + fn on_removal() { + CounterFor::::mutate(|value| value.saturating_dec()); + } +} + +impl + CountedStorageMap +where + Prefix: CountedStorageMapInstance, + Hasher: crate::hash::StorageHasher, + Key: FullCodec, + Value: FullCodec, + QueryKind: QueryKindTrait, + OnEmpty: Get + 'static, + MaxValues: Get>, +{ + /// Get the storage key used to fetch a value corresponding to a specific key. + pub fn hashed_key_for>(key: KeyArg) -> Vec { + ::Map::hashed_key_for(key) + } + + /// Does the value (explicitly) exist in storage? + pub fn contains_key>(key: KeyArg) -> bool { + ::Map::contains_key(key) + } + + /// Load the value associated with the given key from the map. + pub fn get>(key: KeyArg) -> QueryKind::Query { + ::Map::get(key) + } + + /// Try to get the value for the given key from the map. + /// + /// Returns `Ok` if it exists, `Err` if not. + pub fn try_get>(key: KeyArg) -> Result { + ::Map::try_get(key) + } + + /// Swap the values of two keys. + pub fn swap, KeyArg2: EncodeLike>(key1: KeyArg1, key2: KeyArg2) { + ::Map::swap(key1, key2) + } + + /// Store a value to be associated with the given key from the map. + pub fn insert + Clone, ValArg: EncodeLike>(key: KeyArg, val: ValArg) { + if !::Map::contains_key(Ref::from(&key)) { + CounterFor::::mutate(|value| value.saturating_inc()); + } + ::Map::insert(key, val) + } + + /// Remove the value under a key. + pub fn remove + Clone>(key: KeyArg) { + if ::Map::contains_key(Ref::from(&key)) { + CounterFor::::mutate(|value| value.saturating_dec()); + } + ::Map::remove(key) + } + + /// Mutate the value under a key. + pub fn mutate + Clone, R, F: FnOnce(&mut QueryKind::Query) -> R>( + key: KeyArg, + f: F + ) -> R { + Self::try_mutate(key, |v| Ok::(f(v))) + .expect("`Never` can not be constructed; qed") + } + + /// Mutate the item, only if an `Ok` value is returned. + pub fn try_mutate(key: KeyArg, f: F) -> Result + where + KeyArg: EncodeLike + Clone, + F: FnOnce(&mut QueryKind::Query) -> Result, + { + Self::try_mutate_exists(key, |option_value_ref| { + let option_value = core::mem::replace(option_value_ref, None); + let mut query = ::Map::from_optional_value_to_query(option_value); + let res = f(&mut query); + let option_value = ::Map::from_query_to_optional_value(query); + let _ = core::mem::replace(option_value_ref, option_value); + res + }) + } + + /// Mutate the value under a key. Deletes the item if mutated to a `None`. + pub fn mutate_exists + Clone, R, F: FnOnce(&mut Option) -> R>( + key: KeyArg, + f: F + ) -> R { + Self::try_mutate_exists(key, |v| Ok::(f(v))) + .expect("`Never` can not be constructed; qed") + } + + /// Mutate the item, only if an `Ok` value is returned. Deletes the item if mutated to a `None`. + pub fn try_mutate_exists(key: KeyArg, f: F) -> Result + where + KeyArg: EncodeLike + Clone, + F: FnOnce(&mut Option) -> Result, + { + ::Map::try_mutate_exists(key, |option_value| { + let existed = option_value.is_some(); + let res = f(option_value); + let exist = option_value.is_some(); + + if res.is_ok() { + if existed && !exist { + // Value was deleted + CounterFor::::mutate(|value| value.saturating_dec()); + } else if !existed && exist { + // Value was added + CounterFor::::mutate(|value| value.saturating_inc()); + } + } + res + }) + } + + /// Take the value under a key. + pub fn take + Clone>(key: KeyArg) -> QueryKind::Query { + let removed_value = ::Map::mutate_exists( + key, + |value| core::mem::replace(value, None) + ); + if removed_value.is_some() { + CounterFor::::mutate(|value| value.saturating_dec()); + } + ::Map::from_optional_value_to_query(removed_value) + } + + /// Append the given items to the value in the storage. + /// + /// `Value` is required to implement `codec::EncodeAppend`. + /// + /// # Warning + /// + /// If the storage item is not encoded properly, the storage will be overwritten and set to + /// `[item]`. Any default value set for the storage item will be ignored on overwrite. + pub fn append(key: EncodeLikeKey, item: EncodeLikeItem) + where + EncodeLikeKey: EncodeLike + Clone, + Item: Encode, + EncodeLikeItem: EncodeLike, + Value: StorageAppend + { + if !::Map::contains_key(Ref::from(&key)) { + CounterFor::::mutate(|value| value.saturating_inc()); + } + ::Map::append(key, item) + } + + /// Read the length of the storage value without decoding the entire value under the given + /// `key`. + /// + /// `Value` is required to implement [`StorageDecodeLength`]. + /// + /// If the value does not exists or it fails to decode the length, `None` is returned. Otherwise + /// `Some(len)` is returned. + /// + /// # Warning + /// + /// `None` does not mean that `get()` does not return a value. The default value is completly + /// ignored by this function. + pub fn decode_len>(key: KeyArg) -> Option + where Value: StorageDecodeLength, + { + ::Map::decode_len(key) + } + + /// Migrate an item with the given `key` from a defunct `OldHasher` to the current hasher. + /// + /// If the key doesn't exist, then it's a no-op. If it does, then it returns its value. + pub fn migrate_key>( + key: KeyArg + ) -> Option { + ::Map::migrate_key::(key) + } + + /// Remove all value of the storage. + pub fn remove_all() { + CounterFor::::set(0u32); + ::Map::remove_all(None); + } + + /// Iter over all value of the storage. + /// + /// NOTE: If a value failed to decode becaues storage is corrupted then it is skipped. + pub fn iter_values() -> crate::storage::PrefixIterator> { + let map_iterator = ::Map::iter_values(); + crate::storage::PrefixIterator { + prefix: map_iterator.prefix, + previous_key: map_iterator.previous_key, + drain: map_iterator.drain, + closure: map_iterator.closure, + phantom: Default::default(), + } + } + + /// Translate the values of all elements by a function `f`, in the map in no particular order. + /// + /// By returning `None` from `f` for an element, you'll remove it from the map. + /// + /// NOTE: If a value fail to decode because storage is corrupted then it is skipped. + /// + /// # Warning + /// + /// This function must be used with care, before being updated the storage still contains the + /// old type, thus other calls (such as `get`) will fail at decoding it. + /// + /// # Usage + /// + /// This would typically be called inside the module implementation of on_runtime_upgrade. + pub fn translate_values Option>(mut f: F) { + ::Map::translate_values(|old_value| { + let res = f(old_value); + if res.is_none() { + CounterFor::::mutate(|value| value.saturating_dec()); + } + res + }) + } + + /// Try and append the given item to the value in the storage. + /// + /// Is only available if `Value` of the storage implements [`StorageTryAppend`]. + pub fn try_append( + key: KArg, + item: EncodeLikeItem, + ) -> Result<(), ()> + where + KArg: EncodeLike + Clone, + Item: Encode, + EncodeLikeItem: EncodeLike, + Value: StorageTryAppend, + { + let bound = Value::bound(); + let current = ::Map::decode_len(Ref::from(&key)).unwrap_or_default(); + if current < bound { + CounterFor::::mutate(|value| value.saturating_inc()); + let key = ::Map::hashed_key_for(key); + sp_io::storage::append(&key, item.encode()); + Ok(()) + } else { + Err(()) + } + } + + /// Initialize the counter with the actual number of items in the map. + /// + /// This function iterates through all the items in the map and sets the counter. This operation + /// can be very heavy, so use with caution. + /// + /// Returns the number of items in the map which is used to set the counter. + pub fn initialize_counter() -> u32 { + let count = Self::iter_values().count() as u32; + CounterFor::::set(count); + count + } + + /// Return the count. + pub fn count() -> u32 { + CounterFor::::get() + } +} + +impl + CountedStorageMap +where + Prefix: CountedStorageMapInstance, + Hasher: crate::hash::StorageHasher + crate::ReversibleStorageHasher, + Key: FullCodec, + Value: FullCodec, + QueryKind: QueryKindTrait, + OnEmpty: Get + 'static, + MaxValues: Get>, +{ + /// Enumerate all elements in the map in no particular order. + /// + /// If you alter the map while doing this, you'll get undefined results. + pub fn iter() -> crate::storage::PrefixIterator<(Key, Value), OnRemovalCounterUpdate> { + let map_iterator = ::Map::iter(); + crate::storage::PrefixIterator { + prefix: map_iterator.prefix, + previous_key: map_iterator.previous_key, + drain: map_iterator.drain, + closure: map_iterator.closure, + phantom: Default::default(), + } + } + + /// Remove all elements from the map and iterate through them in no particular order. + /// + /// If you add elements to the map while doing this, you'll get undefined results. + pub fn drain() -> crate::storage::PrefixIterator<(Key, Value), OnRemovalCounterUpdate> { + let map_iterator = ::Map::drain(); + crate::storage::PrefixIterator { + prefix: map_iterator.prefix, + previous_key: map_iterator.previous_key, + drain: map_iterator.drain, + closure: map_iterator.closure, + phantom: Default::default(), + } + } + + /// Translate the values of all elements by a function `f`, in the map in no particular order. + /// + /// By returning `None` from `f` for an element, you'll remove it from the map. + /// + /// NOTE: If a value fail to decode because storage is corrupted then it is skipped. + pub fn translate Option>(mut f: F) { + ::Map::translate(|key, old_value| { + let res = f(key, old_value); + if res.is_none() { + CounterFor::::mutate(|value| value.saturating_dec()); + } + res + }) + } +} + +/// Part of storage metadata for a counted storage map. +pub trait CountedStorageMapMetadata { + const MODIFIER: StorageEntryModifier; + const NAME: &'static str; + const DEFAULT: DefaultByteGetter; + const HASHER: frame_metadata::StorageHasher; + const COUNTER_NAME: &'static str; + const COUNTER_MODIFIER: StorageEntryModifier; + const COUNTER_TY: &'static str; + const COUNTER_DEFAULT: DefaultByteGetter; + const COUNTER_DOC: &'static str; +} + +impl CountedStorageMapMetadata + for CountedStorageMap where + Prefix: CountedStorageMapInstance, + Hasher: crate::hash::StorageHasher, + Key: FullCodec, + Value: FullCodec, + QueryKind: QueryKindTrait, + OnEmpty: Get + 'static, + MaxValues: Get>, +{ + const MODIFIER: StorageEntryModifier = ::Map::MODIFIER; + const HASHER: frame_metadata::StorageHasher = ::Map::HASHER; + const NAME: &'static str = ::Map::NAME; + const DEFAULT: DefaultByteGetter = ::Map::DEFAULT; + const COUNTER_NAME: &'static str = CounterFor::::NAME; + const COUNTER_MODIFIER: StorageEntryModifier = CounterFor::::MODIFIER; + const COUNTER_TY: &'static str = "u32"; + const COUNTER_DEFAULT: DefaultByteGetter = CounterFor::::DEFAULT; + const COUNTER_DOC: &'static str = &"Counter for the related counted storage map"; +} + +impl + crate::traits::StorageInfoTrait for + CountedStorageMap +where + Prefix: CountedStorageMapInstance, + Hasher: crate::hash::StorageHasher, + Key: FullCodec + MaxEncodedLen, + Value: FullCodec + MaxEncodedLen, + QueryKind: QueryKindTrait, + OnEmpty: Get + 'static, + MaxValues: Get>, +{ + fn storage_info() -> Vec { + [::Map::storage_info(), CounterFor::::storage_info()].concat() + } +} + +#[cfg(test)] +mod test { + use super::*; + use sp_io::{TestExternalities, hashing::twox_128}; + use crate::hash::*; + use crate::storage::types::ValueQuery; + use crate::storage::bounded_vec::BoundedVec; + use crate::traits::ConstU32; + + struct Prefix; + impl StorageInstance for Prefix { + fn pallet_prefix() -> &'static str { "test" } + const STORAGE_PREFIX: &'static str = "foo"; + } + + struct CounterPrefix; + impl StorageInstance for CounterPrefix { + fn pallet_prefix() -> &'static str { "test" } + const STORAGE_PREFIX: &'static str = "counter_for_foo"; + } + impl CountedStorageMapInstance for Prefix { + type CounterPrefix = CounterPrefix; + } + + struct ADefault; + impl crate::traits::Get for ADefault { + fn get() -> u32 { + 97 + } + } + + #[test] + fn test_value_query() { + type A = CountedStorageMap; + + TestExternalities::default().execute_with(|| { + let mut k: Vec = vec![]; + k.extend(&twox_128(b"test")); + k.extend(&twox_128(b"foo")); + k.extend(&3u16.twox_64_concat()); + assert_eq!(A::hashed_key_for(3).to_vec(), k); + + assert_eq!(A::contains_key(3), false); + assert_eq!(A::get(3), ADefault::get()); + assert_eq!(A::try_get(3), Err(())); + assert_eq!(A::count(), 0); + + // Insert non-existing. + A::insert(3, 10); + + assert_eq!(A::contains_key(3), true); + assert_eq!(A::get(3), 10); + assert_eq!(A::try_get(3), Ok(10)); + assert_eq!(A::count(), 1); + + // Swap non-existing with existing. + A::swap(4, 3); + + assert_eq!(A::contains_key(3), false); + assert_eq!(A::get(3), ADefault::get()); + assert_eq!(A::try_get(3), Err(())); + assert_eq!(A::contains_key(4), true); + assert_eq!(A::get(4), 10); + assert_eq!(A::try_get(4), Ok(10)); + assert_eq!(A::count(), 1); + + // Swap existing with non-existing. + A::swap(4, 3); + + assert_eq!(A::try_get(3), Ok(10)); + assert_eq!(A::contains_key(4), false); + assert_eq!(A::get(4), ADefault::get()); + assert_eq!(A::try_get(4), Err(())); + assert_eq!(A::count(), 1); + + A::insert(4, 11); + + assert_eq!(A::try_get(3), Ok(10)); + assert_eq!(A::try_get(4), Ok(11)); + assert_eq!(A::count(), 2); + + // Swap 2 existing. + A::swap(3, 4); + + assert_eq!(A::try_get(3), Ok(11)); + assert_eq!(A::try_get(4), Ok(10)); + assert_eq!(A::count(), 2); + + // Insert an existing key, shouldn't increment counted values. + A::insert(3, 11); + + assert_eq!(A::count(), 2); + + // Remove non-existing. + A::remove(2); + + assert_eq!(A::contains_key(2), false); + assert_eq!(A::count(), 2); + + // Remove existing. + A::remove(3); + + assert_eq!(A::try_get(3), Err(())); + assert_eq!(A::count(), 1); + + // Mutate non-existing to existing. + A::mutate(3, |query| { + assert_eq!(*query, ADefault::get()); + *query = 40; + }); + + assert_eq!(A::try_get(3), Ok(40)); + assert_eq!(A::count(), 2); + + // Mutate existing to existing. + A::mutate(3, |query| { + assert_eq!(*query, 40); + *query = 40; + }); + + assert_eq!(A::try_get(3), Ok(40)); + assert_eq!(A::count(), 2); + + // Try fail mutate non-existing to existing. + A::try_mutate(2, |query| { + assert_eq!(*query, ADefault::get()); + *query = 4; + Result::<(), ()>::Err(()) + }).err().unwrap(); + + assert_eq!(A::try_get(2), Err(())); + assert_eq!(A::count(), 2); + + // Try succeed mutate non-existing to existing. + A::try_mutate(2, |query| { + assert_eq!(*query, ADefault::get()); + *query = 41; + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(A::try_get(2), Ok(41)); + assert_eq!(A::count(), 3); + + // Try succeed mutate existing to existing. + A::try_mutate(2, |query| { + assert_eq!(*query, 41); + *query = 41; + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(A::try_get(2), Ok(41)); + assert_eq!(A::count(), 3); + + // Try fail mutate non-existing to existing. + A::try_mutate_exists(1, |query| { + assert_eq!(*query, None); + *query = Some(4); + Result::<(), ()>::Err(()) + }).err().unwrap(); + + assert_eq!(A::try_get(1), Err(())); + assert_eq!(A::count(), 3); + + // Try succeed mutate non-existing to existing. + A::try_mutate_exists(1, |query| { + assert_eq!(*query, None); + *query = Some(43); + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(A::try_get(1), Ok(43)); + assert_eq!(A::count(), 4); + + // Try succeed mutate existing to existing. + A::try_mutate_exists(1, |query| { + assert_eq!(*query, Some(43)); + *query = Some(43); + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(A::try_get(1), Ok(43)); + assert_eq!(A::count(), 4); + + // Try succeed mutate existing to non-existing. + A::try_mutate_exists(1, |query| { + assert_eq!(*query, Some(43)); + *query = None; Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(A::try_get(1), Err(())); + assert_eq!(A::count(), 3); + + // Take exsisting. + assert_eq!(A::take(4), 10); + + assert_eq!(A::try_get(4), Err(())); + assert_eq!(A::count(), 2); + + // Take non-exsisting. + assert_eq!(A::take(4), ADefault::get()); + + assert_eq!(A::try_get(4), Err(())); + assert_eq!(A::count(), 2); + + // Remove all. + A::remove_all(); + + assert_eq!(A::count(), 0); + assert_eq!(A::initialize_counter(), 0); + + A::insert(1, 1); + A::insert(2, 2); + + // Iter values. + assert_eq!(A::iter_values().collect::>(), vec![2, 1]); + + // Iter drain values. + assert_eq!(A::iter_values().drain().collect::>(), vec![2, 1]); + assert_eq!(A::count(), 0); + + A::insert(1, 1); + A::insert(2, 2); + + // Test initialize_counter. + assert_eq!(A::initialize_counter(), 2); + }) + } + + #[test] + fn test_option_query() { + type B = CountedStorageMap; + + TestExternalities::default().execute_with(|| { + let mut k: Vec = vec![]; + k.extend(&twox_128(b"test")); + k.extend(&twox_128(b"foo")); + k.extend(&3u16.twox_64_concat()); + assert_eq!(B::hashed_key_for(3).to_vec(), k); + + assert_eq!(B::contains_key(3), false); + assert_eq!(B::get(3), None); + assert_eq!(B::try_get(3), Err(())); + assert_eq!(B::count(), 0); + + // Insert non-existing. + B::insert(3, 10); + + assert_eq!(B::contains_key(3), true); + assert_eq!(B::get(3), Some(10)); + assert_eq!(B::try_get(3), Ok(10)); + assert_eq!(B::count(), 1); + + // Swap non-existing with existing. + B::swap(4, 3); + + assert_eq!(B::contains_key(3), false); + assert_eq!(B::get(3), None); + assert_eq!(B::try_get(3), Err(())); + assert_eq!(B::contains_key(4), true); + assert_eq!(B::get(4), Some(10)); + assert_eq!(B::try_get(4), Ok(10)); + assert_eq!(B::count(), 1); + + // Swap existing with non-existing. + B::swap(4, 3); + + assert_eq!(B::try_get(3), Ok(10)); + assert_eq!(B::contains_key(4), false); + assert_eq!(B::get(4), None); + assert_eq!(B::try_get(4), Err(())); + assert_eq!(B::count(), 1); + + B::insert(4, 11); + + assert_eq!(B::try_get(3), Ok(10)); + assert_eq!(B::try_get(4), Ok(11)); + assert_eq!(B::count(), 2); + + // Swap 2 existing. + B::swap(3, 4); + + assert_eq!(B::try_get(3), Ok(11)); + assert_eq!(B::try_get(4), Ok(10)); + assert_eq!(B::count(), 2); + + // Insert an existing key, shouldn't increment counted values. + B::insert(3, 11); + + assert_eq!(B::count(), 2); + + // Remove non-existing. + B::remove(2); + + assert_eq!(B::contains_key(2), false); + assert_eq!(B::count(), 2); + + // Remove existing. + B::remove(3); + + assert_eq!(B::try_get(3), Err(())); + assert_eq!(B::count(), 1); + + // Mutate non-existing to existing. + B::mutate(3, |query| { + assert_eq!(*query, None); + *query = Some(40) + }); + + assert_eq!(B::try_get(3), Ok(40)); + assert_eq!(B::count(), 2); + + // Mutate existing to existing. + B::mutate(3, |query| { + assert_eq!(*query, Some(40)); + *query = Some(40) + }); + + assert_eq!(B::try_get(3), Ok(40)); + assert_eq!(B::count(), 2); + + // Mutate existing to non-existing. + B::mutate(3, |query| { + assert_eq!(*query, Some(40)); + *query = None + }); + + assert_eq!(B::try_get(3), Err(())); + assert_eq!(B::count(), 1); + + B::insert(3, 40); + + // Try fail mutate non-existing to existing. + B::try_mutate(2, |query| { + assert_eq!(*query, None); + *query = Some(4); + Result::<(), ()>::Err(()) + }).err().unwrap(); + + assert_eq!(B::try_get(2), Err(())); + assert_eq!(B::count(), 2); + + // Try succeed mutate non-existing to existing. + B::try_mutate(2, |query| { + assert_eq!(*query, None); + *query = Some(41); + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(B::try_get(2), Ok(41)); + assert_eq!(B::count(), 3); + + // Try succeed mutate existing to existing. + B::try_mutate(2, |query| { + assert_eq!(*query, Some(41)); + *query = Some(41); + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(B::try_get(2), Ok(41)); + assert_eq!(B::count(), 3); + + // Try succeed mutate existing to non-existing. + B::try_mutate(2, |query| { + assert_eq!(*query, Some(41)); + *query = None; + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(B::try_get(2), Err(())); + assert_eq!(B::count(), 2); + + B::insert(2, 41); + + // Try fail mutate non-existing to existing. + B::try_mutate_exists(1, |query| { + assert_eq!(*query, None); + *query = Some(4); + Result::<(), ()>::Err(()) + }).err().unwrap(); + + assert_eq!(B::try_get(1), Err(())); + assert_eq!(B::count(), 3); + + // Try succeed mutate non-existing to existing. + B::try_mutate_exists(1, |query| { + assert_eq!(*query, None); + *query = Some(43); + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(B::try_get(1), Ok(43)); + assert_eq!(B::count(), 4); + + // Try succeed mutate existing to existing. + B::try_mutate_exists(1, |query| { + assert_eq!(*query, Some(43)); + *query = Some(43); + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(B::try_get(1), Ok(43)); + assert_eq!(B::count(), 4); + + // Try succeed mutate existing to non-existing. + B::try_mutate_exists(1, |query| { + assert_eq!(*query, Some(43)); + *query = None; + Result::<(), ()>::Ok(()) + }).unwrap(); + + assert_eq!(B::try_get(1), Err(())); + assert_eq!(B::count(), 3); + + // Take exsisting. + assert_eq!(B::take(4), Some(10)); + + assert_eq!(B::try_get(4), Err(())); + assert_eq!(B::count(), 2); + + // Take non-exsisting. + assert_eq!(B::take(4), None); + + assert_eq!(B::try_get(4), Err(())); + assert_eq!(B::count(), 2); + + // Remove all. + B::remove_all(); + + assert_eq!(B::count(), 0); + assert_eq!(B::initialize_counter(), 0); + + B::insert(1, 1); + B::insert(2, 2); + + // Iter values. + assert_eq!(B::iter_values().collect::>(), vec![2, 1]); + + // Iter drain values. + assert_eq!(B::iter_values().drain().collect::>(), vec![2, 1]); + assert_eq!(B::count(), 0); + + B::insert(1, 1); + B::insert(2, 2); + + // Test initialize_counter. + assert_eq!(B::initialize_counter(), 2); + }) + } + + #[test] + fn append_decode_len_works() { + type B = CountedStorageMap>; + + TestExternalities::default().execute_with(|| { + assert_eq!(B::decode_len(0), None); + B::append(0, 3); + assert_eq!(B::decode_len(0), Some(1)); + B::append(0, 3); + assert_eq!(B::decode_len(0), Some(2)); + B::append(0, 3); + assert_eq!(B::decode_len(0), Some(3)); + }) + } + + #[test] + fn try_append_decode_len_works() { + type B = CountedStorageMap>>; + + TestExternalities::default().execute_with(|| { + assert_eq!(B::decode_len(0), None); + B::try_append(0, 3).unwrap(); + assert_eq!(B::decode_len(0), Some(1)); + B::try_append(0, 3).unwrap(); + assert_eq!(B::decode_len(0), Some(2)); + B::try_append(0, 3).unwrap(); + assert_eq!(B::decode_len(0), Some(3)); + B::try_append(0, 3).err().unwrap(); + assert_eq!(B::decode_len(0), Some(3)); + }) + } + + #[test] + fn migrate_keys_works() { + type A = CountedStorageMap; + type B = CountedStorageMap; + TestExternalities::default().execute_with(|| { + A::insert(1, 1); + assert_eq!(B::migrate_key::(1), Some(1)); + assert_eq!(B::get(1), Some(1)); + }) + } + + #[test] + fn translate_values() { + type A = CountedStorageMap; + TestExternalities::default().execute_with(|| { + A::insert(1, 1); + A::insert(2, 2); + A::translate_values::(|old_value| if old_value == 1 { + None + } else { + Some(1) + }); + assert_eq!(A::count(), 1); + assert_eq!(A::get(2), Some(1)); + }) + } + + #[test] + fn test_iter_drain_translate() { + type A = CountedStorageMap; + TestExternalities::default().execute_with(|| { + A::insert(1, 1); + A::insert(2, 2); + + assert_eq!(A::iter().collect::>(), vec![(2, 2), (1, 1)]); + + assert_eq!(A::count(), 2); + + A::translate::(|key, value| if key == 1 { + None + } else { + Some(key as u32 * value) + }); + + assert_eq!(A::count(), 1); + + assert_eq!(A::drain().collect::>(), vec![(2, 4)]); + + assert_eq!(A::count(), 0); + }) + } + +} diff --git a/frame/support/src/storage/types/mod.rs b/frame/support/src/storage/types/mod.rs index f61065671315f..b259e50fbec34 100644 --- a/frame/support/src/storage/types/mod.rs +++ b/frame/support/src/storage/types/mod.rs @@ -26,6 +26,7 @@ mod key; mod map; mod nmap; mod value; +mod counted_map; pub use double_map::{StorageDoubleMap, StorageDoubleMapMetadata}; pub use key::{ @@ -35,6 +36,7 @@ pub use key::{ pub use map::{StorageMap, StorageMapMetadata}; pub use nmap::{StorageNMap, StorageNMapMetadata}; pub use value::{StorageValue, StorageValueMetadata}; +pub use counted_map::{CountedStorageMap, CountedStorageMapMetadata, CountedStorageMapInstance}; /// Trait implementing how the storage optional value is converted into the queried type. /// diff --git a/frame/support/test/tests/pallet.rs b/frame/support/test/tests/pallet.rs index 4f1e66a868947..9134d4af5a3f6 100644 --- a/frame/support/test/tests/pallet.rs +++ b/frame/support/test/tests/pallet.rs @@ -277,6 +277,15 @@ pub mod pallet { u32, >; + #[pallet::storage] + #[pallet::storage_prefix = "RenamedCountedMap"] + #[pallet::getter(fn counted_storage_map)] + pub type SomeCountedStorageMap = CountedStorageMap< + Hasher = Twox64Concat, + Key = u8, + Value = u32, + >; + #[pallet::genesis_config] #[derive(Default)] pub struct GenesisConfig { @@ -801,6 +810,13 @@ fn storage_expand() { pallet::ConditionalDoubleMap::::insert(1, 2, 3); pallet::ConditionalNMap::::insert((1, 2), 3); } + + pallet::SomeCountedStorageMap::::insert(1, 2); + let mut k = [twox_128(b"Example"), twox_128(b"RenamedCountedMap")].concat(); + k.extend(1u8.using_encoded(twox_64_concat)); + assert_eq!(unhashed::get::(&k), Some(2u32)); + let k = [twox_128(b"Example"), twox_128(b"CounterForRenamedCountedMap")].concat(); + assert_eq!(unhashed::get::(&k), Some(1u32)); }) } @@ -1008,6 +1024,27 @@ fn metadata() { default: DecodeDifferent::Decoded(vec![0]), documentation: DecodeDifferent::Decoded(vec![]), }, + StorageEntryMetadata { + name: DecodeDifferent::Decoded("RenamedCountedMap".to_string()), + modifier: StorageEntryModifier::Optional, + ty: StorageEntryType::Map { + key: DecodeDifferent::Decoded("u8".to_string()), + value: DecodeDifferent::Decoded("u32".to_string()), + hasher: StorageHasher::Twox64Concat, + unused: false, + }, + default: DecodeDifferent::Decoded(vec![0]), + documentation: DecodeDifferent::Decoded(vec![]), + }, + StorageEntryMetadata { + name: DecodeDifferent::Decoded("CounterForRenamedCountedMap".to_string()), + modifier: StorageEntryModifier::Default, + ty: StorageEntryType::Plain(DecodeDifferent::Decoded("u32".to_string())), + default: DecodeDifferent::Decoded(vec![0, 0, 0, 0]), + documentation: DecodeDifferent::Decoded(vec![ + "Counter for the related counted storage map".to_string(), + ]), + }, ]), })), calls: Some(DecodeDifferent::Decoded(vec![ @@ -1245,6 +1282,16 @@ fn test_storage_info() { max_size: Some(7 + 16 + 8), } }, + StorageInfo { + prefix: prefix(b"Example", b"RenamedCountedMap"), + max_values: None, + max_size: Some(1 + 4 + 8), + }, + StorageInfo { + prefix: prefix(b"Example", b"CounterForRenamedCountedMap"), + max_values: Some(1), + max_size: Some(4), + }, ], ); } diff --git a/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.rs b/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.rs index d103fa09d991b..5e99c84050c95 100644 --- a/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.rs +++ b/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.rs @@ -1,6 +1,6 @@ #[frame_support::pallet] mod pallet { - use frame_support::pallet_prelude::StorageValue; + use frame_support::pallet_prelude::*; #[pallet::config] pub trait Config: frame_system::Config {} @@ -12,9 +12,15 @@ mod pallet { #[pallet::storage] type Foo = StorageValue<_, u8>; - #[pallet::storage] - #[pallet::storage_prefix = "Foo"] - type NotFoo = StorageValue<_, u16>; + #[pallet::storage] + #[pallet::storage_prefix = "Foo"] + type NotFoo = StorageValue<_, u16>; + + #[pallet::storage] + type CounterForBar = StorageValue<_, u16>; + + #[pallet::storage] + type Bar = CountedStorageMap<_, Twox64Concat, u16, u16>; } fn main() { diff --git a/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.stderr b/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.stderr index 63a6e71e44045..716888c9d8b65 100644 --- a/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.stderr +++ b/frame/support/test/tests/pallet_ui/duplicate_storage_prefix.stderr @@ -1,17 +1,47 @@ error: Duplicate storage prefixes found for `Foo` - --> $DIR/duplicate_storage_prefix.rs:16:32 + --> $DIR/duplicate_storage_prefix.rs:16:29 | 16 | #[pallet::storage_prefix = "Foo"] | ^^^^^ +error: Duplicate storage prefixes found for `Foo` + --> $DIR/duplicate_storage_prefix.rs:13:7 + | +13 | type Foo = StorageValue<_, u8>; + | ^^^ + +error: Duplicate storage prefixes found for `CounterForBar`, used for counter associated to counted storage map + --> $DIR/duplicate_storage_prefix.rs:23:7 + | +23 | type Bar = CountedStorageMap<_, Twox64Concat, u16, u16>; + | ^^^ + +error: Duplicate storage prefixes found for `CounterForBar` + --> $DIR/duplicate_storage_prefix.rs:20:7 + | +20 | type CounterForBar = StorageValue<_, u16>; + | ^^^^^^^^^^^^^ + error[E0412]: cannot find type `_GeneratedPrefixForStorageFoo` in this scope --> $DIR/duplicate_storage_prefix.rs:13:7 | 13 | type Foo = StorageValue<_, u8>; | ^^^ not found in this scope -error[E0121]: the type placeholder `_` is not allowed within types on item signatures - --> $DIR/duplicate_storage_prefix.rs:17:35 +error[E0412]: cannot find type `_GeneratedPrefixForStorageNotFoo` in this scope + --> $DIR/duplicate_storage_prefix.rs:17:7 | 17 | type NotFoo = StorageValue<_, u16>; - | ^ not allowed in type signatures + | ^^^^^^ not found in this scope + +error[E0412]: cannot find type `_GeneratedPrefixForStorageCounterForBar` in this scope + --> $DIR/duplicate_storage_prefix.rs:20:7 + | +20 | type CounterForBar = StorageValue<_, u16>; + | ^^^^^^^^^^^^^ not found in this scope + +error[E0412]: cannot find type `_GeneratedPrefixForStorageBar` in this scope + --> $DIR/duplicate_storage_prefix.rs:23:7 + | +23 | type Bar = CountedStorageMap<_, Twox64Concat, u16, u16>; + | ^^^ not found in this scope