♻️ Reworking delegation gen macro and making it available for #[bundle]

Author: notdanilo

Cargo.lock

@@ -12,6 +12,4 @@ edition = { workspace = true }
 proc-macro = true
 
 [dependencies]
-syn = { workspace = true }
-quote = { workspace = true }
-proc-macro2 = { workspace = true }
+bolt-attribute.workspace = true

crates/bolt-lang/attribute/delegate/Cargo.toml

@@ -1,210 +1,18 @@
 use proc_macro::TokenStream;
 
-use proc_macro2::TokenStream as TokenStream2;
-use quote::quote;
-use syn::{parse_macro_input, AttributeArgs, ItemMod, NestedMeta, Type};
-
-/// This macro attribute is used to inject instructions and struct needed to delegate BOLT component.
+/// This macro attribute is used to inject instructions and struct needed to delegate bolt components.
 ///
-/// Components can be delegate in order to be updated in an Ephemeral Rollup
+/// Components can be delegated in order to be updated in an Ephemeral Rollup validator.
 ///
 /// # Example
 /// ```ignore
 ///
-/// #[component(delegate)]
-/// pub struct Position {
-///     pub x: i64,
-///     pub y: i64,
-///     pub z: i64,
+/// #[delegate]
+/// #[anchor_lang::program]
+/// mod program {
 /// }
 /// ```
 #[proc_macro_attribute]
 pub fn delegate(args: TokenStream, input: TokenStream) -> TokenStream {
-    let ast = parse_macro_input!(input as syn::ItemMod);
-    let args = parse_macro_input!(args as syn::AttributeArgs);
-    let component_type =
-        extract_type_name(&args).expect("Expected a component type in macro arguments");
-    let modified = modify_component_module(ast, &component_type);
-    TokenStream::from(quote! {
-        #modified
-    })
-}
-
-/// Modifies the component module and adds the necessary functions and structs.
-fn modify_component_module(mut module: ItemMod, component_type: &Type) -> ItemMod {
-    let (delegate_fn, delegate_struct) = generate_delegate(component_type);
-    let (process_undelegation_fn, process_undelegation_struct) = generate_process_undelegation();
-    let (undelegate_fn, undelegate_struct) = generate_undelegate();
-    module.content = module.content.map(|(brace, mut items)| {
-        items.extend(
-            vec![
-                delegate_fn,
-                delegate_struct,
-                process_undelegation_fn,
-                process_undelegation_struct,
-                undelegate_fn,
-                undelegate_struct,
-            ]
-            .into_iter()
-            .map(|item| syn::parse2(item).unwrap())
-            .collect::<Vec<_>>(),
-        );
-        (brace, items)
-    });
-    module
-}
-
-/// Generates the allow_undelegate function and struct.
-fn generate_undelegate() -> (TokenStream2, TokenStream2) {
-    (
-        quote! {
-            #[automatically_derived]
-            pub fn undelegate(ctx: Context<Undelegate>) -> Result<()> {
-                ::bolt_lang::commit_and_undelegate_accounts(
-                    &ctx.accounts.payer,
-                    vec![&ctx.accounts.delegated_account.to_account_info()],
-                    &ctx.accounts.magic_context,
-                    &ctx.accounts.magic_program,
-                )?;
-                Ok(())
-            }
-        },
-        quote! {
-            #[automatically_derived]
-            #[derive(Accounts)]
-            pub struct Undelegate<'info> {
-                #[account(mut)]
-                pub payer: Signer<'info>,
-                #[account(mut)]
-                /// CHECK: The delegated component
-                pub delegated_account: AccountInfo<'info>,
-                #[account(mut, address = ::bolt_lang::MAGIC_CONTEXT_ID)]
-                /// CHECK:`
-                pub magic_context: AccountInfo<'info>,
-                #[account()]
-                /// CHECK:`
-                pub magic_program: Program<'info, MagicProgram>
-            }
-        },
-    )
-}
-
-/// Generates the undelegate function and struct.
-fn generate_process_undelegation() -> (TokenStream2, TokenStream2) {
-    (
-        quote! {
-            #[automatically_derived]
-            pub fn process_undelegation(ctx: Context<ProcessUndelegation>, account_seeds: Vec<Vec<u8>>) -> Result<()> {
-                let [delegated_account, buffer, payer, system_program] = [
-                    &ctx.accounts.delegated_account,
-                    &ctx.accounts.buffer,
-                    &ctx.accounts.payer,
-                    &ctx.accounts.system_program,
-                ];
-                ::bolt_lang::undelegate_account(
-                    delegated_account,
-                    &id(),
-                    buffer,
-                    payer,
-                    system_program,
-                    account_seeds,
-                )?;
-                Ok(())
-            }
-        },
-        quote! {
-            #[automatically_derived]
-            #[derive(Accounts)]
-            pub struct ProcessUndelegation<'info> {
-                /// CHECK:`
-                #[account(mut)]
-                pub delegated_account: AccountInfo<'info>,
-                /// CHECK:`
-                #[account()]
-                pub buffer: AccountInfo<'info>,
-                /// CHECK:
-                #[account(mut)]
-                pub payer: AccountInfo<'info>,
-                /// CHECK:
-                pub system_program: AccountInfo<'info>,
-            }
-        },
-    )
-}
-
-/// Generates the delegate instruction and related structs to inject in the component.
-fn generate_delegate(component_type: &Type) -> (TokenStream2, TokenStream2) {
-    (
-        quote! {
-            #[automatically_derived]
-            pub fn delegate(ctx: Context<DelegateInput>, commit_frequency_ms: u32, validator: Option<Pubkey>) -> Result<()> {
-                let pda_seeds: &[&[u8]] = &[<#component_type>::seed(), &ctx.accounts.entity.key().to_bytes()];
-
-                let del_accounts = ::bolt_lang::DelegateAccounts {
-                    payer: &ctx.accounts.payer,
-                    pda: &ctx.accounts.account,
-                    owner_program: &ctx.accounts.owner_program,
-                    buffer: &ctx.accounts.buffer,
-                    delegation_record: &ctx.accounts.delegation_record,
-                    delegation_metadata: &ctx.accounts.delegation_metadata,
-                    delegation_program: &ctx.accounts.delegation_program,
-                    system_program: &ctx.accounts.system_program,
-                };
-
-                let config = ::bolt_lang::DelegateConfig {
-                    commit_frequency_ms,
-                    validator,
-                };
-
-                ::bolt_lang::delegate_account(
-                    del_accounts,
-                    pda_seeds,
-                    config,
-                )?;
-
-                Ok(())
-            }
-        },
-        quote! {
-            #[automatically_derived]
-            #[derive(Accounts)]
-            pub struct DelegateInput<'info> {
-                pub payer: Signer<'info>,
-                #[account()]
-                pub entity: Account<'info, Entity>,
-                /// CHECK:
-                #[account(mut)]
-                pub account: AccountInfo<'info>,
-                /// CHECK:`
-                pub owner_program: AccountInfo<'info>,
-                /// CHECK:
-                #[account(mut)]
-                pub buffer: AccountInfo<'info>,
-                /// CHECK:`
-                #[account(mut)]
-                pub delegation_record: AccountInfo<'info>,
-                /// CHECK:`
-                #[account(mut)]
-                pub delegation_metadata: AccountInfo<'info>,
-                /// CHECK:`
-                pub delegation_program: AccountInfo<'info>,
-                /// CHECK:`
-                pub system_program: AccountInfo<'info>,
-            }
-        },
-    )
-}
-
-/// Extracts the type name from attribute arguments.
-fn extract_type_name(args: &AttributeArgs) -> Option<Type> {
-    args.iter().find_map(|arg| {
-        if let NestedMeta::Meta(syn::Meta::Path(path)) = arg {
-            Some(Type::Path(syn::TypePath {
-                qself: None,
-                path: path.clone(),
-            }))
-        } else {
-            None
-        }
-    })
+    bolt_attribute::delegate::process(args, input)
 }

