feat: macro token account support

sdk-libs/macros/src/light_pdas/accounts/builder.rs

@@ -11,6 +11,7 @@ use super::{
     mint::{InfraRefs, LightMintsBuilder},
     parse::{InfraFieldType, ParsedLightAccountsStruct},
     pda::generate_pda_compress_blocks,
+    token::TokenAccountsBuilder,
 };
 
 /// Builder for RentFree derive macro code generation.
@@ -46,15 +47,20 @@ impl LightAccountsBuilder {
 
     /// Validate constraints (e.g., account count < 255).
     pub fn validate(&self) -> Result<(), syn::Error> {
-        let total = self.parsed.rentfree_fields.len() + self.parsed.light_mint_fields.len();
+        let total = self.parsed.rentfree_fields.len()
+            + self.parsed.light_mint_fields.len()
+            + self.parsed.token_account_fields.len()
+            + self.parsed.ata_fields.len();
         if total > 255 {
             return Err(syn::Error::new_spanned(
                 &self.parsed.struct_name,
                 format!(
-                    "Too many compression fields ({} PDAs + {} mints = {} total, maximum 255). \
+                    "Too many compression fields ({} PDAs + {} mints + {} tokens + {} ATAs = {} total, maximum 255). \
                      Light Protocol uses u8 for account indices.",
                     self.parsed.rentfree_fields.len(),
                     self.parsed.light_mint_fields.len(),
+                    self.parsed.token_account_fields.len(),
+                    self.parsed.ata_fields.len(),
                     total
                 ),
             ));
@@ -70,9 +76,11 @@ impl LightAccountsBuilder {
     fn validate_infra_fields(&self) -> Result<(), syn::Error> {
         let has_pdas = self.has_pdas();
         let has_mints = self.has_mints();
+        let has_token_accounts = self.has_token_accounts();
+        let has_atas = self.has_atas();
 
         // Skip validation if no light_account fields
-        if !has_pdas && !has_mints {
+        if !has_pdas && !has_mints && !has_token_accounts && !has_atas {
             return Ok(());
         }
 
@@ -88,27 +96,38 @@ impl LightAccountsBuilder {
             missing.push(InfraFieldType::CompressionConfig);
         }
 
-        // Mints require light_token_config, light_token_rent_sponsor, light_token_cpi_authority
-        if has_mints {
+        // Mints, token accounts, and ATAs require light_token infrastructure
+        let needs_token_infra = has_mints || has_token_accounts || has_atas;
+        if needs_token_infra {
             if self.parsed.infra_fields.light_token_config.is_none() {
                 missing.push(InfraFieldType::LightTokenConfig);
             }
             if self.parsed.infra_fields.light_token_rent_sponsor.is_none() {
                 missing.push(InfraFieldType::LightTokenRentSponsor);
             }
-            if self.parsed.infra_fields.light_token_cpi_authority.is_none() {
+            // CPI authority is required for mints and token accounts (PDA-based signing)
+            if (has_mints || has_token_accounts)
+                && self.parsed.infra_fields.light_token_cpi_authority.is_none()
+            {
                 missing.push(InfraFieldType::LightTokenCpiAuthority);
             }
         }
 
         if !missing.is_empty() {
-            let context = if has_pdas && has_mints {
-                "PDA and mint"
-            } else if has_mints {
-                "mint"
-            } else {
-                "PDA"
-            };
+            let mut types = Vec::new();
+            if has_pdas {
+                types.push("PDA");
+            }
+            if has_mints {
+                types.push("mint");
+            }
+            if has_token_accounts {
+                types.push("token account");
+            }
+            if has_atas {
+                types.push("ATA");
+            }
+            let context = types.join(", ");
 
             let mut msg = format!(
                 "#[derive(LightAccounts)] with {} fields requires the following infrastructure fields:\n",
@@ -129,16 +148,26 @@ impl LightAccountsBuilder {
         Ok(())
     }
 
-    /// Query: any #[light_account(init)] fields?
+    /// Query: any #[light_account(init)] PDA fields?
     pub fn has_pdas(&self) -> bool {
         !self.parsed.rentfree_fields.is_empty()
     }
 
-    /// Query: any #[light_account(init)] fields?
+    /// Query: any #[light_account(init, mint, ...)] fields?
     pub fn has_mints(&self) -> bool {
         !self.parsed.light_mint_fields.is_empty()
     }
 
+    /// Query: any #[light_account(init, token, ...)] fields?
+    pub fn has_token_accounts(&self) -> bool {
+        !self.parsed.token_account_fields.is_empty()
+    }
+
+    /// Query: any #[light_account(init, associated_token, ...)] fields?
+    pub fn has_atas(&self) -> bool {
+        !self.parsed.ata_fields.is_empty()
+    }
+
     /// Query: #[instruction(...)] present?
     pub fn has_instruction_args(&self) -> bool {
         self.parsed.instruction_args.is_some()
@@ -374,4 +403,24 @@ impl LightAccountsBuilder {
             }
         })
     }
+
+    /// Check if token accounts or ATAs need finalize code generation.
+    pub fn needs_token_finalize(&self) -> bool {
+        TokenAccountsBuilder::new(
+            &self.parsed.token_account_fields,
+            &self.parsed.ata_fields,
+            &self.infra,
+        )
+        .needs_finalize()
+    }
+
+    /// Generate finalize body for token accounts and ATAs.
+    pub fn generate_token_finalize_body(&self) -> TokenStream {
+        TokenAccountsBuilder::new(
+            &self.parsed.token_account_fields,
+            &self.parsed.ata_fields,
+            &self.infra,
+        )
+        .generate_finalize_body()
+    }
 }

sdk-libs/macros/src/light_pdas/accounts/derive.rs

@@ -46,10 +46,205 @@ pub(super) fn derive_light_accounts(input: &DeriveInput) -> Result<TokenStream,
 
     // Generate trait implementations
     let pre_init_impl = builder.generate_pre_init_impl(pre_init)?;
-    let finalize_impl = builder.generate_finalize_impl(quote! { Ok(()) })?;
+
+    // Generate finalize body - token accounts and ATAs are created here
+    let finalize_body = if builder.needs_token_finalize() {
+        builder.generate_token_finalize_body()
+    } else {
+        quote! { Ok(()) }
+    };
+    let finalize_impl = builder.generate_finalize_impl(finalize_body)?;
 
     Ok(quote! {
         #pre_init_impl
         #finalize_impl
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use syn::parse_quote;
+
+    use super::*;
+
+    #[test]
+    fn test_token_account_with_init_generates_create_cpi() {
+        // Token account with init should generate CreateTokenAccountCpi in finalize
+        let input: DeriveInput = parse_quote! {
+            #[instruction(params: CreateVaultParams)]
+            pub struct CreateVault<'info> {
+                #[account(mut)]
+                pub fee_payer: Signer<'info>,
+
+                #[light_account(init, token, authority = [b"authority"], mint = my_mint, owner = fee_payer)]
+                pub vault: Account<'info, CToken>,
+
+                pub light_token_compressible_config: Account<'info, CompressibleConfig>,
+                pub light_token_rent_sponsor: Account<'info, RentSponsor>,
+                pub light_token_cpi_authority: AccountInfo<'info>,
+            }
+        };
+
+        let result = derive_light_accounts(&input);
+        assert!(result.is_ok(), "Token account derive should succeed");
+
+        let output = result.unwrap().to_string();
+
+        // Verify finalize generates token account creation
+        assert!(
+            output.contains("LightFinalize"),
+            "Should generate LightFinalize impl"
+        );
+        assert!(
+            output.contains("CreateTokenAccountCpi"),
+            "Should generate CreateTokenAccountCpi call"
+        );
+        assert!(
+            output.contains("rent_free"),
+            "Should call rent_free on CreateTokenAccountCpi"
+        );
+        assert!(
+            output.contains("invoke_signed"),
+            "Should call invoke_signed with seeds"
+        );
+    }
+
+    #[test]
+    fn test_ata_with_init_generates_create_cpi() {
+        // ATA with init should generate create_associated_token_account_idempotent in finalize
+        let input: DeriveInput = parse_quote! {
+            #[instruction(params: CreateAtaParams)]
+            pub struct CreateAta<'info> {
+                #[account(mut)]
+                pub fee_payer: Signer<'info>,
+
+                #[light_account(init, associated_token, owner = wallet, mint = my_mint)]
+                pub user_ata: Account<'info, CToken>,
+
+                pub wallet: AccountInfo<'info>,
+                pub my_mint: AccountInfo<'info>,
+                pub light_token_compressible_config: Account<'info, CompressibleConfig>,
+                pub light_token_rent_sponsor: Account<'info, RentSponsor>,
+            }
+        };
+
+        let result = derive_light_accounts(&input);
+        assert!(result.is_ok(), "ATA derive should succeed");
+
+        let output = result.unwrap().to_string();
+
+        // Verify finalize generates ATA creation
+        assert!(
+            output.contains("LightFinalize"),
+            "Should generate LightFinalize impl"
+        );
+        assert!(
+            output.contains("CreateTokenAtaCpi"),
+            "Should generate CreateTokenAtaCpi call"
+        );
+    }
+
+    #[test]
+    fn test_token_mark_only_generates_noop_finalize() {
+        // Token without init should NOT generate creation code (mark-only mode)
+        // Mark-only returns None from parsing, so token_account_fields is empty
+        let input: DeriveInput = parse_quote! {
+            #[instruction(params: UseVaultParams)]
+            pub struct UseVault<'info> {
+                #[account(mut)]
+                pub fee_payer: Signer<'info>,
+
+                // Mark-only: no init keyword
+                #[light_account(token, authority = [b"authority"])]
+                pub vault: Account<'info, CToken>,
+            }
+        };
+
+        let result = derive_light_accounts(&input);
+        assert!(result.is_ok(), "Mark-only token derive should succeed");
+
+        let output = result.unwrap().to_string();
+
+        // Mark-only should NOT have token account creation
+        assert!(
+            !output.contains("CreateTokenAccountCpi"),
+            "Mark-only should NOT generate CreateTokenAccountCpi"
+        );
+
+        // Should still have LightFinalize but with Ok(())
+        assert!(
+            output.contains("LightFinalize"),
+            "Should still generate LightFinalize impl"
+        );
+    }
+
+    #[test]
+    fn test_mixed_token_and_ata_generates_both() {
+        // Mixed token account + ATA should generate both creation codes
+        let input: DeriveInput = parse_quote! {
+            #[instruction(params: CreateBothParams)]
+            pub struct CreateBoth<'info> {
+                #[account(mut)]
+                pub fee_payer: Signer<'info>,
+
+                #[light_account(init, token, authority = [b"authority"], mint = my_mint, owner = fee_payer)]
+                pub vault: Account<'info, CToken>,
+
+                #[light_account(init, associated_token, owner = wallet, mint = my_mint)]
+                pub user_ata: Account<'info, CToken>,
+
+                pub wallet: AccountInfo<'info>,
+                pub my_mint: AccountInfo<'info>,
+                pub light_token_compressible_config: Account<'info, CompressibleConfig>,
+                pub light_token_rent_sponsor: Account<'info, RentSponsor>,
+                pub light_token_cpi_authority: AccountInfo<'info>,
+            }
+        };
+
+        let result = derive_light_accounts(&input);
+        assert!(result.is_ok(), "Mixed token+ATA derive should succeed");
+
+        let output = result.unwrap().to_string();
+
+        // Should have both creation types
+        assert!(
+            output.contains("CreateTokenAccountCpi"),
+            "Should generate CreateTokenAccountCpi for vault"
+        );
+        assert!(
+            output.contains("CreateTokenAtaCpi"),
+            "Should generate CreateTokenAtaCpi for ATA"
+        );
+    }
+
+    #[test]
+    fn test_no_instruction_args_generates_noop() {
+        // No #[instruction] attribute should generate no-op impls
+        let input: DeriveInput = parse_quote! {
+            pub struct NoInstruction<'info> {
+                #[account(mut)]
+                pub fee_payer: Signer<'info>,
+            }
+        };
+
+        let result = derive_light_accounts(&input);
+        assert!(result.is_ok(), "No instruction args should succeed");
+
+        let output = result.unwrap().to_string();
+
+        // Should generate no-op impls with () param type
+        assert!(
+            output.contains("LightPreInit"),
+            "Should generate LightPreInit impl"
+        );
+        assert!(
+            output.contains("LightFinalize"),
+            "Should generate LightFinalize impl"
+        );
+        // No-op returns Ok(false) in pre_init and Ok(()) in finalize
+        assert!(
+            output.contains("Ok (false)") || output.contains("Ok(false)"),
+            "Should return Ok(false) in pre_init"
+        );
+    }
+}