Skip to content

Commit

Permalink
feat: Add #[instruction] attribute (#1220)
Browse files Browse the repository at this point in the history
Allow to reference the instruction parameters in `seed =` and
`constraint =` parametes inside `light_account` attribute by exposing
them through `#[instruction]` attribute. It works the same way as
in Anchor.

Code example:

    ```rust
    #[instruction]
    pub struct MyInstruction {
        #[light_account(init, seeds = [b"my-seed", name.as_bytes()])]
        pub my_cpda: MyCpdaType,

        [...]
    }
    ```
  • Loading branch information
vadorovsky authored Sep 16, 2024
1 parent 5586f7e commit a6d3b56
Show file tree
Hide file tree
Showing 3 changed files with 315 additions and 38 deletions.
3 changes: 2 additions & 1 deletion examples/name-service/programs/name-service/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ pub enum CustomError {
}

#[light_accounts]
#[instruction(name: String)]
pub struct CreateRecord<'info> {
#[account(mut)]
#[fee_payer]
Expand All @@ -96,7 +97,7 @@ pub struct CreateRecord<'info> {
#[authority]
pub cpi_signer: AccountInfo<'info>,

#[light_account(init, seeds = [b"name-service", record.name.as_bytes()])]
#[light_account(init, seeds = [b"name-service", name.as_bytes()])]
pub record: LightAccount<NameRecord>,
}

Expand Down
170 changes: 151 additions & 19 deletions macros/light-sdk-macros/src/accounts.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
use proc_macro2::{Span, TokenStream};
use quote::quote;
use quote::{quote, ToTokens};
use syn::{
parse::{Parse, ParseStream},
parse_quote,
punctuated::Punctuated,
token::PathSep,
Error, Expr, Fields, Ident, ItemStruct, Meta, Path, PathSegment, Result, Token, Type, TypePath,
Error, Expr, Fields, Ident, ItemStruct, Meta, Path, PathSegment, Result, Stmt, Token, Type,
TypePath,
};

pub(crate) fn process_light_system_accounts(input: ItemStruct) -> Result<TokenStream> {
Expand Down Expand Up @@ -75,13 +76,83 @@ pub(crate) fn process_light_system_accounts(input: ItemStruct) -> Result<TokenSt
Ok(expanded)
}

struct ParamTypeCheck {
ident: Ident,
ty: Type,
}

impl ToTokens for ParamTypeCheck {
fn to_tokens(&self, tokens: &mut TokenStream) {
let Self { ident, ty } = self;
let stmt: Stmt = parse_quote! {
let #ident: &#ty = #ident;
};
stmt.to_tokens(tokens);
}
}

pub struct InstructionArgs {
param_type_checks: Vec<ParamTypeCheck>,
param_names: Vec<Ident>,
}

impl Parse for InstructionArgs {
fn parse(input: ParseStream) -> Result<Self> {
let mut param_type_checks = Vec::new();
let mut param_names = Vec::new();

while !input.is_empty() {
let ident = input.parse::<Ident>()?;
input.parse::<Token![:]>()?;
let ty = input.parse::<Type>()?;

param_names.push(ident.clone());
param_type_checks.push(ParamTypeCheck { ident, ty });

if input.peek(Token![,]) {
input.parse::<Token![,]>()?;
}
}

Ok(InstructionArgs {
param_type_checks,
param_names,
})
}
}

