refactor: decompress runtime to zero copy

Author: ananas-block

program-libs/compressed-account/src/instruction_data/with_account_info.rs

@@ -117,6 +117,27 @@ pub struct OutAccountInfo {
     pub data: Vec<u8>,
 }
 
+/// Fixed-size output metadata WITHOUT the data Vec.
+/// Used for zero-allocation PDA collection in two-pass decompression.
+#[derive(Debug, Default, Clone, Copy, PartialEq)]
+pub struct OutAccountMeta {
+    pub discriminator: [u8; 8],
+    pub data_hash: [u8; 32],
+    pub output_merkle_tree_index: u8,
+    pub lamports: u64,
+}
+
+impl From<&OutAccountInfo> for OutAccountMeta {
+    fn from(info: &OutAccountInfo) -> Self {
+        Self {
+            discriminator: info.discriminator,
+            data_hash: info.data_hash,
+            output_merkle_tree_index: info.output_merkle_tree_index,
+            lamports: info.lamports,
+        }
+    }
+}
+
 impl TryFrom<OutputCompressedAccountWithPackedContext> for OutAccountInfo {
     type Error = CompressedAccountError;
 

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<TokenStream> {
-    // 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<Self::CompressedData>,
+                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<bool, solana_program_error::ProgramError> {
+                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)]