Skip to content

refactor: zero copy decompress runtime #2216

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ananas-block wants to merge 2 commits into
base: main
Choose a base branch
from
+611 −232
Open

refactor: zero copy decompress runtime #2216

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e2f927e
first step
ananas-block Jan 24, 2026
ad7c229
refactor: decompress runtime to zero copy
ananas-block Jan 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 12 additions & 27 deletions program-libs/compressed-account/src/instruction_data/with_account_info.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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]>,
Expand Down Expand Up @@ -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.
Expand Down
172 changes: 110 additions & 62 deletions sdk-libs/macros/src/light_pdas/account/decompress_context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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};
Expand All @@ -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());
}
}
Comment on lines +23 to +40
Copy link
Contributor

@coderabbitai coderabbitai bot Jan 24, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Add bounds check before writing to pda_indices

If more than MAX_DECOMPRESS_ACCOUNTS PDAs are encountered, the current code will panic on out‑of‑bounds write. Guard the write and return a deterministic error instead.

Proposed fix 
-                    if solana_accounts[i].data_is_empty() {
-                        pda_indices[pda_count] = i;
-                        pda_count += 1;
-                    }
+                    if solana_accounts[i].data_is_empty() {
+                        if pda_count >= light_sdk::interface::MAX_DECOMPRESS_ACCOUNTS {
+                            return std::result::Result::Err(
+                                light_sdk::error::LightSdkError::ConstraintViolation.into()
+                            );
+                        }
+                        pda_indices[pda_count] = i;
+                        pda_count += 1;
+                    }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// 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());
}
}
// 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() {
if pda_count >= light_sdk::interface::MAX_DECOMPRESS_ACCOUNTS {
return std::result::Result::Err(
light_sdk::error::LightSdkError::ConstraintViolation.into()
);
}
pda_indices[pda_count] = i;
pda_count += 1;
}
}
LightAccountVariant::`#variant_name` { .. } => {
return std::result::Result::Err(light_sdk::error::LightSdkError::UnexpectedUnpackedVariant.into());
}
}
🤖 Prompt for AI Agents 
In `@sdk-libs/macros/src/light_pdas/account/decompress_context.rs` around lines 23
- 40, The match arm that writes into pda_indices in collect_layout_and_tokens
currently assumes pda_count < MAX_DECOMPRESS_ACCOUNTS and can panic; add a
bounds check before assigning pda_indices[pda_count] and incrementing pda_count:
if pda_count >= MAX_DECOMPRESS_ACCOUNTS return a deterministic error (e.g. a new
or existing LightSdkError variant indicating too many PDAs) instead of writing
out-of-bounds. Update the arm for LightAccountVariant::`#packed_variant_name` { ..
} to perform the check, and reference pda_indices, pda_count and
MAX_DECOMPRESS_ACCOUNTS when implementing the guard so the function never panics
on overflow.

})
.collect();

Comment on lines +23 to +43
Copy link
Contributor

@coderabbitai coderabbitai bot Jan 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Guard against >MAX_DECOMPRESS_ACCOUNTS to avoid out‑of‑bounds writes.

The generated code writes pda_indices[pda_count] without a capacity check. More than 16 PDAs will panic at runtime. Add a guard and return an error when the limit is exceeded.

🛠️ Suggested fix 
-                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::`#packed_variant_name` { .. } => {
+                    // PDA variant: only count if not already initialized (idempotent check)
+                    if solana_accounts[i].data_is_empty() {
+                        if pda_count >= light_sdk::interface::MAX_DECOMPRESS_ACCOUNTS {
+                            return std::result::Result::Err(
+                                light_sdk::error::LightSdkError::ConstraintViolation.into(),
+                            );
+                        }
+                        pda_indices[pda_count] = i;
+                        pda_count += 1;
+                    }
+                }

Also applies to: 187-214

🤖 Prompt for AI Agents 
In `@sdk-libs/macros/src/light_pdas/account/decompress_context.rs` around lines 23
- 43, The generated match arms write into pda_indices[pda_count] without
checking bounds, which can panic if more than MAX_DECOMPRESS_ACCOUNTS PDAs are
encountered; update the code in the match arms generated for
collect_layout_and_tokens (the block using pda_indices, pda_count) to first
check if pda_count >= MAX_DECOMPRESS_ACCOUNTS and return an appropriate error
(e.g., LightSdkError::TooManyDecompressAccounts or map to an existing error)
before assigning to pda_indices and incrementing pda_count; apply the same guard
to the other similar generated block around lines 187-214 so both sites perform
the capacity check and return on overflow instead of writing out of bounds.

// 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);
Expand All @@ -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;
}
Expand All @@ -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());
Expand All @@ -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);
}
Expand Down Expand Up @@ -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
}
}
Comment on lines +216 to 246
Copy link
Contributor

@coderabbitai coderabbitai bot Jan 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Prevent panics when slicing post‑system accounts.

&all_infos[post_system_offset..] will panic on malformed account lists. Use get() and return a ProgramError instead.

🛠️ Suggested fix 
-                let post_system_accounts = &all_infos[post_system_offset..];
+                let post_system_accounts = all_infos
+                    .get(post_system_offset..)
+                    .ok_or(solana_program_error::ProgramError::NotEnoughAccountKeys)?;
🤖 Prompt for AI Agents 
In `@sdk-libs/macros/src/light_pdas/account/decompress_context.rs` around lines
216 - 246, The code panics when slicing all_infos with
&all_infos[post_system_offset..]; change that to check bounds via
all_infos.get(post_system_offset..) and handle the None case by returning an
appropriate ProgramError instead of panicking. In the create_and_write_pda
function, replace let post_system_accounts = &all_infos[post_system_offset..];
with a match or if-let on all_infos.get(post_system_offset..) binding
post_system_accounts on Some and returning Err(...) on None (for example
Err(light_sdk::error::LightSdkError::UnexpectedAccountListLength.into()) or
another suitable ProgramError).


#[inline(never)]
Expand Down
Loading
Loading