light-sdk-macros 0.23.0

Macros for Programs using the Light SDK for ZK Compression
Documentation
//! Seed provider generation for PDA and Light Token accounts.

use std::collections::HashSet;

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

use super::{
    instructions::{InstructionDataSpec, TokenSeedSpec},
    seed_utils::generate_seed_derivation_body,
    visitors::{classify_seed, generate_client_seed_code},
};

/// Generate seed-related helper functions for token variants.
///
/// Currently generates client seed functions only. The legacy TokenSeedProvider
/// trait impls have been removed; seed derivation is now handled directly
/// via `LightAccountVariantTrait` impls generated in `variant_enum.rs`.
pub fn generate_ctoken_seed_provider_implementation(
    _token_seeds: &[TokenSeedSpec],
) -> Result<TokenStream> {
    // TokenSeedProvider is legacy - seed derivation is now handled by
    // LightAccountVariantTrait impls generated in variant_enum.rs
    Ok(quote! {})
}

#[inline(never)]
pub fn generate_client_seed_functions(
    pda_seeds: &Option<Vec<TokenSeedSpec>>,
    token_seeds: &Option<Vec<TokenSeedSpec>>,
    instruction_data: &[InstructionDataSpec],
    is_pinocchio: bool,
) -> Result<TokenStream> {
    let mut functions = Vec::new();

    // Choose the correct pubkey path based on framework
    let pubkey_path = if is_pinocchio {
        quote! { light_account_pinocchio::solana_pubkey::Pubkey }
    } else {
        quote! { solana_pubkey::Pubkey }
    };

    if let Some(pda_seed_specs) = pda_seeds {
        for spec in pda_seed_specs {
            let variant_name = &spec.variant;
            let snake_case = camel_to_snake_case(&variant_name.to_string());
            let function_name = format_ident!("get_{}_seeds", snake_case);

            let (parameters, seed_expressions) =
                analyze_seed_spec_for_client(spec, instruction_data, is_pinocchio)?;

            let fn_body = generate_seed_derivation_body(
                &seed_expressions,
                quote! { &#pubkey_path::from(crate::LIGHT_CPI_SIGNER.program_id) },
                is_pinocchio,
            );
            let function = quote! {
                pub fn #function_name(#(#parameters),*) -> (Vec<Vec<u8>>, #pubkey_path) {
                    #fn_body
                }
            };
            functions.push(function);
        }
    }

    if let Some(token_seed_specs) = token_seeds {
        for spec in token_seed_specs {
            let variant_name = &spec.variant;

            let function_name =
                format_ident!("get_{}_seeds", variant_name.to_string().to_lowercase());

            let (parameters, seed_expressions) =
                analyze_seed_spec_for_client(spec, instruction_data, is_pinocchio)?;

            let fn_body = generate_seed_derivation_body(
                &seed_expressions,
                quote! { &#pubkey_path::from(crate::LIGHT_CPI_SIGNER.program_id) },
                is_pinocchio,
            );
            let function = quote! {
                pub fn #function_name(#(#parameters),*) -> (Vec<Vec<u8>>, #pubkey_path) {
                    #fn_body
                }
            };
            functions.push(function);

            if let Some(owner_seeds_list) = &spec.owner_seeds {
                let owner_seeds_function_name = format_ident!(
                    "get_{}_owner_seeds",
                    variant_name.to_string().to_lowercase()
                );

                let mut owner_seeds_spec = TokenSeedSpec {
                    variant: spec.variant.clone(),
                    _eq: spec._eq,
                    is_token: spec.is_token,
                    seeds: syn::punctuated::Punctuated::new(),
                    owner_seeds: None,
                    inner_type: spec.inner_type.clone(),
                    is_zero_copy: spec.is_zero_copy,
                };

                for owner_seed in owner_seeds_list {
                    owner_seeds_spec.seeds.push(owner_seed.clone());
                }

                let (owner_parameters, owner_seed_expressions) = analyze_seed_spec_for_client(
                    &owner_seeds_spec,
                    instruction_data,
                    is_pinocchio,
                )?;

                let (fn_params, fn_body) = if owner_parameters.is_empty() {
                    (
                        quote! { _program_id: &#pubkey_path },
                        generate_seed_derivation_body(
                            &owner_seed_expressions,
                            quote! { _program_id },
                            is_pinocchio,
                        ),
                    )
                } else {
                    (
                        quote! { #(#owner_parameters),* },
                        generate_seed_derivation_body(
                            &owner_seed_expressions,
                            quote! { &#pubkey_path::from(crate::LIGHT_CPI_SIGNER.program_id) },
                            is_pinocchio,
                        ),
                    )
                };
                let owner_seeds_function = quote! {
                    pub fn #owner_seeds_function_name(#fn_params) -> (Vec<Vec<u8>>, #pubkey_path) {
                        #fn_body
                    }
                };
                functions.push(owner_seeds_function);
            }
        }
    }

    Ok(quote! {
        #[cfg(not(target_os = "solana"))]
        mod __client_seed_functions {
            use super::*;
            #(#functions)*
        }

        #[cfg(not(target_os = "solana"))]
        pub use __client_seed_functions::*;
    })
}

/// Analyze a seed spec and generate client function parameters and seed expressions.
///
/// Uses the classification-based approach: first classify each seed, then generate code.
/// This separates "what kind of seed is this?" from "what code to generate?".
#[inline(never)]
fn analyze_seed_spec_for_client(
    spec: &TokenSeedSpec,
    instruction_data: &[InstructionDataSpec],
    is_pinocchio: bool,
) -> Result<(Vec<TokenStream>, Vec<TokenStream>)> {
    let mut parameters = Vec::new();
    let mut expressions = Vec::new();
    let mut seen_params = HashSet::new();

    for seed in &spec.seeds {
        // Phase 1: Classification
        let info = classify_seed(seed)?;

        // Phase 2: Code generation (modifies parameters and expressions in place)
        generate_client_seed_code(
            &info,
            instruction_data,
            &mut seen_params,
            &mut parameters,
            &mut expressions,
            is_pinocchio,
        )?;
    }

    Ok((parameters, expressions))
}

fn camel_to_snake_case(s: &str) -> String {
    let mut result = String::new();
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() && i > 0 {
            result.push('_');
        }
        result.push(c.to_lowercase().next().unwrap());
    }
    result
}