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

Convert impl_key_prefix_for to proc macro #8844

Merged
10 commits merged into from
May 26, 2021
Merged

Convert impl_key_prefix_for to proc macro #8844

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
947a7de
Convert impl_key_prefix_for to proc macro
KiChjang May 18, 2021
bcb9d48
Reduce the number of let bindings
KiChjang May 18, 2021
b60d70c
Remove parsing of inputs for impl_key_prefix_for_tuples
KiChjang May 18, 2021
bc762cd
Replace unwrap with expect
KiChjang May 18, 2021
67c0caf
Remove unnecessary array of idents
KiChjang May 18, 2021
d9f473e
Use numeric identifiers
KiChjang May 19, 2021
432b224
Simplify ident generation
KiChjang May 24, 2021
4ea62d2
Fix whitespacing
KiChjang May 24, 2021
9f2c30b
Add documentation for impl_key_prefix_for_tuple proc macro
KiChjang May 25, 2021
8673078
Merge remote-tracking branch 'origin' into kckyeung/key-prefix-proc-m…
KiChjang May 26, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 75 additions & 80 deletions frame/support/procedural/src/key_prefix.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,97 +15,92 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use proc_macro2::TokenStream;
use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, format_ident, quote};
use std::collections::BTreeSet;
use syn::{Ident, Result, Token, parse::Parser, punctuated::Punctuated};

