light-sdk-macros 0.23.0

//! PDA block code generation for rent-free accounts.
//!
//! This module handles the generation of compression blocks for PDA fields
//! marked with `#[light_account(init)]`. Each PDA field generates code for:
//! - Account extraction (get account info and key)
//! - Address derivation and registration via prepare_compressed_account_on_init
//! - Rent reimbursement from rent sponsor PDA to fee payer

use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::Ident;

use super::{mint::InfraRefs, parse::ParsedPdaField};

/// Generated identifier names for a PDA field.
pub(super) struct PdaIdents {
    pub idx: u8,
    pub account_info: Ident,
    pub account_key: Ident,
    pub address_tree_pubkey: Ident,
    pub account_data: Ident,
}

impl PdaIdents {
    pub fn new(idx: usize) -> Self {
        Self {
            idx: idx as u8,
            account_info: format_ident!("__account_info_{}", idx),
            account_key: format_ident!("__account_key_{}", idx),
            address_tree_pubkey: format_ident!("__address_tree_pubkey_{}", idx),
            account_data: format_ident!("__account_data_{}", idx),
        }
    }
}

/// Builder for PDA compression block code generation.
pub(super) struct PdaBlockBuilder<'a> {
    field: &'a ParsedPdaField,
    idents: PdaIdents,
}

impl<'a> PdaBlockBuilder<'a> {
    pub fn new(field: &'a ParsedPdaField, idx: usize) -> Self {
        Self {
            field,
            idents: PdaIdents::new(idx),
        }
    }

    /// Generate account extraction (get account info and key).
    fn account_extraction(&self) -> TokenStream {
        let field_name = &self.field.field_name;
        let account_info = &self.idents.account_info;
        let account_key = &self.idents.account_key;

        quote! {
            let #account_info = self.#field_name.to_account_info();
            let #account_key = #account_info.key.to_bytes();
        }
    }

    /// Generate address tree pubkey extraction.
    fn address_tree_extraction(&self) -> TokenStream {
        let addr_tree_info = self
            .field
            .address_tree_info
            .as_ref()
            .expect("address_tree_info required for derive macro");
        let address_tree_pubkey = &self.idents.address_tree_pubkey;

        quote! {
            let #address_tree_pubkey: [u8; 32] = {
                // Explicit type annotation ensures clear error if wrong type is provided.
                let tree_info: &light_account::PackedAddressTreeInfo = &#addr_tree_info;
                let __tree_account = cpi_accounts
                    .get_tree_account_info(tree_info.address_merkle_tree_pubkey_index as usize)?;
                light_account::AccountInfoTrait::key(__tree_account)
            };
        }
    }

    /// Generate account data initialization (set CompressionInfo).
    fn account_data_init(&self) -> TokenStream {
        let ident = &self.field.field_name;
        let account_data = &self.idents.account_data;

        if self.field.is_zero_copy {
            // AccountLoader uses load_init() for newly initialized accounts
            let account_guard = format_ident!("{}_guard", ident);
            quote! {
                {
                    let current_slot = anchor_lang::solana_program::sysvar::clock::Clock::get()
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?.slot;
                    let mut #account_guard = self.#ident.load_init()
                        .map_err(|_| light_account::LightSdkTypesError::InvalidInstructionData)?;
                    let #account_data = &mut *#account_guard;
                    // For zero-copy Pod accounts, set compression_info directly
                    #account_data.compression_info =
                        light_account::CompressionInfo::new_from_config(
                            &compression_config_data,
                            current_slot,
                        );
                }
            }
        } else if self.field.is_boxed {
            quote! {
                {
                    use light_account::LightAccount;
                    use anchor_lang::AnchorSerialize;
                    let current_slot = anchor_lang::solana_program::sysvar::clock::Clock::get()
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?.slot;
                    // Get account info BEFORE mutable borrow
                    let account_info = self.#ident.to_account_info();
                    // Scope the mutable borrow
                    {
                        let #account_data = &mut **self.#ident;
                        // Initialize CompressionInfo using v2 LightAccount trait
                        #account_data.set_decompressed(&compression_config_data, current_slot);
                    }
                    // Now serialize - the mutable borrow above is released
                    let mut data = account_info
                        .try_borrow_mut_data()
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?;
                    self.#ident.serialize(&mut &mut data[8..])
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?;
                }
            }
        } else {
            quote! {
                {
                    use light_account::LightAccount;
                    use anchor_lang::AnchorSerialize;
                    let current_slot = anchor_lang::solana_program::sysvar::clock::Clock::get()
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?.slot;
                    // Get account info BEFORE mutable borrow
                    let account_info = self.#ident.to_account_info();
                    // Scope the mutable borrow
                    {
                        let #account_data = &mut *self.#ident;
                        // Initialize CompressionInfo using v2 LightAccount trait
                        #account_data.set_decompressed(&compression_config_data, current_slot);
                    }
                    // Now serialize - the mutable borrow above is released
                    let mut data = account_info
                        .try_borrow_mut_data()
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?;
                    self.#ident.serialize(&mut &mut data[8..])
                        .map_err(|_| light_account::LightSdkTypesError::ConstraintViolation)?;
                }
            }
        }
    }

    /// Generate the call to prepare_compressed_account_on_init.
    fn prepare_call(&self) -> TokenStream {
        let addr_tree_info = self
            .field
            .address_tree_info
            .as_ref()
            .expect("address_tree_info required for derive macro");
        let output_tree = self
            .field
            .output_tree
            .as_ref()
            .expect("output_tree required for derive macro");
        let account_key = &self.idents.account_key;
        let address_tree_pubkey = &self.idents.address_tree_pubkey;
        let idx = self.idents.idx;

        quote! {
            {
                // Explicit type annotation for tree_info
                let tree_info: &light_account::PackedAddressTreeInfo = &#addr_tree_info;

                ::light_account::prepare_compressed_account_on_init(
                    &#account_key,
                    &#address_tree_pubkey,
                    tree_info,
                    #output_tree,
                    #idx,
                    &crate::LIGHT_CPI_SIGNER.program_id,
                    &mut all_new_address_params,
                    &mut all_compressed_infos,
                )?;
            }
        }
    }

    /// Build the complete compression block for this PDA field.
    pub fn build(&self) -> TokenStream {
        let account_extraction = self.account_extraction();
        let address_tree_extraction = self.address_tree_extraction();
        let account_data_init = self.account_data_init();
        let prepare_call = self.prepare_call();

        quote! {
            // Get account info early before any mutable borrows
            #account_extraction
            // Extract address tree pubkey
            #address_tree_extraction
            // Initialize CompressionInfo in account data
            #account_data_init
            // Register compressed address
            #prepare_call
        }
    }
}

