fix tests, fix: account compression cpi size calculation

Author: ananas-block

anchor-programs/system/src/cpi_context_account.rs

@@ -41,15 +41,15 @@ impl CpiContextAccount {
     }
 }
 
-pub struct ZCpiContextAccount<'a> {
+pub struct ZCpiContextAccount2<'a> {
     pub fee_payer: Ref<&'a mut [u8], light_compressed_account::pubkey::Pubkey>,
     pub associated_merkle_tree: Ref<&'a mut [u8], light_compressed_account::pubkey::Pubkey>,
     pub context: Vec<ZInstructionDataInvokeCpi<'a>>,
 }
 
 pub fn deserialize_cpi_context_account<'info, 'a>(
     account_info: &AccountInfo<'info>,
-) -> std::result::Result<ZCpiContextAccount<'a>, ZeroCopyError> {
+) -> std::result::Result<ZCpiContextAccount2<'a>, ZeroCopyError> {
     let mut account_data = account_info.try_borrow_mut_data().unwrap();
     let data = unsafe { slice::from_raw_parts_mut(account_data.as_mut_ptr(), account_data.len()) };
     let (fee_payer, data) =
@@ -64,7 +64,7 @@ pub fn deserialize_cpi_context_account<'info, 'a>(
         context.push(context_item);
         data = new_data;
     }
