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

Commit a6b0d43

Browse files
Mr-LeshiymelekesKiChjanggui1117
authored
CountedNMap implementation (#10621)
* 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 <[email protected]> * Update frame/support/src/storage/types/counted_nmap.rs Co-authored-by: Anton <[email protected]> * Update frame/support/src/storage/types/counted_nmap.rs Co-authored-by: Anton <[email protected]> * 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 <[email protected]> * removed StorageNMap and StoragePrefixedMap traits impl * fix tests * Update frame/support/src/storage/types/counted_nmap.rs Co-authored-by: Guillaume Yu Thiolliere <[email protected]> * 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 <[email protected]> Co-authored-by: Keith Yeung <[email protected]> Co-authored-by: Guillaume Yu Thiolliere <[email protected]> Co-authored-by: parity-processbot <>
1 parent 2e7932f commit a6b0d43

File tree

11 files changed

+1840
-59
lines changed

11 files changed

+1840
-59
lines changed

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

Lines changed: 116 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -194,18 +194,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
194194
let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
195195
args.args.push(syn::GenericArgument::Type(on_empty));
196196
},
197-
StorageGenerics::Map { hasher, key, value, query_kind, on_empty, max_values } => {
198-
args.args.push(syn::GenericArgument::Type(hasher));
199-
args.args.push(syn::GenericArgument::Type(key));
200-
args.args.push(syn::GenericArgument::Type(value.clone()));
201-
let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
202-
set_result_query_type_parameter(&mut query_kind)?;
203-
args.args.push(syn::GenericArgument::Type(query_kind));
204-
let on_empty = on_empty.unwrap_or_else(|| default_on_empty(value));
205-
args.args.push(syn::GenericArgument::Type(on_empty));
206-
let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
207-
args.args.push(syn::GenericArgument::Type(max_values));
208-
},
197+
StorageGenerics::Map { hasher, key, value, query_kind, on_empty, max_values } |
209198
StorageGenerics::CountedMap {
210199
hasher,
211200
key,
@@ -248,7 +237,14 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
248237
let max_values = max_values.unwrap_or_else(|| default_max_values.clone());
249238
args.args.push(syn::GenericArgument::Type(max_values));
250239
},
251-
StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } => {
240+
StorageGenerics::NMap { keygen, value, query_kind, on_empty, max_values } |
241+
StorageGenerics::CountedNMap {
242+
keygen,
243+
value,
244+
query_kind,
245+
on_empty,
246+
max_values,
247+
} => {
252248
args.args.push(syn::GenericArgument::Type(keygen));
253249
args.args.push(syn::GenericArgument::Type(value.clone()));
254250
let mut query_kind = query_kind.unwrap_or_else(|| default_query_kind.clone());
@@ -265,7 +261,7 @@ pub fn process_generics(def: &mut Def) -> syn::Result<Vec<ResultOnEmptyStructMet
265261

266262
let (value_idx, query_idx, on_empty_idx) = match storage_def.metadata {
267263
Metadata::Value { .. } => (1, 2, 3),
268-
Metadata::NMap { .. } => (2, 3, 4),
264+
Metadata::NMap { .. } | Metadata::CountedNMap { .. } => (2, 3, 4),
269265
Metadata::Map { .. } | Metadata::CountedMap { .. } => (3, 4, 5),
270266
Metadata::DoubleMap { .. } => (5, 6, 7),
271267
};
@@ -359,6 +355,17 @@ fn augment_final_docs(def: &mut Def) {
359355
);
360356
push_string_literal(&doc_line, storage);
361357
},
358+
Metadata::CountedNMap { keys, value, .. } => {
359+
let doc_line = format!(
360+
"Storage type is [`CountedStorageNMap`] with keys type ({}) and value type {}.",
361+
keys.iter()
362+
.map(|k| k.to_token_stream().to_string())
363+
.collect::<Vec<_>>()
364+
.join(", "),
365+
value.to_token_stream()
366+
);
367+
push_string_literal(&doc_line, storage);
368+
},
362369
Metadata::CountedMap { key, value } => {
363370
let doc_line = format!(
364371
"Storage type is [`CountedStorageMap`] with key type {} and value type {}.",
@@ -579,6 +586,36 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
579586
}
580587
)
581588
},
589+
Metadata::CountedNMap { keygen, value, .. } => {
590+
let query = match storage.query_kind.as_ref().expect("Checked by def") {
591+
QueryKind::OptionQuery => quote::quote_spanned!(storage.attr_span =>
592+
Option<#value>
593+
),
594+
QueryKind::ResultQuery(error_path, _) => {
595+
quote::quote_spanned!(storage.attr_span =>
596+
Result<#value, #error_path>
597+
)
598+
},
599+
QueryKind::ValueQuery => quote::quote!(#value),
600+
};
601+
quote::quote_spanned!(storage.attr_span =>
602+
#(#cfg_attrs)*
603+
impl<#type_impl_gen> #pallet_ident<#type_use_gen> #completed_where_clause {
604+
#[doc = #getter_doc_line]
605+
pub fn #getter<KArg>(key: KArg) -> #query
606+
where
607+
KArg: #frame_support::storage::types::EncodeLikeTuple<
608+
<#keygen as #frame_support::storage::types::KeyGenerator>::KArg
609+
>
610+
+ #frame_support::storage::types::TupleToEncodedIter,
611+
{
612+
// NOTE: we can't use any trait here because CountedStorageNMap
613+
// doesn't implement any.
614+
<#full_ident>::get(key)
615+
}
616+
}
617+
)
618+
},
582619
}
583620
} else {
584621
Default::default()
@@ -595,40 +632,72 @@ pub fn expand_storages(def: &mut Def) -> proc_macro2::TokenStream {
595632

596633
let cfg_attrs = &storage_def.cfg_attrs;
597634

598-
let maybe_counter = if let Metadata::CountedMap { .. } = storage_def.metadata {
599-
let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident);
600-
let counter_prefix_struct_const = counter_prefix(&prefix_struct_const);
601-
602-
quote::quote_spanned!(storage_def.attr_span =>
603-
#(#cfg_attrs)*
604-
#[doc(hidden)]
605-
#prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>(
606-
core::marker::PhantomData<(#type_use_gen,)>
607-
);
608-
#(#cfg_attrs)*
609-
impl<#type_impl_gen> #frame_support::traits::StorageInstance
610-
for #counter_prefix_struct_ident<#type_use_gen>
611-
#config_where_clause
612-
{
613-
fn pallet_prefix() -> &'static str {
614-
<
615-
<T as #frame_system::Config>::PalletInfo
616-
as #frame_support::traits::PalletInfo
617-
>::name::<Pallet<#type_use_gen>>()
618-
.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
635+
let maybe_counter = match storage_def.metadata {
636+
Metadata::CountedMap { .. } => {
637+
let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident);
638+
let counter_prefix_struct_const = counter_prefix(&prefix_struct_const);
639+
quote::quote_spanned!(storage_def.attr_span =>
640+
#(#cfg_attrs)*
641+
#[doc(hidden)]
642+
#prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>(
643+
core::marker::PhantomData<(#type_use_gen,)>
644+
);
645+
#(#cfg_attrs)*
646+
impl<#type_impl_gen> #frame_support::traits::StorageInstance
647+
for #counter_prefix_struct_ident<#type_use_gen>
648+
#config_where_clause
649+
{
650+
fn pallet_prefix() -> &'static str {
651+
<
652+
<T as #frame_system::Config>::PalletInfo
653+
as #frame_support::traits::PalletInfo
654+
>::name::<Pallet<#type_use_gen>>()
655+
.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
656+
}
657+
const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
619658
}
620-
const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
621-
}
622-
#(#cfg_attrs)*
623-
impl<#type_impl_gen> #frame_support::storage::types::CountedStorageMapInstance
624-
for #prefix_struct_ident<#type_use_gen>
625-
#config_where_clause
626-
{
627-
type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
628-
}
629-
)
630-
} else {
631-
proc_macro2::TokenStream::default()
659+
#(#cfg_attrs)*
660+
impl<#type_impl_gen> #frame_support::storage::types::CountedStorageMapInstance
661+
for #prefix_struct_ident<#type_use_gen>
662+
#config_where_clause
663+
{
664+
type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
665+
}
666+
)
667+
},
668+
Metadata::CountedNMap { .. } => {
669+
let counter_prefix_struct_ident = counter_prefix_ident(&storage_def.ident);
670+
let counter_prefix_struct_const = counter_prefix(&prefix_struct_const);
671+
quote::quote_spanned!(storage_def.attr_span =>
672+
#(#cfg_attrs)*
673+
#[doc(hidden)]
674+
#prefix_struct_vis struct #counter_prefix_struct_ident<#type_use_gen>(
675+
core::marker::PhantomData<(#type_use_gen,)>
676+
);
677+
#(#cfg_attrs)*
678+
impl<#type_impl_gen> #frame_support::traits::StorageInstance
679+
for #counter_prefix_struct_ident<#type_use_gen>
680+
#config_where_clause
681+
{
682+
fn pallet_prefix() -> &'static str {
683+
<
684+
<T as #frame_system::Config>::PalletInfo
685+
as #frame_support::traits::PalletInfo
686+
>::name::<Pallet<#type_use_gen>>()
687+
.expect("No name found for the pallet in the runtime! This usually means that the pallet wasn't added to `construct_runtime!`.")
688+
}
689+
const STORAGE_PREFIX: &'static str = #counter_prefix_struct_const;
690+
}
691+
#(#cfg_attrs)*
692+
impl<#type_impl_gen> #frame_support::storage::types::CountedStorageNMapInstance
693+
for #prefix_struct_ident<#type_use_gen>
694+
#config_where_clause
695+
{
696+
type CounterPrefix = #counter_prefix_struct_ident<#type_use_gen>;
697+
}
698+
)
699+
},
700+
_ => proc_macro2::TokenStream::default(),
632701
};
633702

