Skip to content
This repository has been archived by the owner on Nov 15, 2023. It is now read-only.

Commit

Permalink
CountedNMap implementation (#10621)
Browse files Browse the repository at this point in the history 
* add initial CountedDoubleMap implementation

* extend CountedDoubleMap functionality

* add some traits implementation for CountedStorageDoubleMap

* add basic tests for CountedStorageDoubleMap

* add mutate functions implementation

* add additional tests

* add test_option_query test

* add try_append_decode_len_works, append_decode_len_works tests

* add migrate_keys_works, translate_values tests

* add test_iter_drain_translate test

* add test_metadata test

* add remove_prefix implementation,  add test_iter_drain_prefix test

* update

* refactor PrefixIterator usage

* Fix CI build

* fix storage_ensure_span_are_ok_wrong_gen.rs storage_ensure_span_are_ok_wrong_gen_unnamed.rs

* add counted_nmap implementation

* add tests, fixes

* remove counted double map impl

* fix metadata checks

* update clear func

* fix clear, clear with prefix

* fix set function

* update

* final fix

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Anton <anton.kalyaev@gmail.com>

* fix comments

* fix suggestion

* cargo update

* Relocate impl of Sealed for Ref to module root

* fix StorageEntryMetadata type

* Update frame/support/src/storage/types/nmap.rs

Co-authored-by: Guillaume Yu Thiolliere <gui.thiolliere@gmail.com>

* removed StorageNMap and StoragePrefixedMap traits impl

* fix tests

* Update frame/support/src/storage/types/counted_nmap.rs

Co-authored-by: Guillaume Yu Thiolliere <gui.thiolliere@gmail.com>

* extend pallet::storage macro with CountedStorageNMap usage

* fix

* add tests

* fix

* fix

* Add counter_storage_final_key(), map_storage_final_prefix() functions

* update tests

* fix

* fix

* fix

* update tests

* fix fmt

* fix fmt

---------

Co-authored-by: Anton <anton.kalyaev@gmail.com>
Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>
Co-authored-by: Guillaume Yu Thiolliere <gui.thiolliere@gmail.com>
Co-authored-by: parity-processbot <>
  • Loading branch information
@Mr-Leshiy @melekes
@KiChjang @gui1117
4 people authored Aug 4, 2023
1 parent 2e7932f commit a6b0d43
Show file tree
Hide file tree
Showing 11 changed files with 1,840 additions and 59 deletions.
163 changes: 116 additions & 47 deletions frame/support/procedural/src/pallet/expand/storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,18 +194,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
args.args.push(syn::GenericArgument::Type(on_empty));
},
StorageGenerics::Map { 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.clone()));
let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
set_result_query_type_parameter(&mut query_kind)?;
args.args.push(syn::GenericArgument::Type(query_kind));
let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
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::Map { hasher, key, value, query_kind, on_empty, max_values } |
StorageGenerics::CountedMap {
hasher,
key,
Expand Down Expand Up @@ -248,7 +237,14 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
args.args.push(syn::GenericArgument::Type(max_values));
},
StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } => {
StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } |
StorageGenerics::CountedNMap {
keygen,
value,
query_kind,
on_empty,
max_values,
} => {
args.args.push(syn::GenericArgument::Type(keygen));
args.args.push(syn::GenericArgument::Type(value.clone()));
let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
Expand All @@ -265,7 +261,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet

let (value_idx, query_idx, on_empty_idx) = match storage_def.metadata {
Metadata::Value { .. } => (1, 2, 3),
Metadata::NMap { .. } => (2, 3, 4),
Metadata::NMap { .. } | Metadata::CountedNMap { .. } => (2, 3, 4),
Metadata::Map { .. } | Metadata::CountedMap { .. } => (3, 4, 5),
Metadata::DoubleMap { .. } => (5, 6, 7),
};
Expand Down Expand Up @@ -359,6 +355,17 @@ fn augment_final_docs(def: &mut Def) {
);
push_string_literal(&doc_line, storage);
},
Metadata::CountedNMap { keys, value, .. } => {
let doc_line = format!(
"Storage type is [`CountedStorageNMap`] with keys type ({}) and value type {}.",
keys.iter()
.map(|k| k.to_token_stream().to_string())
.collect::<Vec<_>>()
.join(", "),
value.to_token_stream()
);
push_string_literal(&doc_line, storage);
},
Metadata::CountedMap { key, value } => {
let doc_line = format!(
"Storage type is [`CountedStorageMap`] with key type {} and value type {}.",
Expand Down Expand Up @@ -579,6 +586,36 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
}
)
},
Metadata::CountedNMap { keygen, value, .. } => {
let query = match storage.query_kind.as_ref().expect("Checked by def") {
QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
Option<#value>
),
QueryKind::ResultQuery(error_path, _) => {
quote::quote_spanned!(storage.attr_span =>
Result<#value, #error_path>
)
},
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 {
#[doc = #getter_doc_line]
pub fn #getter<KArg>(key: KArg) -> #query
where
KArg: #frame_support::storage::types::EncodeLikeTuple<
<#keygen as #frame_support::storage::types::KeyGenerator>::KArg
>
+ #frame_support::storage::types::TupleToEncodedIter,
{
// NOTE: we can't use any trait here because CountedStorageNMap
// doesn't implement any.
<#full_ident>::get(key)
}
}
)
},
}
} else {
Default::default()
Expand All @@ -595,40 +632,72 @@ 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)*
#[doc(hidden)]
#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 {
<
<T as #frame_system::Config>::PalletInfo
as #frame_support::traits::PalletInfo
>::name::<Pallet<#type_use_gen>>()
.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
let maybe_counter = match storage_def.metadata {
Metadata::CountedMap { .. } => {
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)*
#[doc(hidden)]
#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 {
<
<T as #frame_system::Config>::PalletInfo
as #frame_support::traits::PalletInfo
>::name::<Pallet<#type_use_gen>>()
.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
}
const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
}
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()
#(#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>;
}
)
},
Metadata::CountedNMap { .. } => {
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)*
#[doc(hidden)]
#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 {
<
<T as #frame_system::Config>::PalletInfo
as #frame_support::traits::PalletInfo
>::name::<Pallet<#type_use_gen>>()
.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
}
const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
}
#(#cfg_attrs)*
impl<#type_impl_gen> #frame_support::storage::types::CountedStorageNMapInstance
for #prefix_struct_ident<#type_use_gen>
#config_where_clause
{
type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
}
)
},
_ => proc_macro2::TokenStream::default(),
};

