delegation_sdk_attribute_delegate/
lib.rs

1use proc_macro::TokenStream;
2
3use proc_macro2::TokenStream as TokenStream2;
4use quote::quote;
5use syn::{parse_macro_input, ItemMod};
6
7/// This macro attribute is used to inject instructions and struct needed from the delegation program.
8///
9/// Components can be delegate and undelegated to allow fast udpates in the Ephemeral Rollups.
10///
11/// # Example
12/// ```ignore
13///
14/// pub fn delegate(ctx: Context<DelegateInput>) -> Result<()> {
15///     let pda_seeds: &[&[u8]] = &[TEST_PDA_SEED];
16///
17///     let [payer, pda, owner_program, buffer, delegation_record, delegation_metadata, delegation_program, system_program] = [
18///         &ctx.accounts.payer,
19///         &ctx.accounts.pda,
20///         &ctx.accounts.owner_program,
21///         &ctx.accounts.buffer,
22///         &ctx.accounts.buffer,
23///         &ctx.accounts.delegation_record,
24///         &ctx.accounts.delegation_metadata,
25///         &ctx.accounts.delegation_program,
26///         &ctx.accounts.system_program,
27///     ];
28///
29///     delegate_account(
30///         payer,
31///         pda,
32///         owner_program,
33///         buffer,
34///         delegation_record,
35///         delegation_metadata,
36///         delegation_program,
37///         system_program,
38///         pda_seeds,
39///         0,
40///         30000
41///     )?;
42///
43///     Ok(())
44/// }
45///
46/// #[derive(Accounts)]
47/// pub struct DelegateInput<'info> {
48///     pub payer: Signer<'info>,
49///     /// CHECK: The pda to delegate
50///     #[account(mut)]
51///     pub pda: AccountInfo<'info>,
52///     /// CHECK: The owner program of the pda
53///     pub owner_program: AccountInfo<'info>,
54///     /// CHECK: The buffer to temporarily store the pda data
55///     #[account(mut)]
56///     pub buffer: AccountInfo<'info>,
57///     /// CHECK: The delegation record
58///     #[account(mut)]
59///     pub delegation_record: AccountInfo<'info>,
60///     /// CHECK: The delegation account seeds
61///     #[account(mut)]
62///     pub delegation_metadata: AccountInfo<'info>,
63///     /// CHECK: The delegation program
64///     pub delegation_program: AccountInfo<'info>,
65///     /// CHECK: The system program
66///     pub system_program: AccountInfo<'info>,
67/// }
68/// ```
69#[proc_macro_attribute]
70pub fn delegate(_args: TokenStream, input: TokenStream) -> TokenStream {
71    let ast = parse_macro_input!(input as syn::ItemMod);
72    let modified = modify_component_module(ast);
73    TokenStream::from(quote! {
74        #modified
75    })
76}
77
78/// Modifies the component module and adds the necessary functions and structs.
79fn modify_component_module(mut module: ItemMod) -> ItemMod {
80    let (imports, undelegate_fn, undelegate_struct) = generate_undelegate();
81    module.content = module.content.map(|(brace, mut items)| {
82        items.extend(
83            vec![imports, undelegate_fn, undelegate_struct]
84                .into_iter()
85                .map(|item| {
86                    syn::parse2(item).unwrap_or_else(|e| panic!("Failed to parse item: {}", e))
87                })
88                .collect::<Vec<_>>(),
89        );
90        (brace, items)
91    });
92    module
93}
94
95/// Generates the undelegate function and struct.
96fn generate_undelegate() -> (TokenStream2, TokenStream2, TokenStream2) {
97    (
98        quote! {
99            use delegation_program_sdk::undelegate_account;
100        },
101        quote! {
102            #[automatically_derived]
103            pub fn process_undelegation(ctx: Context<InitializeAfterUndelegation>, account_seeds: Vec<Vec<u8>>) -> Result<()> {
104                let [delegated_account, buffer, payer, system_program] = [
105                    &ctx.accounts.base_account,
106                    &ctx.accounts.buffer,
107                    &ctx.accounts.payer,
108                    &ctx.accounts.system_program,
109                ];
110                undelegate_account(
111                    delegated_account,
112                    &id(),
113                    buffer,
114                    payer,
115                    system_program,
116                    account_seeds,
117                )?;
118                Ok(())
119            }
120        },
121        quote! {
122            #[automatically_derived]
123            #[derive(Accounts)]
124                pub struct InitializeAfterUndelegation<'info> {
125                /// CHECK:`
126                #[account(mut)]
127                pub base_account: AccountInfo<'info>,
128                /// CHECK:`
129                #[account()]
130                pub buffer: AccountInfo<'info>,
131                /// CHECK:`
132                #[account(mut)]
133                pub payer: AccountInfo<'info>,
134                /// CHECK:`
135                pub system_program: AccountInfo<'info>,
136            }
137        },
138    )
139}
delegation_sdk_attribute_delegate/lib.rs

delegation_sdk_attribute_delegate/
lib.rs
