light_token/compressed_token/v2/
account2.rs

1use std::ops::Deref;
2
3use light_program_profiler::profile;
4use light_token_interface::instructions::transfer2::{
5    Compression, CompressionMode, MultiInputTokenDataWithContext, MultiTokenTransferOutputData,
6};
7use solana_account_info::AccountInfo;
8use solana_pubkey::Pubkey;
9
10use crate::{error::TokenSdkError, utils::get_token_account_balance};
11
12#[derive(Debug, PartialEq, Clone)]
13pub struct CTokenAccount2 {
14    pub inputs: Vec<MultiInputTokenDataWithContext>,
15    pub output: MultiTokenTransferOutputData,
16    pub compression: Option<Compression>,
17    pub delegate_is_set: bool,
18    pub method_used: bool,
19}
20
21impl CTokenAccount2 {
22    #[profile]
23    pub fn new(token_data: Vec<MultiInputTokenDataWithContext>) -> Result<Self, TokenSdkError> {
24        // all mint indices must be the same
25        // all owners must be the same
26        let amount = token_data.iter().map(|data| data.amount).sum();
27        // Check if token_data is empty
28        if token_data.is_empty() {
29            return Err(TokenSdkError::NoInputAccounts);
30        }
31
32        // Use the indices from the first token data (assuming they're all the same mint/owner)
33        let mint_index = token_data[0].mint;
34        let owner_index = token_data[0].owner;
35        let version = token_data[0].version; // Take version from input
36        let output = MultiTokenTransferOutputData {
37            owner: owner_index,
38            amount,
39            delegate: 0, // Default delegate index
40            mint: mint_index,
41            version, // Use version from input accounts
42            has_delegate: false,
43        };
44        Ok(Self {
45            inputs: token_data,
46            output,
47            delegate_is_set: false,
48            compression: None,
49            method_used: false,
50        })
51    }
52
53    /// Input token accounts are delegated and delegate is signer
54    /// The change output account is also delegated.
55    /// (with new change output account is not delegated even if inputs were)
56    #[profile]
57    pub fn new_delegated(
58        token_data: Vec<MultiInputTokenDataWithContext>,
59    ) -> Result<Self, TokenSdkError> {
60        // all mint indices must be the same
61        // all owners must be the same
62        let amount = token_data.iter().map(|data| data.amount).sum();
63        // Check if token_data is empty
64        if token_data.is_empty() {
65            return Err(TokenSdkError::NoInputAccounts);
66        }
67
68        // Use the indices from the first token data (assuming they're all the same mint/owner)
69        let mint_index = token_data[0].mint;
70        let owner_index = token_data[0].owner;
71        let version = token_data[0].version; // Take version from input
72        let output = MultiTokenTransferOutputData {
73            owner: owner_index,
74            amount,
75            delegate: token_data[0].delegate, // Default delegate index
76            mint: mint_index,
77            version, // Use version from input accounts
78            has_delegate: true,
79        };
80        Ok(Self {
81            inputs: token_data,
82            output,
83            delegate_is_set: false,
84            compression: None,
85            method_used: false,
86        })
87    }
88
89    #[profile]
90    pub fn new_empty(owner_index: u8, mint_index: u8) -> Self {
91        Self {
92            inputs: vec![],
93            output: MultiTokenTransferOutputData {
94                owner: owner_index,
95                amount: 0,
96                delegate: 0, // Default delegate index
97                mint: mint_index,
98                version: 3, // ShaFlat
99                has_delegate: false,
100            },
101            compression: None,
102            delegate_is_set: false,
103            method_used: false,
104        }
105    }
106
107    // TODO: consider this might be confusing because it must not be used in combination with fn transfer()
108    //     could mark the struct as transferred and throw in fn transfer
109    #[profile]
110    pub fn transfer(&mut self, recipient_index: u8, amount: u64) -> Result<Self, TokenSdkError> {
111        if amount > self.output.amount {
112            return Err(TokenSdkError::InsufficientBalance);
113        }
114        // TODO: skip outputs with zero amount when creating the instruction data.
115        self.output.amount -= amount;
116
117        self.method_used = true;
118        Ok(Self {
119            compression: None,
120            inputs: vec![],
121            output: MultiTokenTransferOutputData {
122                owner: recipient_index,
123                amount,
124                delegate: 0,
125                mint: self.output.mint,
126                version: self.output.version,
127                has_delegate: false,
128            },
129            delegate_is_set: false,
130            method_used: false,
131        })
132    }
133
134    /// Approves a delegate for a specified amount of tokens.
135    /// Similar to transfer, this deducts the amount from the current account
136    /// and returns a new CTokenAccount that represents the delegated portion.
137    /// The original account balance is reduced by the delegated amount.
138    #[profile]
139    pub fn approve(&mut self, delegate_index: u8, amount: u64) -> Result<Self, TokenSdkError> {
140        if amount > self.output.amount {
141            return Err(TokenSdkError::InsufficientBalance);
142        }
143
144        // Deduct the delegated amount from current account
145        self.output.amount -= amount;
146
147        self.method_used = true;
148
149        // Create a new delegated account with the specified delegate
150        // Note: In the actual instruction, this will create the proper delegation structure
151        Ok(Self {
152            compression: None,
153            inputs: vec![],
154            output: MultiTokenTransferOutputData {
155                owner: self.output.owner, // Owner remains the same
156                amount,
157                delegate: delegate_index,
158                mint: self.output.mint,
159                version: self.output.version,
160                has_delegate: true,
161            },
162            delegate_is_set: true,
163            method_used: false,
164        })
165    }
166
167    // TODO: consider this might be confusing because it must not be used in combination with fn compress()
168    #[profile]
169    pub fn compress(
170        &mut self,
171        amount: u64,
172        source_or_recipient_index: u8,
173        authority: u8,
174    ) -> Result<(), TokenSdkError> {
175        // Check if there's already a compression set
176        if self.compression.is_some() {
177            return Err(TokenSdkError::CompressionCannotBeSetTwice);
178        }
179
180        self.output.amount += amount;
181        self.compression = Some(Compression::compress(
182            amount,
183            self.output.mint,
184            source_or_recipient_index,
185            authority,
186        ));
187        self.method_used = true;
188
189        Ok(())
190    }
191
192    #[profile]
193    #[allow(clippy::too_many_arguments)]
194    pub fn compress_spl(
195        &mut self,
196        amount: u64,
197        source_or_recipient_index: u8,
198        authority: u8,
199        pool_account_index: u8,
200        pool_index: u8,
201        bump: u8,
202        decimals: u8,
203    ) -> Result<(), TokenSdkError> {
204        // Check if there's already a compression set
205        if self.compression.is_some() {
206            return Err(TokenSdkError::CompressionCannotBeSetTwice);
207        }
208
209        self.output.amount += amount;
210        self.compression = Some(Compression::compress_spl(
211            amount,
212            self.output.mint,
213            source_or_recipient_index,
214            authority,
215            pool_account_index,
216            pool_index,
217            bump,
218            decimals,
219        ));
220        self.method_used = true;
221
222        Ok(())
223    }
224
225    // TODO: consider this might be confusing because it must not be used in combination with fn decompress()
226    #[profile]
227    pub fn decompress(&mut self, amount: u64, source_index: u8) -> Result<(), TokenSdkError> {
228        // Check if there's already a compression set
229        if self.compression.is_some() {
230            return Err(TokenSdkError::CompressionCannotBeSetTwice);
231        }
232
233        if self.output.amount < amount {
234            return Err(TokenSdkError::InsufficientBalance);
235        }
236        self.output.amount -= amount;
237
238        self.compression = Some(Compression::decompress(
239            amount,
240            self.output.mint,
241            source_index,
242        ));
243        self.method_used = true;
244
245        Ok(())
246    }
247
248    #[profile]
249    pub fn decompress_spl(
250        &mut self,
251        amount: u64,
252        source_index: u8,
253        pool_account_index: u8,
254        pool_index: u8,
255        bump: u8,
256        decimals: u8,
257    ) -> Result<(), TokenSdkError> {
258        // Check if there's already a compression set
259        if self.compression.is_some() {
260            return Err(TokenSdkError::CompressionCannotBeSetTwice);
261        }
262
263        if self.output.amount < amount {
264            return Err(TokenSdkError::InsufficientBalance);
265        }
266        self.output.amount -= amount;
267
268        self.compression = Some(Compression::decompress_spl(
269            amount,
270            self.output.mint,
271            source_index,
272            pool_account_index,
273            pool_index,
274            bump,
275            decimals,
276        ));
277        self.method_used = true;
278
279        Ok(())
280    }
281
282    #[profile]
283    pub fn compress_full(
284        &mut self,
285        source_or_recipient_index: u8,
286        authority: u8,
287        token_account_info: &AccountInfo,
288    ) -> Result<(), TokenSdkError> {
289        // Check if there's already a compression set
290        if self.compression.is_some() {
291            return Err(TokenSdkError::CompressionCannotBeSetTwice);
292        }
293
294        // Get the actual token account balance to add to output
295        let token_balance = get_token_account_balance(token_account_info)?;
296
297        // Add the full token balance to the output amount
298        self.output.amount += token_balance;
299
300        // For compress_full, set amount to the actual balance for instruction data
301        self.compression = Some(Compression {
302            amount: token_balance,
303            mode: CompressionMode::Compress, // Use regular compress mode with actual amount
304            mint: self.output.mint,
305            source_or_recipient: source_or_recipient_index,
306            authority,
307            pool_account_index: 0,
308            pool_index: 0,
309            bump: 0,
310            decimals: 0, // Not used for ctoken compression
311        });
312        self.method_used = true;
313
314        Ok(())
315    }
316
317    #[profile]
318    pub fn compress_and_close(
319        &mut self,
320        amount: u64,
321        source_or_recipient_index: u8,
322        authority: u8,
323        rent_sponsor_index: u8,
324        compressed_account_index: u8,
325        destination_index: u8,
326    ) -> Result<(), TokenSdkError> {
327        // Check if there's already a compression set
328        if self.compression.is_some() {
329            return Err(TokenSdkError::CompressionCannotBeSetTwice);
330        }
331
332        // Add the full balance to the output amount
333        self.output.amount += amount;
334
335        // Use the compress_and_close method from Compression
336        self.compression = Some(Compression::compress_and_close(
337            amount,
338            self.output.mint,
339            source_or_recipient_index,
340            authority,
341            rent_sponsor_index,
342            compressed_account_index,
343            destination_index,
344        ));
345        self.method_used = true;
346
347        Ok(())
348    }
349
350    pub fn is_compress(&self) -> bool {
351        self.compression
352            .as_ref()
353            .map(|c| c.mode == CompressionMode::Compress)
354            .unwrap_or(false)
355    }
356
357    pub fn is_decompress(&self) -> bool {
358        self.compression
359            .as_ref()
360            .map(|c| c.mode == CompressionMode::Decompress)
361            .unwrap_or(false)
362    }
363
364    pub fn mint(&self, account_infos: &[AccountInfo]) -> Pubkey {
365        *account_infos[self.mint as usize].key
366    }
367
368    pub fn compression_amount(&self) -> Option<u64> {
369        self.compression.as_ref().map(|c| c.amount)
370    }
371
372    pub fn compression(&self) -> Option<&Compression> {
373        self.compression.as_ref()
374    }
375
376    pub fn owner(&self, account_infos: &[AccountInfo]) -> Pubkey {
377        *account_infos[self.owner as usize].key
378    }
379    // TODO: make option and take from self
380    //pub fn delegate_account<'b>(&self, account_infos: &'b [&'b AccountInfo]) -> &'b Pubkey {
381    //    account_infos[self.output.delegate as usize].key
382    // }
383
384    pub fn input_metas(&self) -> &[MultiInputTokenDataWithContext] {
385        self.inputs.as_slice()
386    }
387
388    /// Consumes token account for instruction creation.
389    pub fn into_inputs_and_outputs(
390        self,
391    ) -> (
392        Vec<MultiInputTokenDataWithContext>,
393        MultiTokenTransferOutputData,
394    ) {
395        (self.inputs, self.output)
396    }
397}
398
399impl Deref for CTokenAccount2 {
400    type Target = MultiTokenTransferOutputData;
401
402    fn deref(&self) -> &Self::Target {
403        &self.output
404    }
405}