634703
quote::quote_spanned!(storage_def.attr_span =>

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

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ pub enum Metadata {
138138
CountedMap { value: syn::Type, key: syn::Type },
139139
DoubleMap { value: syn::Type, key1: syn::Type, key2: syn::Type },
140140
NMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
141+
CountedNMap { keys: Vec<syn::Type>, keygen: syn::Type, value: syn::Type },
141142
}
142143

143144
pub enum QueryKind {
@@ -230,6 +231,13 @@ pub enum StorageGenerics {
230231
on_empty: Option<syn::Type>,
231232
max_values: Option<syn::Type>,
232233
},
234+
CountedNMap {
235+
keygen: syn::Type,
236+
value: syn::Type,
237+
query_kind: Option<syn::Type>,
238+
on_empty: Option<syn::Type>,
239+
max_values: Option<syn::Type>,
240+
},
233241
}
234242

235243
impl StorageGenerics {
@@ -242,6 +250,8 @@ impl StorageGenerics {
242250
Self::Value { value, .. } => Metadata::Value { value },
243251
Self::NMap { keygen, value, .. } =>
244252
Metadata::NMap { keys: collect_keys(&keygen)?, keygen, value },
253+
Self::CountedNMap { keygen, value, .. } =>
254+
Metadata::CountedNMap { keys: collect_keys(&keygen)?, keygen, value },
245255
};
246256

247257
Ok(res)
@@ -254,7 +264,8 @@ impl StorageGenerics {
254264
Self::Map { query_kind, .. } |
255265
Self::CountedMap { query_kind, .. } |
256266
Self::Value { query_kind, .. } |
257-
Self::NMap { query_kind, .. } => query_kind.clone(),
267+
Self::NMap { query_kind, .. } |
268+
Self::CountedNMap { query_kind, .. } => query_kind.clone(),
258269
}
259270
}
260271
}
@@ -265,6 +276,7 @@ enum StorageKind {
265276
CountedMap,
266277
DoubleMap,
267278
NMap,
279+
CountedNMap,
268280
}
269281

