light_program_test/
compressible.rs

1#[cfg(feature = "devenv")]
2use std::collections::HashMap;
3
4#[cfg(feature = "devenv")]
5use anchor_lang::pubkey;
6#[cfg(feature = "devenv")]
7use borsh::BorshDeserialize;
8#[cfg(feature = "devenv")]
9use light_client::rpc::{Rpc, RpcError};
10#[cfg(feature = "devenv")]
11use light_compressible::rent::SLOTS_PER_EPOCH;
12#[cfg(feature = "devenv")]
13use light_compressible::{config::CompressibleConfig, rent::RentConfig};
14#[cfg(feature = "devenv")]
15use light_ctoken_types::{
16    state::{CToken, ExtensionStruct},
17    COMPRESSIBLE_TOKEN_ACCOUNT_SIZE,
18};
19#[cfg(feature = "devenv")]
20use solana_pubkey::Pubkey;
21
22#[cfg(feature = "devenv")]
23use crate::LightProgramTest;
24
25#[cfg(feature = "devenv")]
26pub type CompressibleAccountStore = HashMap<Pubkey, StoredCompressibleAccount>;
27
28#[cfg(feature = "devenv")]
29#[derive(Eq, Hash, PartialEq)]
30pub struct StoredCompressibleAccount {
31    pub pubkey: Pubkey,
32    pub last_paid_slot: u64,
33    pub account: CToken,
34}
35
36#[cfg(feature = "devenv")]
37#[derive(Debug, PartialEq, Copy, Clone)]
38pub struct FundingPoolConfig {
39    pub compressible_config_pda: Pubkey,
40    pub compression_authority_pda: Pubkey,
41    pub compression_authority_pda_bump: u8,
42    /// rent_sponsor == pool pda
43    pub rent_sponsor_pda: Pubkey,
44    pub rent_sponsor_pda_bump: u8,
45}
46
47#[cfg(feature = "devenv")]
48impl FundingPoolConfig {
49    pub fn new(version: u16) -> Self {
50        let config = CompressibleConfig::new_ctoken(
51            version,
52            true,
53            Pubkey::default(),
54            Pubkey::default(),
55            RentConfig::default(),
56        );
57        let compressible_config = CompressibleConfig::derive_pda(
58            &pubkey!("Lighton6oQpVkeewmo2mcPTQQp7kYHr4fWpAgJyEmDX"),
59            version,
60        )
61        .0;
62        Self {
63            compressible_config_pda: compressible_config,
64            rent_sponsor_pda: config.rent_sponsor,
65            rent_sponsor_pda_bump: config.rent_sponsor_bump,
66            compression_authority_pda: config.compression_authority,
67            compression_authority_pda_bump: config.compression_authority_bump,
68        }
69    }
70
71    pub fn get_v1() -> Self {
72        Self::new(1)
73    }
74}
75
76#[cfg(feature = "devenv")]
77pub async fn claim_and_compress(
78    rpc: &mut LightProgramTest,
79    stored_compressible_accounts: &mut CompressibleAccountStore,
80) -> Result<(), RpcError> {
81    use crate::forester::{claim_forester, compress_and_close_forester};
82
83    let forester_keypair = rpc.test_accounts.protocol.forester.insecure_clone();
84    let payer = rpc.get_payer().insecure_clone();
85
86    // Get all compressible token accounts
87    let compressible_ctoken_accounts = rpc
88        .context
89        .get_program_accounts(&light_compressed_token::ID);
90
91    for account in compressible_ctoken_accounts
92        .iter()
93        .filter(|e| e.1.data.len() > 200 && e.1.lamports > 0)
94    {
95        let des_account = CToken::deserialize(&mut account.1.data.as_slice())?;
96        if let Some(extensions) = des_account.extensions.as_ref() {
97            for extension in extensions.iter() {
98                if let ExtensionStruct::Compressible(e) = extension {
99                    let base_lamports = rpc
100                        .get_minimum_balance_for_rent_exemption(
101                            COMPRESSIBLE_TOKEN_ACCOUNT_SIZE as usize,
102                        )
103                        .await
104                        .unwrap();
105                    let last_funded_epoch = e
106                        .get_last_funded_epoch(
107                            account.1.data.len() as u64,
108                            account.1.lamports,
109                            base_lamports,
110                        )
111                        .unwrap();
112                    let last_funded_slot = last_funded_epoch * SLOTS_PER_EPOCH;
113                    stored_compressible_accounts.insert(
114                        account.0,
115                        StoredCompressibleAccount {
116                            pubkey: account.0,
117                            last_paid_slot: last_funded_slot,
118                            account: des_account.clone(),
119                        },
120                    );
121                }
122            }
123        }
124    }
125
126    let current_slot = rpc.get_slot().await?;
127    let compressible_accounts = {
128        stored_compressible_accounts
129            .iter()
130            .filter(|a| a.1.last_paid_slot < current_slot)
131            .map(|e| e.1)
132            .collect::<Vec<_>>()
133    };
134
135    let claim_able_accounts = stored_compressible_accounts
136        .iter()
137        .filter(|a| a.1.last_paid_slot >= current_slot)
138        .map(|e| *e.0)
139        .collect::<Vec<_>>();
140
141    // Process claimable accounts in batches
142    for token_accounts in claim_able_accounts.as_slice().chunks(20) {
143        println!("Claim from : {:?}", token_accounts);
144        // Use the new claim_forester function to claim via registry program
145        claim_forester(rpc, token_accounts, &forester_keypair, &payer).await?;
146    }
147
148    // Process compressible accounts in batches
149    const BATCH_SIZE: usize = 10; // Process up to 10 accounts at a time
150    let mut pubkeys = Vec::with_capacity(compressible_accounts.len());
151    for chunk in compressible_accounts.chunks(BATCH_SIZE) {
152        let chunk_pubkeys: Vec<Pubkey> = chunk.iter().map(|e| e.pubkey).collect();
153        println!("Compress and close: {:?}", chunk_pubkeys);
154
155        // Use the new compress_and_close_forester function via registry program
156        compress_and_close_forester(rpc, &chunk_pubkeys, &forester_keypair, &payer, None).await?;
157
158        // Remove processed accounts from the HashMap
159        for account_pubkey in chunk {
160            pubkeys.push(account_pubkey.pubkey);
161        }
162    }
163    for pubkey in pubkeys {
164        stored_compressible_accounts.remove(&pubkey);
165    }
166
167    Ok(())
168}