/// Takes an input struct annotated with `#[light_accounts]` attribute and
/// then:
///
/// - Creates a separate struct with `Light` prefix and moves compressed
/// account fields (annotated with `#[light_account]` attribute) to it. As a
/// result, the original struct, later processed by Anchor macros, contains
/// only regular accounts.
/// - Creates an extention trait, with `LightContextExt` prefix, which serves
/// as an extension to `LightContext` and defines these methods:
/// - `check_constraints`, where the checks extracted from `#[light_account]`
/// attributes are performed.
/// - `derive_address_seeds`, where the seeds extracted from
/// `#[light_account]` attributes are used to derive the address.
pub(crate) fn process_light_accounts(input: ItemStruct) -> Result<TokenStream> {
let mut anchor_accounts_strct = input.clone();

let (_, type_gen, _) = input.generics.split_for_impl();

let anchor_accounts_name = input.ident.clone();
let light_accounts_name = Ident::new(&format!("Light{}", input.ident), Span::call_site());
let ext_trait_name = Ident::new(
&format!("LightContextExt{}", input.ident),
Span::call_site(),
);
let params_name = Ident::new(&format!("Params{}", input.ident), Span::call_site());

let instruction_params = input
.attrs
.iter()
.find(|attribute| attribute.path().is_ident("instruction"))
.map(|attribute| attribute.parse_args::<InstructionArgs>())
.transpose()?;

let mut light_accounts_fields: Punctuated<syn::Field, Token![,]> = Punctuated::new();

Expand All @@ -94,11 +165,18 @@ pub(crate) fn process_light_accounts(input: ItemStruct) -> Result<TokenStream> {
)),
};

// Fields which should belong to the Anchor instruction struct.
let mut anchor_fields = Punctuated::new();
// Names of fields which should belong to the Anchor instruction struct.
let mut anchor_field_idents = Vec::new();
// Names of fields which should belong to the Light instruction struct.
let mut light_field_idents = Vec::new();
// Names of fields of the Light instruction struct, which should be
// available in constraints.
let mut light_referrable_field_idents = Vec::new();
let mut constraint_calls = Vec::new();
let mut derive_address_seed_calls = Vec::new();
let mut set_address_seed_calls = Vec::new();

for field in fields.named.iter() {
let mut light_account = false;
Expand Down Expand Up @@ -135,6 +213,10 @@ pub(crate) fn process_light_accounts(input: ItemStruct) -> Result<TokenStream> {
}
};

if account_args.action != LightAccountAction::Init {
light_referrable_field_idents.push(field.ident.clone());
}

if let Some(constraint) = account_args.constraint {
let Constraint { expr, error } = constraint;
let error = match error {
Expand All @@ -157,8 +239,10 @@ pub(crate) fn process_light_accounts(input: ItemStruct) -> Result<TokenStream> {
&crate::ID,
&unpacked_address_merkle_context,
);
#field_ident.set_address_seed(address_seed);
});
set_address_seed_calls.push(quote! {
#field_ident.set_address_seed(address_seed);
})
} else {
anchor_fields.push(field.clone());
anchor_field_idents.push(field.ident.clone());
Expand All @@ -181,21 +265,60 @@ pub(crate) fn process_light_accounts(input: ItemStruct) -> Result<TokenStream> {
}
};

let light_referrable_fields = if light_referrable_field_idents.is_empty() {
quote! {}
} else {
quote! {
let #light_accounts_name {
#(#light_referrable_field_idents),*, ..
} = &self.light_accounts;
}
};
let input_fields = match instruction_params {
Some(instruction_params) => {
let param_names = instruction_params.param_names;
let param_type_checks = instruction_params.param_type_checks;
quote! {
let #params_name { #(#param_names),*, .. } = inputs;
#(#param_type_checks)*
}
}
None => quote! {},
};

let expanded = quote! {
#[::light_sdk::light_system_accounts]
#[derive(::anchor_lang::Accounts, ::light_sdk::LightTraits)]
#anchor_accounts_strct

#light_accounts_strct

impl<'a, 'b, 'c, 'info> LightContextExt for ::light_sdk::context::LightContext<
pub trait #ext_trait_name {
fn check_constraints(
&self,
inputs: &#params_name,
) -> Result<()>;
fn derive_address_seeds(
&mut self,
address_merkle_context: ::light_sdk::merkle_context::PackedAddressMerkleContext,
inputs: &#params_name,
);
}