270282
/// Check the generics in the `map` contains the generics in `gen` may contains generics in
@@ -493,6 +505,29 @@ fn process_named_generics(
493505
max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
494506
}
495507
},
508+
StorageKind::CountedNMap => {
509+
check_generics(
510+
&parsed,
511+
&["Key", "Value"],
512+
&["QueryKind", "OnEmpty", "MaxValues"],
513+
"CountedStorageNMap",
514+
args_span,
515+
)?;
516+
517+
StorageGenerics::CountedNMap {
518+
keygen: parsed
519+
.remove("Key")
520+
.map(|binding| binding.ty)
521+
.expect("checked above as mandatory generic"),
522+
value: parsed
523+
.remove("Value")
524+
.map(|binding| binding.ty)
525+
.expect("checked above as mandatory generic"),
526+
query_kind: parsed.remove("QueryKind").map(|binding| binding.ty),
527+
on_empty: parsed.remove("OnEmpty").map(|binding| binding.ty),
528+
max_values: parsed.remove("MaxValues").map(|binding| binding.ty),
529+
}
530+
},
496531
};
497532

498533
let metadata = generics.metadata()?;
@@ -578,6 +613,16 @@ fn process_unnamed_generics(
578613
false,
579614
)
580615
},
616+
StorageKind::CountedNMap => {
617+
let keygen = retrieve_arg(1)?;
618+
let keys = collect_keys(&keygen)?;
619+
(
620+
None,
621+
Metadata::CountedNMap { keys, keygen, value: retrieve_arg(2)? },
622+
retrieve_arg(3).ok(),
623+
false,
624+
)
625+
},
581626
};
582627