crates/bolt-lang/attribute/delegate/src/lib.rs

@@ -7,9 +7,12 @@ use crate::component;
 use crate::system;
 use crate::common::generate_program;
 
-pub fn process(_attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn process(attr: TokenStream, item: TokenStream) -> TokenStream {
     let bundle_mod = parse_macro_input!(item as ItemMod);
     let mut program = generate_program(&bundle_mod.ident.to_string());
+    if attr.to_string().contains("delegate") {
+        program.attrs.insert(0, syn::parse_quote! { #[bolt_lang::delegate] });
+    }
     if let Some((_, items)) = bundle_mod.content {
         for item in items {
             match item {

crates/bolt-lang/attribute/src/bundle/mod.rs

@@ -6,16 +6,24 @@ pub fn generate_program(identifier: &str) -> syn::ItemMod {
     let snake_case_name = identifier.to_snake_case();
     let snake_case_name = syn::Ident::new(&snake_case_name, Span::call_site());
 
-    parse_quote! {
-        #[program]
-        pub mod #snake_case_name {
-            use super::*;
+    let mut module: syn::ItemMod = parse_quote! {
+        pub mod #snake_case_name {}
+    };
+    inject_program(&mut module);
+    module
+}
 
+pub fn inject_program(module: &mut syn::ItemMod) {
+    module.attrs.push(syn::parse_quote! { #[program] });
+    module.content.as_mut().map(|(brace, items)| {
+        items.insert(0, syn::Item::Use(syn::parse_quote! { use super::*; }));
+        items.insert(1, syn::Item::Struct(syn::parse_quote! {
             #[derive(Accounts)]
             pub struct VariadicBoltComponents<'info> {
                 #[account()]
                 pub authority: AccountInfo<'info>,
             }
-        }
-    }
+        }));
+        (brace, items.clone())
+    });
 }