impl<'a, 'b, 'c, 'info> #ext_trait_name for ::light_sdk::context::LightContext<
'a, 'b, 'c, 'info, #anchor_accounts_name #type_gen, #light_accounts_name,
> {
#[allow(unused_parens)]
#[allow(unused_variables)]
fn check_constraints(&self) -> Result<()> {
let #anchor_accounts_name { #(#anchor_field_idents),*, .. } = &self.anchor_context.accounts;
let #light_accounts_name { #(#light_field_idents),* } = &self.light_accounts;
fn check_constraints(
&self,
inputs: &#params_name,
) -> Result<()> {
let #anchor_accounts_name {
#(#anchor_field_idents),*, ..
} = &self.anchor_context.accounts;
#light_referrable_fields
#input_fields

#(#constraint_calls)*

Expand All @@ -206,23 +329,31 @@ pub(crate) fn process_light_accounts(input: ItemStruct) -> Result<TokenStream> {
fn derive_address_seeds(
&mut self,
address_merkle_context: PackedAddressMerkleContext,
inputs: &#params_name,
) {
let #anchor_accounts_name { #(#anchor_field_idents),*, .. } = &self.anchor_context.accounts;
let #light_accounts_name { #(#light_field_idents),* } = &mut self.light_accounts;
let #anchor_accounts_name {
#(#anchor_field_idents),*, ..
} = &self.anchor_context.accounts;
#light_referrable_fields
#input_fields

let unpacked_address_merkle_context =
::light_sdk::program_merkle_context::unpack_address_merkle_context(
address_merkle_context, self.anchor_context.remaining_accounts);

#(#derive_address_seed_calls)*

let #light_accounts_name { #(#light_field_idents),* } = &mut self.light_accounts;

#(#set_address_seed_calls)*
}
}
};

Ok(expanded)
}

mod kw {
mod light_account_kw {
// Action
syn::custom_keyword!(init);
syn::custom_keyword!(close);
Expand All @@ -232,6 +363,7 @@ mod kw {
syn::custom_keyword!(seeds);
}

#[derive(Eq, PartialEq)]
pub(crate) enum LightAccountAction {
Init,
Mut,
Expand Down Expand Up @@ -263,20 +395,20 @@ impl Parse for LightAccountArgs {
let lookahead = input.lookahead1();

// Actions
if lookahead.peek(kw::init) {
input.parse::<kw::init>()?;
if lookahead.peek(light_account_kw::init) {
input.parse::<light_account_kw::init>()?;
action = Some(LightAccountAction::Init);
} else if lookahead.peek(Token![mut]) {
input.parse::<Token![mut]>()?;
action = Some(LightAccountAction::Mut);
} else if lookahead.peek(kw::close) {
input.parse::<kw::close>()?;
} else if lookahead.peek(light_account_kw::close) {
input.parse::<light_account_kw::close>()?;
action = Some(LightAccountAction::Close);
}
// Constraint
else if lookahead.peek(kw::constraint) {
else if lookahead.peek(light_account_kw::constraint) {
// Parse the constraint.
input.parse::<kw::constraint>()?;
input.parse::<light_account_kw::constraint>()?;
input.parse::<Token![=]>()?;
let expr: Expr = input.parse()?;

Expand All @@ -290,8 +422,8 @@ impl Parse for LightAccountArgs {
constraint = Some(Constraint { expr, error });
}
// Seeds
else if lookahead.peek(kw::seeds) {
input.parse::<kw::seeds>()?;
else if lookahead.peek(light_account_kw::seeds) {
input.parse::<light_account_kw::seeds>()?;
input.parse::<Token![=]>()?;
seeds = Some(input.parse::<Expr>()?);
} else {
Expand Down Expand Up @@ -427,7 +559,7 @@ pub(crate) fn process_light_accounts_derive(input: ItemStruct) -> Result<TokenSt
)? {
accounts.push(compressed_account);
}
})
});
}

let expanded = quote! {
Expand Down
Loading

0 comments on commit a6d3b56

Please sign in to comment.