-    Ok(ZCpiContextAccount {
+    Ok(ZCpiContextAccount2 {
         fee_payer,
         associated_merkle_tree,
         context,

cli/accounts/cpi_context_cpi1uHzrEhBG733DoEJNgHCyRS3XmmyVNZx5fonubE4.json

@@ -110,6 +110,15 @@ impl<'info, T: AccountInfoTrait> AccountIterator<'info, T> {
         Ok(account_info)
     }
 
+    #[inline(always)]
+    #[track_caller]
+    pub fn next_signer_non_mut(&mut self, account_name: &str) -> Result<&'info T, AccountError> {
+        let account_info = self.next_signer(account_name)?;
+        check_non_mut(account_info)
+            .inspect_err(|e| self.print_on_error(e, account_name, Location::caller()))?;
+        Ok(account_info)
+    }
+
     #[inline(always)]
     #[track_caller]
     pub fn next_non_mut(&mut self, account_name: &str) -> Result<&'info T, AccountError> {

cli/accounts/cpi_context_cpi2cdhkH5roePvcudTgUL8ppEBfTay1desGh8G8QxK.json

@@ -8,7 +8,7 @@ use light_program_test::{
     ProgramTestConfig,
 };
 use light_sdk::constants::{
-    CPI_CONTEXT_ACCOUNT_DISCRIMINATOR, CPI_CONTEXT_ACCOUNT_DISCRIMINATOR_V1,
+    CPI_CONTEXT_ACCOUNT_2_DISCRIMINATOR, CPI_CONTEXT_ACCOUNT_DISCRIMINATOR,
 };
 use light_system_program_pinocchio::cpi_context::state::deserialize_cpi_context_account;
 use light_test_utils::{legacy_cpi_context_account::get_legacy_cpi_context_account, Rpc};
@@ -63,7 +63,7 @@ async fn test_re_init_cpi_account() {
             .unwrap();
         assert_eq!(
             &pre_account.data[0..8],
-            &CPI_CONTEXT_ACCOUNT_DISCRIMINATOR_V1,
+            &CPI_CONTEXT_ACCOUNT_2_DISCRIMINATOR,
             "Account should have legacy discriminator"
         );
         assert_eq!(

program-libs/account-checks/src/account_iterator.rs

@@ -471,14 +471,16 @@ pub fn cpi_execute_compressed_transaction_transfer<
         );
         data.extend(inputs);
 
-        // 4 static accounts
-        let accounts_len = 4 + remaining_accounts.len() + cpi_context.is_some() as usize;
+        // 6 static accounts
+        let accounts_len = 6 + remaining_accounts.len() + cpi_context.is_some() as usize;
         let mut account_infos = Vec::with_capacity(accounts_len);
         let mut account_metas = Vec::with_capacity(accounts_len);
         account_infos.push(ctx.get_fee_payer().to_account_info());
         account_infos.push(cpi_authority_pda);
         account_infos.push(ctx.get_registered_program_pda().to_account_info());
         account_infos.push(ctx.get_account_compression_authority().to_account_info());
+        account_infos.push(ctx.get_account_compression_program().to_account_info());
+        account_infos.push(ctx.get_system_program().to_account_info());
 
         account_metas.push(AccountMeta {
             pubkey: account_infos[0].key(),

program-tests/system-test/tests/test_re_init_cpi_account.rs

@@ -14,7 +14,7 @@ use light_program_profiler::profile;
 use pinocchio::{account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey};
 
 use crate::{
-    cpi_context::state::ZCpiContextAccount,
+    cpi_context::state::ZCpiContextAccount2,
     processor::sol_compression::{SOL_POOL_PDA_BUMP, SOL_POOL_PDA_SEED},
     Result,
 };
@@ -104,7 +104,7 @@ pub fn check_anchor_option_cpi_context_account(
     } else {
         {
             check_owner(&crate::ID, option_cpi_context_account)?;
-            check_discriminator::<ZCpiContextAccount>(
+            check_discriminator::<ZCpiContextAccount2>(
                 option_cpi_context_account.try_borrow_data()?.as_ref(),
             )?;
         }
@@ -149,7 +149,7 @@ pub fn check_option_cpi_context_account<'a>(
                 location.column()
             )
         })?;
-        check_discriminator::<ZCpiContextAccount>(account_info.try_borrow_data()?.as_ref())
+        check_discriminator::<ZCpiContextAccount2>(account_info.try_borrow_data()?.as_ref())
             .inspect_err(|_| {
                 let location = Location::caller();
                 solana_msg::msg!(

programs/compressed-token/src/process_transfer.rs

@@ -11,9 +11,7 @@ use light_program_profiler::profile;
 use pinocchio::{account_info::AccountInfo, program_error::ProgramError};
 
 use crate::{
-    cpi_context::state::{
-        cpi_context_account_new, CpiContextAccountInitParams, CpiContextAccountLegacy,
-    },
+    cpi_context::state::{cpi_context_account_new, CpiContextAccount, CpiContextAccountInitParams},
     errors::SystemProgramError,
     Result,
 };
@@ -62,7 +60,7 @@ pub fn init_cpi_context_account(accounts: &[AccountInfo]) -> Result<()> {
     let ctx = InitializeCpiContextAccount::from_account_infos(accounts)?;
     let params: CpiContextAccountInitParams =
         CpiContextAccountInitParams::new(*ctx.associated_merkle_tree.key());
-    cpi_context_account_new::<false>(ctx.cpi_context_account, params)?;
+    cpi_context_account_new::<false>(ctx.cpi_context_account, params).map(|_| ())?;
 
     Ok(())
 }
@@ -75,13 +73,13 @@ pub fn reinit_cpi_context_account(accounts: &[AccountInfo]) -> Result<()> {
     let cpi_context_account = &accounts[0];
     let associated_merkle_tree = {
         let data = cpi_context_account.try_borrow_data()?;
-        CpiContextAccountLegacy::deserialize(&mut &data[8..])
+        CpiContextAccount::deserialize(&mut &data[8..])
             .map_err(|_| ProgramError::BorshIoError)?
             .associated_merkle_tree
     };
     let params: CpiContextAccountInitParams =
         CpiContextAccountInitParams::new(associated_merkle_tree);
-    cpi_context_account_new::<true>(cpi_context_account, params)?;
+    cpi_context_account_new::<true>(cpi_context_account, params).map(|_| ())?;
 
     Ok(())
 }

programs/system/src/accounts/account_checks.rs

@@ -12,7 +12,7 @@ pub const INVOKE_CPI_INSTRUCTION: [u8; 8] = [49, 212, 191, 129, 39, 194, 43, 196
 pub const INVOKE_CPI_WITH_READ_ONLY_INSTRUCTION: [u8; 8] = [86, 47, 163, 166, 21, 223, 92, 8];
 // discriminator of CpiContextAccount2
 pub const CPI_CONTEXT_ACCOUNT_DISCRIMINATOR: [u8; 8] = [34, 184, 183, 14, 100, 80, 183, 124];
-pub const CPI_CONTEXT_ACCOUNT_DISCRIMINATOR_V1: [u8; 8] = [22, 20, 149, 218, 74, 204, 128, 166];
+pub const CPI_CONTEXT_ACCOUNT_2_DISCRIMINATOR: [u8; 8] = [22, 20, 149, 218, 74, 204, 128, 166];
 pub const INVOKE_CPI_WITH_ACCOUNT_INFO_INSTRUCTION: [u8; 8] = [228, 34, 128, 84, 47, 139, 86, 240];
 pub const RE_INIT_CPI_CONTEXT_ACCOUNT_INSTRUCTION: [u8; 8] =
     [187, 147, 22, 142, 104, 180, 136, 190];

programs/system/src/accounts/init_context_account.rs

@@ -10,7 +10,7 @@ use light_program_profiler::profile;
 use pinocchio::{account_info::AccountInfo, instruction::AccountMeta, pubkey::Pubkey};
 
 use crate::{
-    cpi_context::state::ZCpiContextAccount, errors::SystemProgramError,
+    cpi_context::state::ZCpiContextAccount2, errors::SystemProgramError,
     utils::transfer_lamports_invoke, Result, MAX_OUTPUT_ACCOUNTS,
 };
 
@@ -140,7 +140,7 @@ impl<'info> SystemContext<'info> {
 #[derive(Debug)]
 pub struct WrappedInstructionData<'a, T: InstructionData<'a>> {
     instruction_data: T,
-    cpi_context: Option<ZCpiContextAccount<'a>>,
+    cpi_context: Option<ZCpiContextAccount2<'a>>,
     address_len: usize,
     input_len: usize,
     outputs_len: usize,
@@ -181,34 +181,36 @@ where
     }
 
     #[profile]
-    pub fn set_cpi_context(&mut self, cpi_context: ZCpiContextAccount<'a>) -> Result<()> {
+    pub fn set_cpi_context(&mut self, cpi_context: ZCpiContextAccount2<'a>) -> Result<()> {
         if self.cpi_context.is_none() {
             self.outputs_len += cpi_context.out_accounts.len();
             if self.outputs_len > MAX_OUTPUT_ACCOUNTS {
                 return Err(SystemProgramError::TooManyOutputAccounts.into());
             }
             self.address_len += cpi_context.new_addresses.len();
             self.input_len += cpi_context.in_accounts.len();
+
+            // Calculate offsets from the CPI context account
+            let (outputs_start_offset, outputs_end_offset) = cpi_context.calculate_output_offsets();
+
             self.cpi_context = Some(cpi_context);
-            // TODO: set correctly for correct cpi data allocation
-            // self.cpi_context_outputs_start_offset = outputs_start_offset;
-            // self.cpi_context_outputs_end_offset = outputs_end_offset;
+            self.cpi_context_outputs_start_offset = outputs_start_offset;
+            self.cpi_context_outputs_end_offset = outputs_end_offset;
         } else {
             return Err(SystemProgramError::CpiContextAlreadySet.into());
         }
         Ok(())
     }
-    // TODO: hardcode will be a standard value
+
     pub fn get_cpi_context_outputs_start_offset(&self) -> usize {
         self.cpi_context_outputs_start_offset
     }
 
-    // TODO: hardcode will be a standard value
     pub fn get_cpi_context_outputs_end_offset(&self) -> usize {
         self.cpi_context_outputs_end_offset
     }
 
-    pub fn get_cpi_context_account(&'b self) -> &'b Option<ZCpiContextAccount<'a>> {
+    pub fn get_cpi_context_account(&'b self) -> &'b Option<ZCpiContextAccount2<'a>> {
         &self.cpi_context
     }