quote::quote_spanned!(storage_def.attr_span =>
Expand Down
50 changes: 48 additions & 2 deletions frame/support/procedural/src/pallet/parse/storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ pub enum Metadata {
CountedMap { value: syn::Type, key: syn::Type },
DoubleMap { value: syn::Type, key1: syn::Type, key2: syn::Type },
NMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
CountedNMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
}

pub enum QueryKind {
Expand Down Expand Up @@ -230,6 +231,13 @@ pub enum StorageGenerics {
on_empty: Option<syn::Type>,
max_values: Option<syn::Type>,
},
CountedNMap {
keygen: syn::Type,
value: syn::Type,
query_kind: Option<syn::Type>,
on_empty: Option<syn::Type>,
max_values: Option<syn::Type>,
},
}

impl StorageGenerics {
Expand All @@ -242,6 +250,8 @@ impl StorageGenerics {
Self::Value { value, .. } => Metadata::Value { value },
Self::NMap { keygen, value, .. } =>
Metadata::NMap { keys: collect_keys(&keygen)?, keygen, value },
Self::CountedNMap { keygen, value, .. } =>
Metadata::CountedNMap { keys: collect_keys(&keygen)?, keygen, value },
};

Ok(res)
Expand All @@ -254,7 +264,8 @@ impl StorageGenerics {
Self::Map { query_kind, .. } |
Self::CountedMap { query_kind, .. } |
Self::Value { query_kind, .. } |
Self::NMap { query_kind, .. } => query_kind.clone(),
Self::NMap { query_kind, .. } |
Self::CountedNMap { query_kind, .. } => query_kind.clone(),
}
}
}
Expand All @@ -265,6 +276,7 @@ enum StorageKind {
CountedMap,
DoubleMap,
NMap,
CountedNMap,
}

