bolt_attribute_bolt_delegate/
lib.rs

1use proc_macro::TokenStream;
2
3use proc_macro2::TokenStream as TokenStream2;
4use quote::quote;
5use syn::{parse_macro_input, AttributeArgs, ItemMod, NestedMeta, Type};
6
7/// This macro attribute is used to inject instructions and struct needed to delegate BOLT component.
8///
9/// Components can be delegate in order to be updated in an Ephemeral Rollup
10///
11/// # Example
12/// ```ignore
13///
14/// #[component(delegate)]
15/// pub struct Position {
16///     pub x: i64,
17///     pub y: i64,
18///     pub z: i64,
19/// }
20/// ```
21#[proc_macro_attribute]
22pub fn delegate(args: TokenStream, input: TokenStream) -> TokenStream {
23    let ast = parse_macro_input!(input as syn::ItemMod);
24    let args = parse_macro_input!(args as syn::AttributeArgs);
25    let component_type =
26        extract_type_name(&args).expect("Expected a component type in macro arguments");
27    let modified = modify_component_module(ast, &component_type);
28    TokenStream::from(quote! {
29        #modified
30    })
31}
32
33/// Modifies the component module and adds the necessary functions and structs.
34fn modify_component_module(mut module: ItemMod, component_type: &Type) -> ItemMod {
35    let (delegate_fn, delegate_struct) = generate_delegate(component_type);
36    let (reinit_undelegate_fn, reinit_undelegate_struct) = generate_reinit_after_undelegate();
37    let (undelegate_fn, undelegate_struct) = generate_undelegate();
38    module.content = module.content.map(|(brace, mut items)| {
39        items.extend(
40            vec![
41                delegate_fn,
42                delegate_struct,
43                reinit_undelegate_fn,
44                reinit_undelegate_struct,
45                undelegate_fn,
46                undelegate_struct,
47            ]
48            .into_iter()
49            .map(|item| syn::parse2(item).unwrap())
50            .collect::<Vec<_>>(),
51        );
52        (brace, items)
53    });
54    module
55}
56
57/// Generates the allow_undelegate function and struct.
58fn generate_undelegate() -> (TokenStream2, TokenStream2) {
59    (
60        quote! {
61            #[automatically_derived]
62            pub fn undelegate(ctx: Context<Undelegate>) -> Result<()> {
63                ::bolt_lang::commit_and_undelegate_accounts(
64                    &ctx.accounts.payer,
65                    vec![&ctx.accounts.delegated_account.to_account_info()],
66                    &ctx.accounts.magic_context,
67                    &ctx.accounts.magic_program,
68                )?;
69                Ok(())
70            }
71        },
72        quote! {
73            #[automatically_derived]
74            #[derive(Accounts)]
75            pub struct Undelegate<'info> {
76                #[account(mut)]
77                pub payer: Signer<'info>,
78                #[account(mut)]
79                /// CHECK: The delegated component
80                pub delegated_account: AccountInfo<'info>,
81                #[account(mut, address = ::bolt_lang::MAGIC_CONTEXT_ID)]
82                /// CHECK:`
83                pub magic_context: AccountInfo<'info>,
84                #[account()]
85                /// CHECK:`
86                pub magic_program: Program<'info, MagicProgram>
87            }
88        },
89    )
90}
91
92/// Generates the undelegate function and struct.
93fn generate_reinit_after_undelegate() -> (TokenStream2, TokenStream2) {
94    (
95        quote! {
96            #[automatically_derived]
97            pub fn process_undelegation(ctx: Context<InitializeAfterUndelegation>, account_seeds: Vec<Vec<u8>>) -> Result<()> {
98                let [delegated_account, buffer, payer, system_program] = [
99                    &ctx.accounts.delegated_account,
100                    &ctx.accounts.buffer,
101                    &ctx.accounts.payer,
102                    &ctx.accounts.system_program,
103                ];
104                ::bolt_lang::undelegate_account(
105                    delegated_account,
106                    &id(),
107                    buffer,
108                    payer,
109                    system_program,
110                    account_seeds,
111                )?;
112                Ok(())
113            }
114        },
115        quote! {
116            #[automatically_derived]
117            #[derive(Accounts)]
118            pub struct InitializeAfterUndelegation<'info> {
119                /// CHECK:`
120                #[account(mut)]
121                pub delegated_account: AccountInfo<'info>,
122                /// CHECK:`
123                #[account()]
124                pub buffer: AccountInfo<'info>,
125                /// CHECK:
126                #[account(mut)]
127                pub payer: AccountInfo<'info>,
128                /// CHECK:
129                pub system_program: AccountInfo<'info>,
130            }
131        },
132    )
133}
134
135/// Generates the delegate instruction and related structs to inject in the component.
136fn generate_delegate(component_type: &Type) -> (TokenStream2, TokenStream2) {
137    (
138        quote! {
139            #[automatically_derived]
140            pub fn delegate(ctx: Context<DelegateInput>, commit_frequency_ms: u32, validator: Option<Pubkey>) -> Result<()> {
141                let pda_seeds: &[&[u8]] = &[<#component_type>::seed(), &ctx.accounts.entity.key().to_bytes()];
142
143                let del_accounts = ::bolt_lang::DelegateAccounts {
144                    payer: &ctx.accounts.payer,
145                    pda: &ctx.accounts.account,
146                    owner_program: &ctx.accounts.owner_program,
147                    buffer: &ctx.accounts.buffer,
148                    delegation_record: &ctx.accounts.delegation_record,
149                    delegation_metadata: &ctx.accounts.delegation_metadata,
150                    delegation_program: &ctx.accounts.delegation_program,
151                    system_program: &ctx.accounts.system_program,
152                };
153
154                let config = ::bolt_lang::DelegateConfig {
155                    commit_frequency_ms,
156                    validator,
157                };
158
159                ::bolt_lang::delegate_account(
160                    del_accounts,
161                    pda_seeds,
162                    config,
163                )?;
164
165                Ok(())
166            }
167        },
168        quote! {
169            #[automatically_derived]
170            #[derive(Accounts)]
171            pub struct DelegateInput<'info> {
172                pub payer: Signer<'info>,
173                #[account()]
174                pub entity: Account<'info, Entity>,
175                /// CHECK:
176                #[account(mut)]
177                pub account: AccountInfo<'info>,
178                /// CHECK:`
179                pub owner_program: AccountInfo<'info>,
180                /// CHECK:
181                #[account(mut)]
182                pub buffer: AccountInfo<'info>,
183                /// CHECK:`
184                #[account(mut)]
185                pub delegation_record: AccountInfo<'info>,
186                /// CHECK:`
187                #[account(mut)]
188                pub delegation_metadata: AccountInfo<'info>,
189                /// CHECK:`
190                pub delegation_program: AccountInfo<'info>,
191                /// CHECK:`
192                pub system_program: AccountInfo<'info>,
193            }
194        },
195    )
196}
197
198/// Extracts the type name from attribute arguments.
199fn extract_type_name(args: &AttributeArgs) -> Option<Type> {
200    args.iter().find_map(|arg| {
201        if let NestedMeta::Meta(syn::Meta::Path(path)) = arg {
202            Some(Type::Path(syn::TypePath {
203                qself: None,
204                path: path.clone(),
205            }))
206        } else {
207            None
208        }
209    })
210}