refactor: extract into functions

Author: ananas-block

programs/compressed-token/program/src/compressed_token/transfer2/processor.rs

@@ -1,7 +1,9 @@
 use anchor_compressed_token::ErrorCode;
 use anchor_lang::prelude::ProgramError;
 use light_array_map::ArrayMap;
-use light_compressed_account::instruction_data::with_readonly::InstructionDataInvokeCpiWithReadOnly;
+use light_compressed_account::instruction_data::with_readonly::{
+    InstructionDataInvokeCpiWithReadOnly, ZInstructionDataInvokeCpiWithReadOnlyMut,
+};
 use light_program_profiler::profile;
 use light_token_interface::{
     hash_cache::HashCache,
@@ -256,59 +258,10 @@ fn process_with_system_program_cpi<'a>(
         accounts,
         mint_cache,
     )?;
-
-    // ATA decompress is permissionless and idempotent.
-    // Detect from: exactly 1 input, 1 Decompress compression, CompressedOnly with is_ata=true.
-    // Multi-input batches (including mixed ATA + non-ATA) are not idempotent.
-    let is_ata_decompress = inputs.in_token_data.len() == 1
-        && inputs
-            .compressions
-            .as_ref()
-            .is_some_and(|c| c.len() == 1 && c.iter().any(|c| c.mode.is_decompress()))
-        && inputs.in_tlv.as_ref().is_some_and(|tlvs| {
-            tlvs.iter().flatten().any(|ext| {
-                matches!(ext, ZExtensionInstructionData::CompressedOnly(data) if data.is_ata())
-            })
-        });
-
-    if is_ata_decompress {
-        let input_data = &inputs.in_token_data[0];
-        let merkle_context = &input_data.merkle_context;
-        let input_account = cpi_instruction_struct
-            .input_compressed_accounts
-            .first()
-            .ok_or(ProgramError::InvalidAccountData)?;
-
-        let owner_hashed = light_hasher::hash_to_field_size::hash_to_bn254_field_size_be(
-            &crate::LIGHT_CPI_SIGNER.program_id,
-        );
-        let tree_account = validated_accounts
-            .packed_accounts
-            .get_u8(merkle_context.merkle_tree_pubkey_index, "idempotent: tree")?;
-        let merkle_tree_hashed =
-            light_hasher::hash_to_field_size::hash_to_bn254_field_size_be(tree_account.key());
-
-        let lamports: u64 = (*input_account.lamports).into();
-        let account_hash = light_compressed_account::compressed_account::hash_with_hashed_values(
-            &lamports,
-            input_account.address.as_ref().map(|x| x.as_slice()),
-            Some((
-                input_account.discriminator.as_slice(),
-                input_account.data_hash.as_slice(),
-            )),
-            &owner_hashed,
-            &merkle_tree_hashed,
-            &merkle_context.leaf_index.get(),
-            true,
-        )
-        .map_err(ProgramError::from)?;
-
-        let mut tree =
-            light_batched_merkle_tree::merkle_tree::BatchedMerkleTreeAccount::state_from_account_info(tree_account)
-                .map_err(ProgramError::from)?;
-
-        if tree.check_input_queue_non_inclusion(&account_hash).is_err() {
-            // Account is in bloom filter -- already spent. Idempotent no-op.
+    let is_idempotent_ata_decompress = is_idempotent_ata_decompress(inputs);
+    #[allow(clippy::collapsible_if)]
+    if is_idempotent_ata_decompress {
+        if check_ata_decompress_idempotent(inputs, &cpi_instruction_struct, validated_accounts)? {
             return Ok(());
         }
     }
@@ -389,3 +342,76 @@ fn process_with_system_program_cpi<'a>(
     }
     Ok(())
 }