pub fn impl_key_prefix_for(input: proc_macro::TokenStream) -> Result<TokenStream> {
let args = Punctuated::<Ident, Token![,]>::parse_terminated.parse(input)?;
let arity = args.len();
if arity > 18 {
let msg = "Only supports implementing for tuples up to 18 elements";
return Err(syn::Error::new(args.last().unwrap().span(), msg));
}
use syn::{Ident, Result};

let tuple_set = args
.into_iter()
.try_fold(BTreeSet::new(), |mut acc, ident| {
let span = ident.span();
if !acc.insert(ident) {
return Err(syn::Error::new(span, "Duplicate identifier name"));
}
Ok(acc)
})?;
let mut prefix_set_iter = tuple_set.iter().rev().cloned().peekable();
let mut suffix_set = Vec::new();
const ALL_IDENTS: [&'static str; 18] =
KiChjang marked this conversation as resolved.
Show resolved Hide resolved
["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R"];

pub fn impl_key_prefix_for_tuples(input: proc_macro::TokenStream) -> Result<TokenStream> {
if !input.is_empty() {
return Err(syn::Error::new(Span::call_site(), "No arguments expected"));
}

let mut all_trait_impls = TokenStream::new();

while let Some(next_ident) = prefix_set_iter.next() {
if prefix_set_iter.peek().is_none() {
break;
}
for i in 2..=18 {
KiChjang marked this conversation as resolved.
Show resolved Hide resolved
let current_tuple = ALL_IDENTS[0..i]
.iter()
.map(|s| Ident::new(s, Span::call_site()))
.collect::<Vec<_>>();
let mut prefix_iter = current_tuple.iter().rev().cloned().peekable();
let mut suffix_set = BTreeSet::new();

while let Some(next_ident) = prefix_iter.next() {
if prefix_iter.peek().is_none() {
break;
}

suffix_set.push(next_ident);
let hashers = tuple_set.iter().map(|ident| format_ident!("Hasher{}", ident)).collect::<Vec<_>>();
let kargs = prefix_set_iter.clone().map(|ident| format_ident!("KArg{}", ident)).collect::<Vec<_>>();
let all_idents = &tuple_set;
let prefixes = prefix_set_iter.clone().collect::<Vec<_>>();
let partial_keygen = if suffix_set.len() == arity - 1 {
let key = prefix_set_iter.peek().unwrap();
let hasher = format_ident!("Hasher{}", key);

quote!(Key<#hasher, #key>)
} else {
let keys = prefix_set_iter.clone();
let hashers = keys.clone().map(|ident| format_ident!("Hasher{}", ident));

quote!((#(Key<#hashers, #keys>),*))
};
let suffixes = if suffix_set.len() == 1 {
suffix_set.iter().next().unwrap().to_token_stream()
} else {
let iter = suffix_set.iter();
quote!((#(#iter),*))
};
let suffix_keygen = if suffix_set.len() == 1 {
let key = suffix_set.iter().next().unwrap();
let hasher = format_ident!("Hasher{}", key);

quote!(Key<#hasher, #key>)
} else {
let keys = &suffix_set;
let hashers = keys.iter().map(|ident| format_ident!("Hasher{}", ident));

quote!((#(Key<#hashers, #keys>),*))
};

let trait_impls = quote!{
impl<
#(#all_idents: FullCodec,)*
#(#hashers: StorageHasher,)*
#(#kargs: EncodeLike<#prefixes>),*
> HasKeyPrefix<(#(#kargs,)*)> for (#(Key<#hashers, #all_idents>,)*) {
type Suffix = #suffixes;

fn partial_key(prefix: (#(#kargs,)*)) -> Vec<u8> {
<#partial_keygen>::final_key(prefix)
suffix_set.insert(next_ident);
let hashers = current_tuple.iter().map(|ident| format_ident!("Hasher{}", ident)).collect::<Vec<_>>();
let kargs = prefix_iter.clone().rev().map(|ident| format_ident!("KArg{}", ident)).collect::<Vec<_>>();
let prefixes = prefix_iter.clone().rev().collect::<Vec<_>>();
let partial_keygen = if suffix_set.len() == current_tuple.len() - 1 {
let key = prefix_iter.peek().unwrap();
let hasher = format_ident!("Hasher{}", key);

quote!(Key<#hasher, #key>)
} else {
let keys = prefix_iter.clone().rev();
let hashers = keys.clone().map(|ident| format_ident!("Hasher{}", ident));

quote!((#(Key<#hashers, #keys>),*))
};
let suffixes = if suffix_set.len() == 1 {
suffix_set.iter().next().unwrap().to_token_stream()
} else {
quote!((#(#suffix_set),*))
};
let suffix_keygen = if suffix_set.len() == 1 {
let key = suffix_set.iter().next().unwrap();
let hasher = format_ident!("Hasher{}", key);

quote!(Key<#hasher, #key>)
} else {
let keys = &suffix_set;
let hashers = keys.iter().map(|ident| format_ident!("Hasher{}", ident));

quote!((#(Key<#hashers, #keys>),*))
};

let trait_impls = quote!{
impl<
#(#current_tuple: FullCodec,)*
#(#hashers: StorageHasher,)*
#(#kargs: EncodeLike<#prefixes>),*
> HasKeyPrefix<(#(#kargs,)*)> for (#(Key<#hashers, #current_tuple>,)*) {
KiChjang marked this conversation as resolved.
Show resolved Hide resolved
type Suffix = #suffixes;

fn partial_key(prefix: (#(#kargs,)*)) -> Vec<u8> {
<#partial_keygen>::final_key(prefix)
}
}
}

impl<
#(#all_idents: FullCodec,)*
#(#hashers: ReversibleStorageHasher,)*
#(#kargs: EncodeLike<#prefixes>),*
> HasReversibleKeyPrefix<(#(#kargs,)*)> for (#(Key<#hashers, #all_idents>,)*) {
fn decode_partial_key(key_material: &[u8]) -> Result<Self::Suffix, codec::Error> {
<#suffix_keygen>::decode_final_key(key_material).map(|k| k.0)
impl<
#(#current_tuple: FullCodec,)*
#(#hashers: ReversibleStorageHasher,)*
#(#kargs: EncodeLike<#prefixes>),*
> HasReversibleKeyPrefix<(#(#kargs,)*)> for (#(Key<#hashers, #current_tuple>,)*) {
fn decode_partial_key(key_material: &[u8]) -> Result<Self::Suffix, codec::Error> {
<#suffix_keygen>::decode_final_key(key_material).map(|k| k.0)
}
}
}
};
};

all_trait_impls.extend(trait_impls);
all_trait_impls.extend(trait_impls);
}
}

Ok(all_trait_impls)
Expand Down
4 changes: 2 additions & 2 deletions frame/support/procedural/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -454,6 +454,6 @@ pub fn derive_max_encoded_len(input: TokenStream) -> TokenStream {
}

#[proc_macro]
KiChjang marked this conversation as resolved.
Show resolved Hide resolved
pub fn impl_key_prefix_for(input: TokenStream) -> TokenStream {
key_prefix::impl_key_prefix_for(input).unwrap_or_else(syn::Error::into_compile_error).into()
pub fn impl_key_prefix_for_tuples(input: TokenStream) -> TokenStream {
key_prefix::impl_key_prefix_for_tuples(input).unwrap_or_else(syn::Error::into_compile_error).into()
}
20 changes: 1 addition & 19 deletions frame/support/src/storage/types/key.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -270,22 +270,4 @@ pub trait HasReversibleKeyPrefix<P>: ReversibleKeyGenerator + HasKeyPrefix<P> {
fn decode_partial_key(key_material: &[u8]) -> Result<Self::Suffix, codec::Error>;
}

frame_support_procedural::impl_key_prefix_for!(A, B);
frame_support_procedural::impl_key_prefix_for!(A, B, C);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K, L);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K, L, M);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K, L, M, N);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);
frame_support_procedural::impl_key_prefix_for!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q);
frame_support_procedural::impl_key_prefix_for!(
A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P, Q, R,
);
frame_support_procedural::impl_key_prefix_for_tuples!();