diff --git a/program-libs/compressed-account/src/instruction_data/with_account_info.rs b/program-libs/compressed-account/src/instruction_data/with_account_info.rs index 01ae244321..94157a100b 100644 --- a/program-libs/compressed-account/src/instruction_data/with_account_info.rs +++ b/program-libs/compressed-account/src/instruction_data/with_account_info.rs @@ -1,6 +1,8 @@ -use core::ops::{Deref, DerefMut}; +use core::ops::Deref; -use light_zero_copy::{errors::ZeroCopyError, slice::ZeroCopySliceBorsh, traits::ZeroCopyAt}; +use light_zero_copy::{ + errors::ZeroCopyError, slice::ZeroCopySliceBorsh, traits::ZeroCopyAt, ZeroCopyMut, +}; use zerocopy::{ little_endian::{U16, U32, U64}, FromBytes, Immutable, IntoBytes, KnownLayout, Ref, Unaligned, @@ -39,7 +41,8 @@ use crate::{ not(feature = "anchor"), derive(borsh::BorshDeserialize, borsh::BorshSerialize) )] -#[derive(Debug, Default, PartialEq, Clone)] +#[repr(C)] +#[derive(Debug, Default, PartialEq, Clone, ZeroCopyMut)] pub struct InAccountInfo { pub discriminator: [u8; 8], /// Data hash @@ -101,7 +104,8 @@ pub struct ZInAccountInfo { not(feature = "anchor"), derive(borsh::BorshDeserialize, borsh::BorshSerialize) )] -#[derive(Debug, Default, PartialEq, Clone)] +#[repr(C)] +#[derive(Debug, Default, PartialEq, Clone, ZeroCopyMut)] pub struct OutAccountInfo { pub discriminator: [u8; 8], /// Data hash @@ -305,27 +309,6 @@ impl Deref for ZOutAccountInfo<'_> { } } -#[derive(Debug, PartialEq)] -pub struct ZOutAccountInfoMut<'a> { - meta: Ref<&'a mut [u8], ZOutAccountInfoMeta>, - /// Account data. - pub data: &'a mut [u8], -} - -impl Deref for ZOutAccountInfoMut<'_> { - type Target = ZOutAccountInfoMeta; - - fn deref(&self) -> &Self::Target { - &self.meta - } -} - -impl DerefMut for ZOutAccountInfoMut<'_> { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.meta - } -} - #[cfg_attr( all(feature = "std", feature = "anchor"), derive(anchor_lang::AnchorDeserialize, anchor_lang::AnchorSerialize) @@ -334,7 +317,8 @@ impl DerefMut for ZOutAccountInfoMut<'_> { not(feature = "anchor"), derive(borsh::BorshDeserialize, borsh::BorshSerialize) )] -#[derive(Debug, PartialEq, Clone, Default)] +#[repr(C)] +#[derive(Debug, PartialEq, Clone, Default, ZeroCopyMut)] pub struct CompressedAccountInfo { /// Address. pub address: Option<[u8; 32]>, @@ -383,7 +367,8 @@ impl<'a> CompressedAccountInfo { not(feature = "anchor"), derive(borsh::BorshDeserialize, borsh::BorshSerialize) )] -#[derive(Debug, PartialEq, Default, Clone)] +#[repr(C)] +#[derive(Debug, PartialEq, Default, Clone, ZeroCopyMut)] pub struct InstructionDataInvokeCpiWithAccountInfo { /// 0 V1 instruction accounts. /// 1 Optimized V2 instruction accounts. diff --git a/sdk-libs/macros/src/light_pdas/account/decompress_context.rs b/sdk-libs/macros/src/light_pdas/account/decompress_context.rs index 269221f96f..e241476121 100644 --- a/sdk-libs/macros/src/light_pdas/account/decompress_context.rs +++ b/sdk-libs/macros/src/light_pdas/account/decompress_context.rs @@ -1,4 +1,9 @@ //! DecompressContext trait generation. +//! +//! Generates the implementation of the DecompressContext trait for the +//! DecompressAccountsIdempotent struct. This uses a zero-allocation two-pass approach: +//! - Pass 1 (collect_layout_and_tokens): Count PDAs, collect output_data_lens, collect tokens +//! - Pass 2 (create_and_write_pda): Create PDA on Solana, return data for zero-copy buffer use proc_macro2::TokenStream; use quote::{format_ident, quote}; @@ -15,23 +20,40 @@ pub fn generate_decompress_context_trait_impl( token_variant_ident: Ident, lifetime: syn::Lifetime, ) -> Result { - // Generate match arms that extract idx fields, resolve Pubkeys, construct CtxSeeds - let pda_match_arms: Vec<_> = pda_ctx_seeds + // Generate match arms for collect_layout_and_tokens - count PDAs that need decompression + let collect_layout_pda_arms: Vec<_> = pda_ctx_seeds + .iter() + .map(|info| { + let variant_name = &info.variant_name; + let packed_variant_name = make_packed_variant_name(variant_name); + quote! { + LightAccountVariant::#packed_variant_name { .. } => { + // PDA variant: only count if not already initialized (idempotent check) + if solana_accounts[i].data_is_empty() { + pda_indices[pda_count] = i; + pda_count += 1; + } + } + LightAccountVariant::#variant_name { .. } => { + return std::result::Result::Err(light_sdk::error::LightSdkError::UnexpectedUnpackedVariant.into()); + } + } + }) + .collect(); + + // Generate match arms for create_and_write_pda - unpack, derive seeds, create PDA, return data + let create_pda_match_arms: Vec<_> = pda_ctx_seeds .iter() .map(|info| { - // Use variant_name for enum variant matching let variant_name = &info.variant_name; - // Use inner_type for type references (generics, trait bounds) - // Qualify with crate:: to ensure it's accessible from generated code let inner_type = qualify_type_with_crate(&info.inner_type); let packed_variant_name = make_packed_variant_name(variant_name); - // Create packed type (also qualified with crate::) let packed_inner_type = make_packed_type(&info.inner_type) .expect("inner_type should be a valid type path"); - // Use variant_name for CtxSeeds struct (matches what decompress.rs generates) let ctx_seeds_struct_name = format_ident!("{}CtxSeeds", variant_name); let ctx_fields = &info.ctx_seed_fields; let params_only_fields = &info.params_only_seed_fields; + // Generate pattern to extract idx fields from packed variant let idx_field_patterns: Vec<_> = ctx_fields.iter().map(|field| { let idx_field = format_ident!("{}_idx", field); @@ -42,11 +64,12 @@ pub fn generate_decompress_context_trait_impl( quote! { #field } }).collect(); // Generate code to resolve idx fields to Pubkeys + // Note: when matching on &compressed_data.data, idx fields are references, so we dereference let resolve_ctx_seeds: Vec<_> = ctx_fields.iter().map(|field| { let idx_field = format_ident!("{}_idx", field); quote! { let #field = *post_system_accounts - .get(#idx_field as usize) + .get(*#idx_field as usize) .ok_or(solana_program_error::ProgramError::InvalidAccountData)? .key; } @@ -61,38 +84,40 @@ pub fn generate_decompress_context_trait_impl( quote! { let ctx_seeds = #ctx_seeds_struct_name { #(#field_inits),* }; } }; // Generate SeedParams update with params-only field values - // Note: variant_seed_params is declared OUTSIDE the match to avoid borrow checker issues - // (the reference passed to handle_packed_pda_variant would outlive the match arm scope) - // params-only fields are stored directly in packed variant (not by reference), - // so we use the value directly without dereferencing + // Note: when matching on &compressed_data.data, params fields are references, so we dereference let seed_params_update = if params_only_fields.is_empty() { - // No update needed - use the default value declared before match quote! {} } else { let field_inits: Vec<_> = params_only_fields.iter().map(|(field, _, _)| { - quote! { #field: std::option::Option::Some(#field) } + quote! { #field: std::option::Option::Some(*#field) } }).collect(); quote! { variant_seed_params = SeedParams { #(#field_inits,)* ..Default::default() }; } }; + quote! { LightAccountVariant::#packed_variant_name { data: packed, #(#idx_field_patterns,)* #(#params_field_patterns,)* .. } => { #(#resolve_ctx_seeds)* #ctx_seeds_construction #seed_params_update - light_sdk::interface::handle_packed_pda_variant::<#inner_type, #packed_inner_type, _, _>( - &*self.rent_sponsor, - cpi_accounts, - address_space, - &solana_accounts[i], - i, - &packed, - &meta, - post_system_accounts, - &mut compressed_pda_infos, + + // Unpack the data + let data: #inner_type = <#packed_inner_type as light_sdk::interface::Unpack>::unpack(&packed, post_system_accounts)?; + + // Use helper function to derive seeds, verify PDA, create account, and write to zero-copy buffer + // Pass data and compressed_meta by reference to reduce caller stack usage + light_sdk::interface::derive_verify_create_and_write_pda::<#inner_type, _, _>( &program_id, + &data, &ctx_seeds, - std::option::Option::Some(&variant_seed_params), - )?; + seed_params, + &variant_seed_params, + compressed_meta, + address_space, + solana_account, + &*self.rent_sponsor, + cpi_accounts, + zc_info, + ) } LightAccountVariant::#variant_name { .. } => { return std::result::Result::Err(light_sdk::error::LightSdkError::UnexpectedUnpackedVariant.into()); @@ -102,9 +127,18 @@ pub fn generate_decompress_context_trait_impl( .collect(); // For mint-only programs (no PDA variants), add an arm for the Empty variant - let empty_variant_arm = if pda_ctx_seeds.is_empty() { + let empty_variant_arm_collect = if pda_ctx_seeds.is_empty() { + quote! { + LightAccountVariant::Empty => { + return std::result::Result::Err(solana_program_error::ProgramError::InvalidAccountData); + } + } + } else { + quote! {} + }; + + let empty_variant_arm_create = if pda_ctx_seeds.is_empty() { quote! { - // Mint-only programs have an Empty variant that should never be decompressed LightAccountVariant::Empty => { return std::result::Result::Err(solana_program_error::ProgramError::InvalidAccountData); } @@ -150,51 +184,65 @@ pub fn generate_decompress_context_trait_impl( self.ctoken_config.as_ref().map(|a| &**a) } - fn collect_pda_and_token<'b>( + #[allow(clippy::type_complexity)] + fn collect_layout_and_tokens( &self, - cpi_accounts: &light_sdk::cpi::v2::CpiAccounts<'b, #lifetime>, - address_space: solana_pubkey::Pubkey, - compressed_accounts: Vec, + compressed_accounts: &[Self::CompressedData], solana_accounts: &[solana_account_info::AccountInfo<#lifetime>], - seed_params: std::option::Option<&Self::SeedParams>, - ) -> std::result::Result<( - Vec<::light_sdk::compressed_account::CompressedAccountInfo>, - Vec<(Self::PackedTokenData, Self::CompressedMeta)>, - ), solana_program_error::ProgramError> { - solana_msg::msg!("collect_pda_and_token: start, {} accounts", compressed_accounts.len()); - let post_system_offset = cpi_accounts.system_accounts_end_offset(); - let all_infos = cpi_accounts.account_infos(); - let post_system_accounts = &all_infos[post_system_offset..]; - let program_id = &crate::ID; - - solana_msg::msg!("collect_pda_and_token: allocating vecs"); - let mut compressed_pda_infos = Vec::with_capacity(compressed_accounts.len()); + pda_indices: &mut [usize; light_sdk::interface::MAX_DECOMPRESS_ACCOUNTS], + ) -> std::result::Result<(usize, Vec<(Self::PackedTokenData, Self::CompressedMeta)>), solana_program_error::ProgramError> { + let mut pda_count: usize = 0; let mut compressed_token_accounts = Vec::with_capacity(compressed_accounts.len()); - solana_msg::msg!("collect_pda_and_token: starting loop"); - for (i, compressed_data) in compressed_accounts.into_iter().enumerate() { - solana_msg::msg!("collect_pda_and_token: processing account {}", i); - let meta = compressed_data.meta; - // Declare variant_seed_params OUTSIDE the match to avoid borrow checker issues - // (reference passed to handle_packed_pda_variant with ? would outlive match arm scope) - let mut variant_seed_params = SeedParams::default(); - match compressed_data.data { - #(#pda_match_arms)* - LightAccountVariant::PackedCTokenData(mut data) => { - solana_msg::msg!("collect_pda_and_token: token variant {}", i); - data.token_data.version = 3; - compressed_token_accounts.push((data, meta)); - solana_msg::msg!("collect_pda_and_token: token {} done", i); + for (i, compressed_data) in compressed_accounts.iter().enumerate() { + let meta = compressed_data.meta.clone(); + match &compressed_data.data { + #(#collect_layout_pda_arms)* + LightAccountVariant::PackedCTokenData(data) => { + let mut token_data = data.clone(); + token_data.token_data.version = 3; + compressed_token_accounts.push((token_data, meta)); } LightAccountVariant::CTokenData(_) => { return std::result::Result::Err(light_sdk::error::LightSdkError::UnexpectedUnpackedVariant.into()); } - #empty_variant_arm + #empty_variant_arm_collect } } - solana_msg::msg!("collect_pda_and_token: loop done, pdas={} tokens={}", compressed_pda_infos.len(), compressed_token_accounts.len()); - std::result::Result::Ok((compressed_pda_infos, compressed_token_accounts)) + std::result::Result::Ok((pda_count, compressed_token_accounts)) + } + + #[inline(never)] + #[allow(clippy::too_many_arguments)] + fn create_and_write_pda<'b, 'c>( + &self, + cpi_accounts: &light_sdk::cpi::v2::CpiAccounts<'b, #lifetime>, + address_space: &solana_pubkey::Pubkey, + compressed_data: &Self::CompressedData, + solana_account: &solana_account_info::AccountInfo<#lifetime>, + seed_params: std::option::Option<&Self::SeedParams>, + zc_info: &mut light_sdk::interface::ZCompressedAccountInfoMut<'c>, + ) -> std::result::Result { + let post_system_offset = cpi_accounts.system_accounts_end_offset(); + let all_infos = cpi_accounts.account_infos(); + let post_system_accounts = &all_infos[post_system_offset..]; + let program_id = crate::ID; + let compressed_meta = &compressed_data.meta; + let mut variant_seed_params = SeedParams::default(); + let _ = &variant_seed_params; // Suppress unused warning when no params-only fields + + match &compressed_data.data { + #(#create_pda_match_arms)* + LightAccountVariant::PackedCTokenData(_) => { + // Tokens are handled separately, skip here + std::result::Result::Ok(false) + } + LightAccountVariant::CTokenData(_) => { + return std::result::Result::Err(light_sdk::error::LightSdkError::UnexpectedUnpackedVariant.into()); + } + #empty_variant_arm_create + } } #[inline(never)] diff --git a/sdk-libs/sdk/src/interface/decompress_idempotent.rs b/sdk-libs/sdk/src/interface/decompress_idempotent.rs index aa41ec0c98..86247c5266 100644 --- a/sdk-libs/sdk/src/interface/decompress_idempotent.rs +++ b/sdk-libs/sdk/src/interface/decompress_idempotent.rs @@ -1,7 +1,8 @@ #![allow(clippy::all)] // TODO: Remove. use light_compressed_account::{ - address::derive_address, instruction_data::with_account_info::OutAccountInfo, + address::derive_address, + instruction_data::with_account_info::{CompressedAccountInfo, OutAccountInfo}, }; use light_compressible::DECOMPRESSED_PDA_DISCRIMINATOR; use light_hasher::{sha256::Sha256BE, Hasher}; @@ -93,6 +94,53 @@ fn invoke_create_account_with_heap<'info>( .map_err(|e| LightSdkError::ProgramError(e)) } +/// Maximum number of seeds for PDA derivation. +pub const MAX_SEEDS: usize = 16; + +/// Convert Vec seeds to fixed array and call PDA creation. +/// Isolated in separate function to reduce stack usage (seed_refs array is on its own frame). +#[inline(never)] +#[cfg(feature = "v2")] +pub fn prepare_account_for_decompression_with_vec_seeds<'a, 'info, T>( + program_id: &Pubkey, + data: T, + compressed_meta: CompressedAccountMeta, + solana_account: &AccountInfo<'info>, + rent_sponsor: &AccountInfo<'info>, + cpi_accounts: &CpiAccounts<'a, 'info>, + seeds_vec: &[Vec], +) -> Result< + Option, + LightSdkError, +> +where + T: Clone + + crate::account::Size + + LightDiscriminator + + Default + + AnchorSerialize + + AnchorDeserialize + + HasCompressionInfo + + 'info, +{ + // Convert Vec seeds to fixed array on this stack frame + let mut seed_refs: [&[u8]; MAX_SEEDS] = [&[]; MAX_SEEDS]; + let len = seeds_vec.len().min(MAX_SEEDS); + for j in 0..len { + seed_refs[j] = seeds_vec[j].as_slice(); + } + + prepare_account_for_decompression_idempotent( + program_id, + data, + compressed_meta, + solana_account, + rent_sponsor, + cpi_accounts, + &seed_refs[..len], + ) +} + /// Helper function to decompress a compressed account into a PDA /// idempotently with seeds. #[inline(never)] @@ -171,3 +219,141 @@ where Ok(Some(account_info_result)) } + +/// Verify derived PDA matches the expected account. +/// Isolated function to reduce stack usage. +#[inline(never)] +#[cfg(feature = "v2")] +pub fn verify_pda_match( + derived_pda: &Pubkey, + expected: &Pubkey, +) -> Result<(), crate::error::LightSdkError> { + if *derived_pda != *expected { + msg!( + "Derived PDA does not match: expected {:?}, got {:?}", + expected, + derived_pda + ); + return Err(crate::error::LightSdkError::ConstraintViolation); + } + Ok(()) +} + +/// Zero-copy variant: derive seeds, verify PDA, create account, and write directly to zero-copy buffer. +/// This avoids returning intermediate structs that cause stack overflow. +/// Takes compressed_meta_no_address by reference and computes compressed_meta internally to reduce caller stack usage. +#[inline(never)] +#[cfg(all(feature = "v2", feature = "cpi-context"))] +pub fn derive_verify_create_and_write_pda<'a, 'info, 'c, T, CtxSeeds, SeedParams>( + program_id: &Pubkey, + data: &T, + ctx_seeds: &CtxSeeds, + seed_params: Option<&SeedParams>, + default_params: &SeedParams, + compressed_meta_no_address: &CompressedAccountMetaNoLamportsNoAddress, + address_space: &Pubkey, + solana_account: &AccountInfo<'info>, + rent_sponsor: &AccountInfo<'info>, + cpi_accounts: &CpiAccounts<'a, 'info>, + zc_info: &mut light_compressed_account::instruction_data::with_account_info::ZCompressedAccountInfoMut<'c>, +) -> Result +where + T: Clone + + crate::account::Size + + LightDiscriminator + + Default + + AnchorSerialize + + AnchorDeserialize + + HasCompressionInfo + + crate::interface::PdaSeedDerivation + + 'info, +{ + // Derive PDA seeds (keeps seeds_vec on this stack frame) + let (seeds_vec, derived_pda) = if let Some(params) = seed_params { + T::derive_pda_seeds_with_accounts(data, program_id, ctx_seeds, params)? + } else { + T::derive_pda_seeds_with_accounts(data, program_id, ctx_seeds, default_params)? + }; + + // Verify PDA matches + verify_pda_match(&derived_pda, solana_account.key)?; + + // Compute compressed_meta with address inside this frame (not in caller) + let compressed_meta = into_compressed_meta_with_address( + compressed_meta_no_address, + solana_account, + *address_space, + program_id, + ); + + // Create PDA (seeds_vec is converted to refs here and dropped after) + // Clone data since prepare_account_for_decompression_with_vec_seeds takes ownership + let result = prepare_account_for_decompression_with_vec_seeds( + program_id, + data.clone(), + compressed_meta, + solana_account, + rent_sponsor, + cpi_accounts, + &seeds_vec, + )?; + + // Write directly to zero-copy buffer if created + if let Some(account_info) = result { + write_to_zero_copy_buffer(account_info, zc_info)?; + Ok(true) + } else { + Ok(false) + } +} + +/// Write CompressedAccountInfo directly to zero-copy buffer. +/// Isolated function to reduce stack usage. +#[inline(never)] +#[cfg(all(feature = "v2", feature = "cpi-context"))] +fn write_to_zero_copy_buffer( + account_info: CompressedAccountInfo, + zc_info: &mut light_compressed_account::instruction_data::with_account_info::ZCompressedAccountInfoMut<'_>, +) -> Result<(), solana_program_error::ProgramError> { + // Extract address + let address = account_info + .address + .ok_or(solana_program_error::ProgramError::InvalidAccountData)?; + + // Write address to zero-copy buffer + if let Some(ref mut zc_addr) = zc_info.address { + zc_addr.copy_from_slice(&address); + } + + // Extract and write input + if let Some(input) = account_info.input { + if let Some(ref mut zc_input) = zc_info.input { + zc_input.discriminator = input.discriminator; + zc_input.data_hash = input.data_hash; + zc_input.merkle_context.merkle_tree_pubkey_index = + input.merkle_context.merkle_tree_pubkey_index; + zc_input.merkle_context.queue_pubkey_index = input.merkle_context.queue_pubkey_index; + zc_input + .merkle_context + .leaf_index + .set(input.merkle_context.leaf_index); + zc_input.merkle_context.prove_by_index = input.merkle_context.prove_by_index as u8; + zc_input.root_index.set(input.root_index); + zc_input.lamports.set(input.lamports); + } + } + + // Extract and write output + if let Some(output) = account_info.output { + if let Some(ref mut zc_output) = zc_info.output { + zc_output.discriminator = output.discriminator; + zc_output.data_hash = output.data_hash; + zc_output.output_merkle_tree_index = output.output_merkle_tree_index; + zc_output.lamports.set(output.lamports); + // Output data is the PDA pubkey (same as address for decompressed PDAs) + zc_output.data.copy_from_slice(&address); + } + } + + Ok(()) +} diff --git a/sdk-libs/sdk/src/interface/decompress_runtime.rs b/sdk-libs/sdk/src/interface/decompress_runtime.rs index 3e4a3859da..da0f2d87a1 100644 --- a/sdk-libs/sdk/src/interface/decompress_runtime.rs +++ b/sdk-libs/sdk/src/interface/decompress_runtime.rs @@ -1,21 +1,35 @@ //! Traits and processor for decompress_accounts_idempotent instruction. -use light_compressed_account::instruction_data::{ - cpi_context::CompressedCpiContext, - with_account_info::{CompressedAccountInfo, InstructionDataInvokeCpiWithAccountInfo}, +//! +//! This module implements a zero-allocation two-pass approach for PDA decompression: +//! - Pass 1: Count PDAs, collect pda_indices, collect tokens +//! - Pass 2: Allocate CPI buffer, create PDAs, write directly to zero-copy buffer +//! +//! Stack usage is minimal (~128 bytes for pda_indices array). +use light_compressed_account::{ + discriminators::INVOKE_CPI_WITH_ACCOUNT_INFO_INSTRUCTION, + instruction_data::{ + compressed_proof::CompressedProof, + cpi_context::CompressedCpiContext, + with_account_info::{ + CompressedAccountInfoConfig, InAccountInfoConfig, + InstructionDataInvokeCpiWithAccountInfo, InstructionDataInvokeCpiWithAccountInfoConfig, + OutAccountInfoConfig, ZCompressedAccountInfoMut, + }, + }, }; use light_sdk_types::{ - cpi_accounts::CpiAccountsConfig, cpi_context_write::CpiContextWriteAccounts, - instruction::account_meta::CompressedAccountMetaNoLamportsNoAddress, CpiSigner, + cpi_accounts::CpiAccountsConfig, cpi_context_write::CpiContextWriteAccounts, CpiSigner, }; +use light_zero_copy::{traits::ZeroCopyAtMut, ZeroCopyNew}; use solana_account_info::AccountInfo; -use solana_msg::msg; use solana_program_error::ProgramError; use solana_pubkey::Pubkey; -use crate::{ - cpi::{v2::CpiAccounts, InvokeLightSystemProgram}, - AnchorDeserialize, AnchorSerialize, LightDiscriminator, -}; +use crate::cpi::v2::CpiAccounts; + +/// Maximum number of accounts that can be decompressed in a single instruction. +/// Stack usage: pda_indices[usize; 16] = 128 bytes +pub const MAX_DECOMPRESS_ACCOUNTS: usize = 16; /// Trait for account variants that can be checked for token or PDA type. pub trait HasTokenVariant { @@ -38,7 +52,11 @@ pub trait TokenSeedProvider: Copy { ) -> Result<(Vec>, Pubkey), ProgramError>; } -/// Context trait for decompression. +/// Context trait for decompression with zero-allocation two-pass approach. +/// +/// The trait provides two key methods for the two-pass approach: +/// - `collect_layout_and_tokens`: Pass 1 - count PDAs, collect output_data_lens, collect tokens +/// - `create_and_write_pda`: Pass 2 - create PDA on Solana, write to zero-copy buffer pub trait DecompressContext<'info> { /// The compressed account data type (wraps program's variant enum) type CompressedData: HasTokenVariant; @@ -61,20 +79,57 @@ pub trait DecompressContext<'info> { fn token_cpi_authority(&self) -> Option<&AccountInfo<'info>>; fn token_config(&self) -> Option<&AccountInfo<'info>>; - /// Collect and unpack compressed accounts into PDAs and tokens. + /// Pass 1: Count non-initialized PDAs and collect tokens. + /// + /// Iterates through compressed accounts and: + /// - For each PDA: checks if already initialized, if not: stores index in pda_indices + /// - For each token: collects into the returned Vec + /// + /// # Arguments + /// * `compressed_accounts` - The compressed accounts to process + /// * `solana_accounts` - The corresponding Solana accounts (to check if already initialized) + /// * `pda_indices` - Array to fill with indices of PDAs that need decompression + /// + /// # Returns + /// * `pda_count` - Number of PDAs that need decompression (not already initialized) + /// * `tokens` - Vec of (PackedTokenData, CompressedMeta) for token accounts #[allow(clippy::type_complexity)] + fn collect_layout_and_tokens( + &self, + compressed_accounts: &[Self::CompressedData], + solana_accounts: &[AccountInfo<'info>], + pda_indices: &mut [usize; MAX_DECOMPRESS_ACCOUNTS], + ) -> Result<(usize, Vec<(Self::PackedTokenData, Self::CompressedMeta)>), ProgramError>; + + /// Pass 2: Create single PDA on Solana and write directly to zero-copy buffer. + /// + /// This method: + /// 1. Unpacks the compressed data to get the account state + /// 2. Derives PDA seeds and verifies against solana_account + /// 3. Creates the PDA account on Solana (via create_account CPI) + /// 4. Writes address, input, and output directly to the zero-copy buffer + /// + /// # Arguments + /// * `cpi_accounts` - CPI accounts for system program calls + /// * `address_space` - The address space for deriving compressed addresses + /// * `compressed_data` - The compressed account data to decompress (includes meta) + /// * `solana_account` - The target Solana PDA account + /// * `seed_params` - Optional seed parameters for PDA derivation + /// * `zc_info` - Mutable reference to zero-copy buffer to write directly into + /// + /// # Returns + /// * `true` if PDA was created and written + /// * `false` if account was already initialized (idempotent skip) #[allow(clippy::too_many_arguments)] - fn collect_pda_and_token<'b>( + fn create_and_write_pda<'b, 'c>( &self, cpi_accounts: &CpiAccounts<'b, 'info>, - address_space: Pubkey, - compressed_accounts: Vec, - solana_accounts: &[AccountInfo<'info>], + address_space: &Pubkey, + compressed_data: &Self::CompressedData, + solana_account: &AccountInfo<'info>, seed_params: Option<&Self::SeedParams>, - ) -> Result<( - Vec, - Vec<(Self::PackedTokenData, Self::CompressedMeta)> - ), ProgramError>; + zc_info: &mut ZCompressedAccountInfoMut<'c>, + ) -> Result; /// Process token decompression. #[allow(clippy::too_many_arguments)] @@ -105,6 +160,134 @@ pub trait PdaSeedDerivation { ) -> Result<(Vec>, Pubkey), ProgramError>; } +/// Output data length for decompressed PDAs (always 32 bytes = PDA pubkey). +pub const PDA_OUTPUT_DATA_LEN: u32 = 32; + +/// Build CPI config for PDA decompression. +/// +/// For PDA decompression, all accounts have identical config structure: +/// - has_address: true +/// - has_input: true +/// - has_output: true +/// - output_data_len: 32 (PDA pubkey) +/// +/// # Arguments +/// * `pda_count` - Number of PDAs to decompress +/// * `has_proof` - Whether a validity proof is included +/// +/// # Returns +/// `InstructionDataInvokeCpiWithAccountInfoConfig` ready for `byte_len()` and `new_zero_copy()` +#[inline(never)] +pub fn build_decompression_cpi_config( + pda_count: usize, + has_proof: bool, +) -> InstructionDataInvokeCpiWithAccountInfoConfig { + let account_infos = (0..pda_count) + .map(|_| CompressedAccountInfoConfig { + address: (true, ()), + input: (true, InAccountInfoConfig { merkle_context: () }), + output: ( + true, + OutAccountInfoConfig { + data: PDA_OUTPUT_DATA_LEN, + }, + ), + }) + .collect(); + + InstructionDataInvokeCpiWithAccountInfoConfig { + cpi_context: (), + proof: (has_proof, ()), + new_address_params: vec![], + account_infos, + read_only_addresses: vec![], + read_only_accounts: vec![], + } +} + +/// Allocate CPI instruction bytes with discriminator. +/// +/// # Arguments +/// * `config` - The CPI config describing byte layout +/// +/// # Returns +/// A zeroed Vec with space for discriminator + instruction data +#[inline(never)] +pub fn allocate_decompress_cpi_bytes( + config: &InstructionDataInvokeCpiWithAccountInfoConfig, +) -> Result, ProgramError> { + let data_len = InstructionDataInvokeCpiWithAccountInfo::byte_len(config) + .map_err(|_| ProgramError::InvalidAccountData)?; + let mut cpi_bytes = vec![0u8; data_len + 8]; + cpi_bytes[0..8].copy_from_slice(&INVOKE_CPI_WITH_ACCOUNT_INFO_INSTRUCTION); + Ok(cpi_bytes) +} + +/// Core CPI invocation to light-system-program. +/// Shared by both direct execution and CPI context write paths. +#[inline(never)] +fn invoke_light_system_cpi( + account_metas: Vec, + account_infos: &[AccountInfo<'_>], + cpi_bytes: Vec, + bump: u8, +) -> Result<(), ProgramError> { + use light_sdk_types::constants::{CPI_AUTHORITY_PDA_SEED, LIGHT_SYSTEM_PROGRAM_ID}; + + let instruction = solana_instruction::Instruction { + program_id: LIGHT_SYSTEM_PROGRAM_ID.into(), + accounts: account_metas, + data: cpi_bytes, + }; + + let signer_seeds = [CPI_AUTHORITY_PDA_SEED, &[bump]]; + solana_cpi::invoke_signed(&instruction, account_infos, &[signer_seeds.as_slice()]) +} + +/// Execute CPI to light-system-program with full account set. +#[inline(never)] +pub fn execute_cpi_invoke_sdk<'info>( + cpi_accounts: &CpiAccounts<'_, 'info>, + cpi_bytes: Vec, + bump: u8, +) -> Result<(), ProgramError> { + let account_metas = crate::cpi::v2::lowlevel::to_account_metas(cpi_accounts)?; + let account_infos = cpi_accounts.to_account_infos(); + invoke_light_system_cpi(account_metas, &account_infos, cpi_bytes, bump) +} + +/// Execute CPI to write to CPI context (minimal accounts). +#[inline(never)] +pub fn execute_cpi_write_to_context<'info>( + accounts: &CpiContextWriteAccounts<'_, AccountInfo<'info>>, + cpi_bytes: Vec, + bump: u8, +) -> Result<(), ProgramError> { + let account_metas = vec![ + crate::AccountMeta { + pubkey: *accounts.fee_payer.key, + is_writable: true, + is_signer: true, + }, + crate::AccountMeta { + pubkey: *accounts.authority.key, + is_writable: false, + is_signer: true, + }, + crate::AccountMeta { + pubkey: *accounts.cpi_context.key, + is_writable: true, + is_signer: false, + }, + ]; + let account_infos = [ + accounts.fee_payer.clone(), + accounts.authority.clone(), + accounts.cpi_context.clone(), + ]; + invoke_light_system_cpi(account_metas, &account_infos, cpi_bytes, bump) +} + /// Check what types of accounts are in the batch. /// Returns (has_tokens, has_pdas). #[inline(never)] @@ -123,89 +306,50 @@ pub fn check_account_types(compressed_accounts: &[T]) -> (bo (has_tokens, has_pdas) } -/// Handler for unpacking and preparing a single PDA variant for decompression. +/// Populate CPI struct metadata (mode, bump, program_id, proof, cpi_context). #[inline(never)] -#[allow(clippy::too_many_arguments)] -pub fn handle_packed_pda_variant<'a, 'b, 'info, T, P, A, S>( - accounts_rent_sponsor: &AccountInfo<'info>, - cpi_accounts: &CpiAccounts<'b, 'info>, - address_space: Pubkey, - solana_account: &AccountInfo<'info>, - index: usize, - packed: &P, - meta: &CompressedAccountMetaNoLamportsNoAddress, - post_system_accounts: &[AccountInfo<'info>], - compressed_pda_infos: &mut Vec, - program_id: &Pubkey, - seed_accounts: &A, - seed_params: Option<&S>, -) -> Result<(), ProgramError> -where - T: PdaSeedDerivation - + Clone - + crate::account::Size - + LightDiscriminator - + Default - + AnchorSerialize - + AnchorDeserialize - + crate::interface::HasCompressionInfo - + 'info, - P: crate::interface::Unpack, - S: Default, -{ - let data: T = P::unpack(packed, post_system_accounts)?; - - let (seeds_vec, derived_pda) = if let Some(params) = seed_params { - data.derive_pda_seeds_with_accounts(program_id, seed_accounts, params)? - } else { - let default_params = S::default(); - data.derive_pda_seeds_with_accounts(program_id, seed_accounts, &default_params)? - }; - if derived_pda != *solana_account.key { - msg!( - "Derived PDA does not match account at index {}: expected {:?}, got {:?}, seeds: {:?}", - index, - solana_account.key, - derived_pda, - seeds_vec - ); - return Err(ProgramError::from( - crate::error::LightSdkError::ConstraintViolation, - )); - } - - let compressed_infos = { - // Use fixed-size array to avoid heap allocation (MAX_SEEDS = 16) - const MAX_SEEDS: usize = 16; - let mut seed_refs: [&[u8]; MAX_SEEDS] = [&[]; MAX_SEEDS]; - let len = seeds_vec.len().min(MAX_SEEDS); - for i in 0..len { - seed_refs[i] = seeds_vec[i].as_slice(); +fn populate_cpi_metadata<'a>( + cpi_struct: &mut >::ZeroCopyAtMut, + bump: u8, + invoking_program_id: &Pubkey, + proof: Option<&CompressedProof>, + cpi_context: &CompressedCpiContext, + with_cpi_context: bool, +) { + cpi_struct.mode = 1; + cpi_struct.bump = bump; + cpi_struct.invoking_program_id = invoking_program_id.to_bytes().into(); + cpi_struct.compress_or_decompress_lamports = 0u64.into(); + cpi_struct.is_compress = 0; + cpi_struct.with_cpi_context = with_cpi_context as u8; + cpi_struct.with_transaction_hash = 0; + cpi_struct.cpi_context.cpi_context_account_index = cpi_context.cpi_context_account_index; + cpi_struct.cpi_context.first_set_context = cpi_context.first_set_context as u8; + cpi_struct.cpi_context.set_context = cpi_context.set_context as u8; + + if let Some(input_proof) = proof { + if let Some(ref mut proof_ref) = cpi_struct.proof { + proof_ref.a = input_proof.a; + proof_ref.b = input_proof.b; + proof_ref.c = input_proof.c; } - crate::interface::decompress_idempotent::prepare_account_for_decompression_idempotent::( - program_id, - data, - crate::interface::decompress_idempotent::into_compressed_meta_with_address( - meta, - solana_account, - address_space, - program_id, - ), - solana_account, - accounts_rent_sponsor, - cpi_accounts, - &seed_refs[..len], - )? - }; - compressed_pda_infos.extend(compressed_infos); - Ok(()) + } } -/// Processor for decompress_accounts_idempotent. +/// Processor for decompress_accounts_idempotent using zero-allocation two-pass approach. +/// +/// This function implements the two-pass approach for minimal stack usage: +/// - Pass 1: Count PDAs, collect pda_indices, collect tokens +/// - Pass 2: Allocate CPI buffer, create PDAs, write directly to zero-copy buffer /// /// CPI context batching rules: /// - Can use inputs from N trees /// - All inputs must use the FIRST CPI context account of the FIRST input +/// +/// # Stack Usage +/// - pda_indices: [usize; 16] = 128 bytes +/// - Counters = ~16 bytes +/// - Total: ~144 bytes (acceptable) #[inline(never)] #[allow(clippy::too_many_arguments)] pub fn process_decompress_accounts_idempotent<'info, Ctx>( @@ -236,7 +380,6 @@ where } // Use CPI context batching when we have both PDAs and tokens - // CPI context can handle inputs from N trees - all use FIRST cpi context of FIRST input let needs_cpi_context = has_tokens && has_pdas; let cpi_accounts = if needs_cpi_context { CpiAccounts::new_with_config( @@ -260,15 +403,14 @@ where .get(pda_accounts_start..) .ok_or_else(|| ProgramError::from(crate::error::LightSdkError::ConstraintViolation))?; - let (compressed_pda_infos, compressed_token_accounts) = ctx.collect_pda_and_token( - &cpi_accounts, - address_space, - compressed_accounts, - solana_accounts, - seed_params, - )?; + // Stack array for tracking PDA indices + let mut pda_indices: [usize; MAX_DECOMPRESS_ACCOUNTS] = [0; MAX_DECOMPRESS_ACCOUNTS]; - let has_pdas = !compressed_pda_infos.is_empty(); + // Pass 1: Collect layout and tokens (checks which PDAs need decompression) + let (pda_count, compressed_token_accounts) = + ctx.collect_layout_and_tokens(&compressed_accounts, solana_accounts, &mut pda_indices)?; + + let has_pdas = pda_count > 0; let has_tokens = !compressed_token_accounts.is_empty(); if !has_pdas && !has_tokens { @@ -277,27 +419,55 @@ where let fee_payer = ctx.fee_payer(); - // Process PDAs (if any) + // Process PDAs (if any) using zero-copy pattern with two-pass approach if has_pdas { + let cpi_signer_config = cpi_accounts.config().cpi_signer; + + // Build CPI config from pda_count + let cpi_config = build_decompression_cpi_config(pda_count, proof.0.is_some()); + + // Allocate CPI bytes (single acceptable allocation) + let mut cpi_bytes = allocate_decompress_cpi_bytes(&cpi_config)?; + + // Get zero-copy mutable struct + let (mut cpi_struct, _remaining) = + InstructionDataInvokeCpiWithAccountInfo::new_zero_copy(&mut cpi_bytes[8..], cpi_config) + .map_err(|_| ProgramError::InvalidAccountData)?; + + // Populate CPI metadata (mode, bump, program_id, proof, cpi_context) + let cpi_context_for_pdas = if has_tokens { + CompressedCpiContext::first() + } else { + CompressedCpiContext::default() + }; + populate_cpi_metadata( + &mut cpi_struct, + cpi_signer_config.bump, + &cpi_signer_config.program_id.into(), + proof.0.as_ref(), + &cpi_context_for_pdas, + has_tokens, // with_cpi_context + ); + + // Pass 2: Create PDAs and write directly to zero-copy buffer + let zc_account_infos = cpi_struct.account_infos.as_mut_slice(); + for (zc_idx, &account_idx) in pda_indices[..pda_count].iter().enumerate() { + let zc_info = &mut zc_account_infos[zc_idx]; + + // Create PDA and write directly to zero-copy buffer + ctx.create_and_write_pda( + &cpi_accounts, + &address_space, + &compressed_accounts[account_idx], + &solana_accounts[account_idx], + seed_params, + zc_info, + )?; + } + + // Execute CPI if !has_tokens { - // PDAs only - execute directly (manual construction to avoid extra allocations) - let cpi_signer_config = cpi_accounts.config().cpi_signer; - let instruction_data = InstructionDataInvokeCpiWithAccountInfo { - mode: 1, - bump: cpi_signer_config.bump, - invoking_program_id: cpi_signer_config.program_id.into(), - compress_or_decompress_lamports: 0, - is_compress: false, - with_cpi_context: false, - with_transaction_hash: false, - cpi_context: CompressedCpiContext::default(), - proof: proof.0, - new_address_params: Vec::new(), - account_infos: compressed_pda_infos, - read_only_addresses: Vec::new(), - read_only_accounts: Vec::new(), - }; - instruction_data.invoke(cpi_accounts.clone())?; + execute_cpi_invoke_sdk(&cpi_accounts, cpi_bytes, cpi_signer_config.bump)?; } else { // PDAs + tokens - write to CPI context first, tokens will execute let authority = cpi_accounts @@ -312,24 +482,7 @@ where cpi_context: cpi_context_account, cpi_signer, }; - - // Manual construction to avoid extra allocations - let instruction_data = InstructionDataInvokeCpiWithAccountInfo { - mode: 1, - bump: cpi_signer.bump, - invoking_program_id: cpi_signer.program_id.into(), - compress_or_decompress_lamports: 0, - is_compress: false, - with_cpi_context: true, - with_transaction_hash: false, - cpi_context: CompressedCpiContext::first(), - proof: proof.0, - new_address_params: Vec::new(), - account_infos: compressed_pda_infos, - read_only_addresses: Vec::new(), - read_only_accounts: Vec::new(), - }; - instruction_data.invoke_write_to_cpi_context_first(system_cpi_accounts)?; + execute_cpi_write_to_context(&system_cpi_accounts, cpi_bytes, cpi_signer.bump)?; } } diff --git a/sdk-libs/sdk/src/interface/mod.rs b/sdk-libs/sdk/src/interface/mod.rs index f6111f1c73..301f390cf6 100644 --- a/sdk-libs/sdk/src/interface/mod.rs +++ b/sdk-libs/sdk/src/interface/mod.rs @@ -34,13 +34,20 @@ pub use config::{ process_update_light_config, LightConfig, COMPRESSIBLE_CONFIG_SEED, MAX_ADDRESS_TREES_PER_SPACE, }; +#[cfg(all(feature = "v2", feature = "cpi-context"))] +pub use decompress_idempotent::derive_verify_create_and_write_pda; #[cfg(feature = "v2")] pub use decompress_idempotent::{ into_compressed_meta_with_address, prepare_account_for_decompression_idempotent, + prepare_account_for_decompression_with_vec_seeds, verify_pda_match, MAX_SEEDS, }; #[cfg(all(feature = "v2", feature = "cpi-context"))] pub use decompress_runtime::{ - check_account_types, handle_packed_pda_variant, process_decompress_accounts_idempotent, + build_decompression_cpi_config, check_account_types, process_decompress_accounts_idempotent, DecompressContext, HasTokenVariant, PdaSeedDerivation, TokenSeedProvider, + MAX_DECOMPRESS_ACCOUNTS, PDA_OUTPUT_DATA_LEN, }; +// Re-export ZCompressedAccountInfoMut for use in macro-generated code +#[cfg(all(feature = "v2", feature = "cpi-context"))] +pub use light_compressed_account::instruction_data::with_account_info::ZCompressedAccountInfoMut; pub use light_compressible::{rent, CreateAccountsProof};