/// Generate compression blocks for PDA fields using PdaBlockBuilder.
///
/// Returns a vector of TokenStreams for compression blocks.
/// The blocks push into `all_new_address_params` and `all_compressed_infos` vectors.
pub(super) fn generate_pda_compress_blocks(fields: &[ParsedPdaField]) -> Vec<TokenStream> {
    fields
        .iter()
        .enumerate()
        .map(|(idx, field)| PdaBlockBuilder::new(field, idx).build())
        .collect()
}

/// Generate rent reimbursement code that calls the SDK reimburse_rent function.
///
/// This generates a single call to `reimburse_rent` with all created PDA account infos,
/// which transfers the total rent-exemption amount from the rent sponsor PDA to the fee payer.
pub(super) fn generate_rent_reimbursement_block(
    fields: &[ParsedPdaField],
    infra: &InfraRefs,
) -> TokenStream {
    if fields.is_empty() {
        return quote! {};
    }

    let fee_payer = &infra.fee_payer;
    let rent_sponsor = &infra.pda_rent_sponsor;

    // Collect account info expressions for all PDA fields
    let account_info_exprs: Vec<TokenStream> = fields
        .iter()
        .map(|field| {
            let field_name = &field.field_name;
            quote! { self.#field_name.to_account_info() }
        })
        .collect();

    let count = fields.len();

    quote! {
        // Reimburse fee_payer for rent paid to Anchor for all PDAs
        {
            let __created_accounts: [solana_account_info::AccountInfo<'info>; #count] = [
                #(#account_info_exprs),*
            ];
            let __rent_sponsor_bump_byte = [compression_config_data.rent_sponsor_bump];
            let __rent_sponsor_seeds: &[&[u8]] = &[
                light_account::RENT_SPONSOR_SEED,
                &__rent_sponsor_bump_byte,
            ];
            ::light_account::reimburse_rent(
                &__created_accounts,
                &self.#fee_payer.to_account_info(),
                &self.#rent_sponsor.to_account_info(),
                __rent_sponsor_seeds,
            )?;
        }
    }
}