+
+/// Detect idempotent associated token account decompress:
+/// - exactly 1 input compressed token account with CompressedOnly extension is_ata=true
+/// - 1 Decompress compression.
+///
+/// Multi-input batches (including mixed ATA + non-ATA) are not idempotent.
+#[inline(always)]
+pub fn is_idempotent_ata_decompress(inputs: &ZCompressedTokenInstructionDataTransfer2) -> bool {
+    inputs.in_token_data.len() == 1
+        && inputs
+            .compressions
+            .as_ref()
+            .is_some_and(|c| c.len() == 1 && c.iter().any(|c| c.mode.is_decompress()))
+        && inputs.in_tlv.as_ref().is_some_and(|tlvs| {
+            tlvs.iter().flatten().any(|ext| {
+                matches!(ext, ZExtensionInstructionData::CompressedOnly(data) if data.is_ata())
+            })
+        })
+}
+
+/// Computes the compressed account hash and checks whether the hash exists in the input queue bloom filters.
+/// The account compression program inserts spent compressed accounts into the respective input queue.
+///
+/// if exists in bloom filter -> exit the ata was already decompressed
+/// else decompress
+#[cold]
+fn check_ata_decompress_idempotent(
+    inputs: &ZCompressedTokenInstructionDataTransfer2,
+    cpi_instruction_struct: &ZInstructionDataInvokeCpiWithReadOnlyMut<'_>,
+    validated_accounts: &Transfer2Accounts,
+) -> Result<bool, ProgramError> {
+    let input_data = inputs
+        .in_token_data
+        .first()
+        .ok_or(ProgramError::InvalidAccountData)?;
+    let merkle_context = &input_data.merkle_context;
+    let input_account = cpi_instruction_struct
+        .input_compressed_accounts
+        .first()
+        .ok_or(ProgramError::InvalidAccountData)?;
+
+    let owner_hashed = light_hasher::hash_to_field_size::hash_to_bn254_field_size_be(
+        &crate::LIGHT_CPI_SIGNER.program_id,
+    );
+    let tree_account = validated_accounts
+        .packed_accounts
+        .get_u8(merkle_context.merkle_tree_pubkey_index, "idempotent: tree")?;
+    let merkle_tree_hashed =
+        light_hasher::hash_to_field_size::hash_to_bn254_field_size_be(tree_account.key());
+
+    let lamports: u64 = input_account.lamports.get();
+    let account_hash = light_compressed_account::compressed_account::hash_with_hashed_values(
+        &lamports,
+        input_account.address.as_ref().map(|x| x.as_slice()),
+        Some((
+            input_account.discriminator.as_slice(),
+            input_account.data_hash.as_slice(),
+        )),
+        &owner_hashed,
+        &merkle_tree_hashed,
+        &merkle_context.leaf_index.get(),
+        true,
+    )
+    .map_err(ProgramError::from)?;
+
+    let mut tree =
+        light_batched_merkle_tree::merkle_tree::BatchedMerkleTreeAccount::state_from_account_info(
+            tree_account,
+        )
+        .map_err(ProgramError::from)?;
+
+    Ok(tree.check_input_queue_non_inclusion(&account_hash).is_err())
+}

programs/compressed-token/program/src/shared/token_input.rs

@@ -80,15 +80,15 @@ pub fn set_input_compressed_account<'a>(
     // For ATA decompress (is_ata=true), verify the wallet owner from owner_index instead
     // of the compressed account owner (which is the ATA pubkey that can't sign).
     // Also verify that owner_account (the ATA) matches the derived ATA from wallet_owner + mint + bump.
-    let (signer_account, is_ata_decompress) = if let Some(exts) = tlv_data {
+    let (signer_account, check_signer) = if let Some(exts) = tlv_data {
         resolve_ata_signer(exts, packed_accounts, mint_account, owner_account)?
     } else {
-        (owner_account, false)
+        (owner_account, true)
     };
 
     // ATA decompress is permissionless -- the destination is a deterministic PDA,
     // so there is no griefing vector. ATA derivation is still validated above.
-    if !is_ata_decompress {
+    if check_signer {
         verify_owner_or_delegate_signer(
             signer_account,
             delegate_account,
@@ -232,13 +232,13 @@ fn resolve_ata_signer<'a>(
                 if !pinocchio::pubkey::pubkey_eq(owner_account.key(), &derived_ata) {
                     return Err(TokenError::InvalidAtaDerivation.into());
                 }
-
-                return Ok((wallet_owner, true));
+                // Do not check signer the recipient token account is the correct ata.
+                return Ok((wallet_owner, false));
             }
         }
     }
 
-    Ok((owner_account, false))
+    Ok((owner_account, true))
 }
 
 #[cold]

programs/compressed-token/program/tests/transfer2_processor.rs