583628
Ok(res)
@@ -594,10 +639,11 @@ fn process_generics(
594639
"CountedStorageMap" => StorageKind::CountedMap,
595640
"StorageDoubleMap" => StorageKind::DoubleMap,
596641
"StorageNMap" => StorageKind::NMap,
642+
"CountedStorageNMap" => StorageKind::CountedNMap,
597643
found => {
598644
let msg = format!(
599645
"Invalid pallet::storage, expected ident: `StorageValue` or \
600-
`StorageMap` or `CountedStorageMap` or `StorageDoubleMap` or `StorageNMap` \
646+
`StorageMap` or `CountedStorageMap` or `StorageDoubleMap` or `StorageNMap` or `CountedStorageNMap` \
601647
in order to expand metadata, found `{}`.",
602648
found,
603649
);

‎frame/support/src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1545,8 +1545,8 @@ pub mod pallet_prelude {
15451545
storage::{
15461546
bounded_vec::BoundedVec,
15471547
types::{
1548-
CountedStorageMap, Key as NMapKey, OptionQuery, ResultQuery, StorageDoubleMap,
1549-
StorageMap, StorageNMap, StorageValue, ValueQuery,
1548+
CountedStorageMap, CountedStorageNMap, Key as NMapKey, OptionQuery, ResultQuery,
1549+
StorageDoubleMap, StorageMap, StorageNMap, StorageValue, ValueQuery,
15501550
},
15511551
StorageList,
15521552
},

‎frame/support/src/storage/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1428,6 +1428,7 @@ mod private {
14281428
impl<K, V, S> Sealed for bounded_btree_map::BoundedBTreeMap<K, V, S> {}
14291429
impl<T, S> Sealed for bounded_btree_set::BoundedBTreeSet<T, S> {}
14301430
impl<T: Encode> Sealed for BTreeSet<T> {}
1431+
impl<'a, T: EncodeLike<U>, U: Encode> Sealed for codec::Ref<'a, T, U> {}
14311432

14321433
macro_rules! impl_sealed_for_tuple {
14331434
($($elem:ident),+) => {

‎frame/support/src/storage/types/counted_map.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -612,8 +612,9 @@ mod test {
612612
assert_eq!(A::count(), 2);
613613

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

617+
assert_eq!(A::try_get(3), Ok(12));
617618
assert_eq!(A::count(), 2);
618619

619620
// Remove non-existing.
@@ -706,17 +707,17 @@ mod test {
706707
// Try succeed mutate existing to existing.
707708
A::try_mutate_exists(1, |query| {
708709
assert_eq!(*query, Some(43));
709-
*query = Some(43);
710+
*query = Some(45);
710711
Result::<(), ()>::Ok(())
711712
})
712713
.unwrap();
713714

714-
assert_eq!(A::try_get(1), Ok(43));
715+
assert_eq!(A::try_get(1), Ok(45));
715716
assert_eq!(A::count(), 4);
716717

717718
// Try succeed mutate existing to non-existing.
718719
A::try_mutate_exists(1, |query| {
719-
assert_eq!(*query, Some(43));
720+
assert_eq!(*query, Some(45));
720721
*query = None;
721722
Result::<(), ()>::Ok(())
722723
})

0 commit comments

Comments
 (0)