/// Check the generics in the `map` contains the generics in `gen` may contains generics in
Expand Down Expand Up @@ -493,6 +505,29 @@ fn process_named_generics(
max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
}
},
StorageKind::CountedNMap => {
check_generics(
&parsed,
&["Key", "Value"],
&["QueryKind", "OnEmpty", "MaxValues"],
"CountedStorageNMap",
args_span,
)?;

StorageGenerics::CountedNMap {
keygen: 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),
}
},
};

let metadata = generics.metadata()?;
Expand Down Expand Up @@ -578,6 +613,16 @@ fn process_unnamed_generics(
false,
)
},
StorageKind::CountedNMap => {
let keygen = retrieve_arg(1)?;
let keys = collect_keys(&keygen)?;
(
None,
Metadata::CountedNMap { keys, keygen, value: retrieve_arg(2)? },
retrieve_arg(3).ok(),
false,
)
},
};

Ok(res)
Expand All @@ -594,10 +639,11 @@ fn process_generics(
"CountedStorageMap" => StorageKind::CountedMap,
"StorageDoubleMap" => StorageKind::DoubleMap,
"StorageNMap" => StorageKind::NMap,
"CountedStorageNMap" => StorageKind::CountedNMap,
found => {
let msg = format!(
"Invalid pallet::storage, expected ident: `StorageValue` or \
`StorageMap` or `CountedStorageMap` or `StorageDoubleMap` or `StorageNMap` \
`StorageMap` or `CountedStorageMap` or `StorageDoubleMap` or `StorageNMap` or `CountedStorageNMap` \
in order to expand metadata, found `{}`.",
found,
);
Expand Down
4 changes: 2 additions & 2 deletions frame/support/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1545,8 +1545,8 @@ pub mod pallet_prelude {
storage::{
bounded_vec::BoundedVec,
types::{
CountedStorageMap, Key as NMapKey, OptionQuery, ResultQuery, StorageDoubleMap,
StorageMap, StorageNMap, StorageValue, ValueQuery,
CountedStorageMap, CountedStorageNMap, Key as NMapKey, OptionQuery, ResultQuery,
StorageDoubleMap, StorageMap, StorageNMap, StorageValue, ValueQuery,
},
StorageList,
},
Expand Down
1 change: 1 addition & 0 deletions frame/support/src/storage/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1428,6 +1428,7 @@ mod private {
impl<K, V, S> Sealed for bounded_btree_map::BoundedBTreeMap<K, V, S> {}
impl<T, S> Sealed for bounded_btree_set::BoundedBTreeSet<T, S> {}
impl<T: Encode> Sealed for BTreeSet<T> {}
impl<'a, T: EncodeLike<U>, U: Encode> Sealed for codec::Ref<'a, T, U> {}

macro_rules! impl_sealed_for_tuple {
($($elem:ident),+) => {
Expand Down
9 changes: 5 additions & 4 deletions frame/support/src/storage/types/counted_map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -612,8 +612,9 @@ mod test {
assert_eq!(A::count(), 2);

// Insert an existing key, shouldn't increment counted values.
A::insert(3, 11);
A::insert(3, 12);

assert_eq!(A::try_get(3), Ok(12));
assert_eq!(A::count(), 2);

// Remove non-existing.
Expand Down Expand Up @@ -706,17 +707,17 @@ mod test {
// Try succeed mutate existing to existing.
A::try_mutate_exists(1, |query| {
assert_eq!(*query, Some(43));
*query = Some(43);
*query = Some(45);
Result::<(), ()>::Ok(())
})
.unwrap();

assert_eq!(A::try_get(1), Ok(43));
assert_eq!(A::try_get(1), Ok(45));
assert_eq!(A::count(), 4);

// Try succeed mutate existing to non-existing.
A::try_mutate_exists(1, |query| {
assert_eq!(*query, Some(43));
assert_eq!(*query, Some(45));
*query = None;
Result::<(), ()>::Ok(())
})
Expand Down
Loading

0 comments on commit a6b0d43

Please sign in to comment.