@@ -0,0 +1,157 @@
+use anchor_lang::AnchorSerialize;
+use light_compressed_token::compressed_token::transfer2::processor::is_idempotent_ata_decompress;
+use light_token_interface::instructions::{
+    extensions::{CompressedOnlyExtensionInstructionData, ExtensionInstructionData},
+    transfer2::{
+        CompressedTokenInstructionDataTransfer2, Compression, CompressionMode,
+        MultiInputTokenDataWithContext,
+    },
+};
+use light_zero_copy::traits::ZeroCopyAt;
+use rand::{rngs::StdRng, Rng, SeedableRng};
+
+fn serialize(data: &CompressedTokenInstructionDataTransfer2) -> Vec<u8> {
+    let mut buf = Vec::new();
+    AnchorSerialize::serialize(data, &mut buf).unwrap();
+    buf
+}
+
+fn base_input() -> MultiInputTokenDataWithContext {
+    MultiInputTokenDataWithContext::default()
+}
+
+fn base_data() -> CompressedTokenInstructionDataTransfer2 {
+    CompressedTokenInstructionDataTransfer2 {
+        with_transaction_hash: false,
+        with_lamports_change_account_merkle_tree_index: false,
+        lamports_change_account_merkle_tree_index: 0,
+        lamports_change_account_owner_index: 0,
+        output_queue: 0,
+        max_top_up: 0,
+        cpi_context: None,
+        compressions: None,
+        proof: None,
+        in_token_data: vec![],
+        out_token_data: vec![],
+        in_lamports: None,
+        out_lamports: None,
+        in_tlv: None,
+        out_tlv: None,
+    }
+}
+
+fn check(data: &CompressedTokenInstructionDataTransfer2) -> bool {
+    let buf = serialize(data);
+    let (z, _) = CompressedTokenInstructionDataTransfer2::zero_copy_at(&buf).unwrap();
+    is_idempotent_ata_decompress(&z)
+}
+
+#[test]
+fn test_is_idempotent_ata_decompress_empty() {
+    assert!(!check(&base_data()));
+}
+
+#[test]
+fn test_is_idempotent_ata_decompress_no_compressions() {
+    let mut data = base_data();
+    data.in_token_data = vec![base_input()];
+    data.in_tlv = Some(vec![vec![ExtensionInstructionData::CompressedOnly(
+        CompressedOnlyExtensionInstructionData {
+            delegated_amount: 0,
+            withheld_transfer_fee: 0,
+            is_frozen: false,
+            compression_index: 0,
+            is_ata: true,
+            bump: 0,
+            owner_index: 0,
+        },
+    )]]);
+    assert!(!check(&data));
+}
+
+#[test]
+fn test_is_idempotent_ata_decompress_multiple_inputs() {
+    let mut data = base_data();
+    data.in_token_data = vec![base_input(), base_input()];
+    data.compressions = Some(vec![Compression::decompress(100, 0, 0)]);
+    assert!(!check(&data));
+}
+
+#[test]
+fn test_is_idempotent_ata_decompress_compress_mode() {
+    let mut data = base_data();
+    data.in_token_data = vec![base_input()];
+    data.compressions = Some(vec![Compression::compress(100, 0, 0, 0)]);
+    data.in_tlv = Some(vec![vec![ExtensionInstructionData::CompressedOnly(
+        CompressedOnlyExtensionInstructionData {
+            delegated_amount: 0,
+            withheld_transfer_fee: 0,
+            is_frozen: false,
+            compression_index: 0,
+            is_ata: true,
+            bump: 0,
+            owner_index: 0,
+        },
+    )]]);
+    assert!(!check(&data));
+}
+
+#[test]
+fn test_is_idempotent_ata_decompress_not_ata() {
+    let mut data = base_data();
+    data.in_token_data = vec![base_input()];
+    data.compressions = Some(vec![Compression::decompress(100, 0, 0)]);
+    data.in_tlv = Some(vec![vec![ExtensionInstructionData::CompressedOnly(
+        CompressedOnlyExtensionInstructionData {
+            delegated_amount: 0,
+            withheld_transfer_fee: 0,
+            is_frozen: false,
+            compression_index: 0,
+            is_ata: false,
+            bump: 0,
+            owner_index: 0,
+        },
+    )]]);
+    assert!(!check(&data));
+}
+
+#[test]
+fn test_is_idempotent_ata_decompress_random_always_false() {
+    let mut rng = StdRng::seed_from_u64(42);
+    for _ in 0..1000 {
+        let mut data = base_data();
+        let num_inputs = rng.gen_range(0..5);
+        data.in_token_data = (0..num_inputs).map(|_| base_input()).collect();
+
+        // Random compressions -- never Decompress mode so result is always false.
+        if rng.gen_bool(0.5) {
+            let num_compressions = rng.gen_range(0..4);
+            data.compressions = Some(
+                (0..num_compressions)
+                    .map(|_| {
+                        if rng.gen_bool(0.5) {
+                            Compression::compress(rng.gen(), 0, 0, 0)
+                        } else {
+                            Compression {
+                                mode: CompressionMode::CompressAndClose,
+                                amount: rng.gen(),
+                                mint: 0,
+                                source_or_recipient: 0,
+                                authority: 0,
+                                pool_account_index: 0,
+                                pool_index: 0,
+                                bump: 0,
+                                decimals: 0,
+                            }
+                        }
+                    })
+                    .collect(),
+            );
+        }
+
+        assert!(
+            !check(&data),
+            "Expected false for random input with {num_inputs} inputs"
+        );
+    }
+}

sdk-libs/client/src/interface/load_accounts.rs

@@ -364,7 +364,7 @@ fn build_transfer2(
             },
         )?;
 
-        let owner_idx = packed.insert_or_get(ctx.wallet_owner);
+        let owner_idx = packed.insert_or_get_read_only(ctx.wallet_owner);
         let ata_idx = packed.insert_or_get(derive_token_ata(&ctx.wallet_owner, &ctx.mint));
         let mint_idx = packed.insert_or_get(token.mint);
         let delegate_idx = token.delegate.map(|d| packed.insert_or_get(d)).unwrap_or(0);