Skip to content

Commit

Permalink
feat: Add #[instruction] attribute
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 committed Sep 16, 2024
1 parent 5586f7e commit bf9d765
Show file tree
Hide file tree
Showing 3 changed files with 309 additions and 37 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
163 changes: 145 additions & 18 deletions macros/light-sdk-macros/src/accounts.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
use proc_macro2::{Span, TokenStream};
use quote::quote;
use quote::{quote, ToTokens};
use syn::{
parse::{Parse, ParseStream},
parse_quote,
punctuated::Punctuated,
spanned::Spanned,
token::PathSep,
Error, Expr, Fields, Ident, ItemStruct, Meta, Path, PathSegment, Result, Token, Type, TypePath,
};
Expand Down Expand Up @@ -75,13 +76,86 @@ pub(crate) fn process_light_system_accounts(input: ItemStruct) -> Result<TokenSt
Ok(expanded)
}

struct InstructionParam {
ident: Ident,
// TODO(vadorovsky): Actually use that field to validate whether types
// provided through `#[instruction]` match the actual arguments in the
// instruction function.
// For now, we just parse that type for syntax compatibility with Anchor.
_ty: Type,
}

impl ToTokens for InstructionParam {
fn to_tokens(&self, tokens: &mut TokenStream) {
self.ident.to_tokens(tokens);
}
}

pub(crate) struct InstructionArgs(Vec<InstructionParam>);

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

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

instructions.push(InstructionParam { ident, _ty: ty });

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

Ok(InstructionArgs(instructions))
}
}

impl ToTokens for InstructionArgs {
fn to_tokens(&self, tokens: &mut TokenStream) {
for param in self.0.iter() {
param.to_tokens(tokens);
Token![,](param.span()).to_tokens(tokens);
}
}
}

/// 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.
/// - ``
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()?;

// .map_or(Vec::new(), |args| args.instruction_params);

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

Expand All @@ -94,11 +168,15 @@ 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 fi
let mut anchor_field_idents = Vec::new();
let mut light_field_idents = Vec::new();
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,55 @@ 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) => quote! {
let #params_name { #instruction_params .. } = inputs;
},
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 +324,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 +358,7 @@ mod kw {
syn::custom_keyword!(seeds);
}

#[derive(Eq, PartialEq)]
pub(crate) enum LightAccountAction {
Init,
Mut,
Expand Down Expand Up @@ -263,20 +390,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 +417,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 +554,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 bf9d765

Please sign in to comment.