lang: Optimize associated constraint check (#425)

This commit is contained in:
Kirill Fomichev 2021-06-27 20:03:26 +03:00 committed by GitHub
parent 96c898ddb4
commit 609c836d63
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 47 additions and 33 deletions

View File

@ -316,26 +316,25 @@ pub fn generate_constraint_associated_init(
}, },
}; };
let associated_seeds_constraint = generate_constraint_associated_seeds(f, c); let associated_seeds_constraint = generate_constraint_associated_seeds(f, c);
let seeds_with_nonce = match c.associated_seeds.len() { let seeds_with_nonce = if c.associated_seeds.is_empty() {
0 => quote! { quote! {
#associated_seeds_constraint #associated_seeds_constraint
let seeds = [ let seeds = [
&b"anchor"[..], &b"anchor"[..],
#associated_target.to_account_info().key.as_ref(), #associated_target.to_account_info().key.as_ref(),
&[nonce], &[nonce],
]; ];
}, }
_ => { } else {
let seeds = to_seeds_tts(&c.associated_seeds); let seeds = to_seeds_tts(&c.associated_seeds);
quote! { quote! {
#associated_seeds_constraint #associated_seeds_constraint
let seeds = [ let seeds = [
&b"anchor"[..], &b"anchor"[..],
#associated_target.to_account_info().key.as_ref(), #associated_target.to_account_info().key.as_ref(),
#seeds #seeds
&[nonce], &[nonce],
]; ];
}
} }
}; };
generate_pda(f, seeds_with_nonce, payer, &c.space, true) generate_pda(f, seeds_with_nonce, payer, &c.space, true)
@ -443,29 +442,44 @@ pub fn generate_constraint_associated_seeds(
) -> proc_macro2::TokenStream { ) -> proc_macro2::TokenStream {
let field = &f.ident; let field = &f.ident;
let associated_target = c.associated_target.clone(); let associated_target = c.associated_target.clone();
let seeds_no_nonce = match c.associated_seeds.len() { let seeds_no_nonce = if c.associated_seeds.is_empty() {
0 => quote! { quote! {
[ &b"anchor"[..],
&b"anchor"[..], #associated_target.to_account_info().key.as_ref(),
#associated_target.to_account_info().key.as_ref(), }
] } else {
}, let seeds = to_seeds_tts(&c.associated_seeds);
_ => { quote! {
let seeds = to_seeds_tts(&c.associated_seeds); &b"anchor"[..],
quote! { #associated_target.to_account_info().key.as_ref(),
[ #seeds
&b"anchor"[..], }
#associated_target.to_account_info().key.as_ref(), };
#seeds let associated_field = if c.is_init {
] quote! {
let (__associated_field, nonce) = Pubkey::find_program_address(
&[#seeds_no_nonce],
program_id,
);
}
} else {
let nonce = match &f.ty {
Ty::ProgramAccount(_) => quote! { #field.__nonce },
Ty::Loader(_) => {
// Zero copy is not deserialized, so the data must be lazy loaded.
quote! { #field.load()?.__nonce }
} }
_ => panic!("Invalid type for initializing a program derived address"),
};
quote! {
let __associated_field = Pubkey::create_program_address(
&[#seeds_no_nonce &[#nonce]],
program_id,
)?;
} }
}; };
quote! { quote! {
let (__associated_field, nonce) = Pubkey::find_program_address( #associated_field
&#seeds_no_nonce,
program_id,
);
if &__associated_field != #field.to_account_info().key { if &__associated_field != #field.to_account_info().key {
return Err(anchor_lang::__private::ErrorCode::ConstraintAssociatedInit.into()); return Err(anchor_lang::__private::ErrorCode::ConstraintAssociatedInit.into());
} }