light_token_interface/instructions/transfer2/
compression.rs

1use std::fmt::Debug;
2
3use light_zero_copy::{
4    errors::ZeroCopyError, traits::ZeroCopyAtMut, ZeroCopy, ZeroCopyMut, ZeroCopyNew,
5};
6use zerocopy::Ref;
7
8use crate::{AnchorDeserialize, AnchorSerialize, TokenError};
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq, AnchorSerialize, AnchorDeserialize, ZeroCopy)]
11#[repr(C)]
12pub enum CompressionMode {
13    Compress,
14    Decompress,
15    /// Compresses token account and closes it
16    /// Signer must be rent authority, token account must be compressible
17    /// Not implemented for spl token accounts.
18    CompressAndClose,
19}
20
21impl ZCompressionMode {
22    pub fn is_compress(&self) -> bool {
23        matches!(self, ZCompressionMode::Compress)
24    }
25
26    pub fn is_decompress(&self) -> bool {
27        matches!(self, ZCompressionMode::Decompress)
28    }
29
30    pub fn is_compress_and_close(&self) -> bool {
31        matches!(self, ZCompressionMode::CompressAndClose)
32    }
33}
34
35pub const COMPRESS: u8 = 0u8;
36pub const DECOMPRESS: u8 = 1u8;
37pub const COMPRESS_AND_CLOSE: u8 = 2u8;
38
39impl<'a> ZeroCopyAtMut<'a> for CompressionMode {
40    type ZeroCopyAtMut = Ref<&'a mut [u8], u8>;
41    fn zero_copy_at_mut(
42        bytes: &'a mut [u8],
43    ) -> Result<(Self::ZeroCopyAtMut, &'a mut [u8]), ZeroCopyError> {
44        let (mode, bytes) = zerocopy::Ref::<&mut [u8], u8>::from_prefix(bytes)?;
45
46        Ok((mode, bytes))
47    }
48}
49
50impl<'a> ZeroCopyNew<'a> for CompressionMode {
51    type ZeroCopyConfig = ();
52    type Output = Ref<&'a mut [u8], u8>;
53
54    fn byte_len(_config: &Self::ZeroCopyConfig) -> Result<usize, ZeroCopyError> {
55        Ok(1) // CompressionMode enum size is always 1 byte
56    }
57
58    fn new_zero_copy(
59        bytes: &'a mut [u8],
60        _config: Self::ZeroCopyConfig,
61    ) -> Result<(Self::Output, &'a mut [u8]), ZeroCopyError> {
62        let (mode, remaining_bytes) = zerocopy::Ref::<&mut [u8], u8>::from_prefix(bytes)?;
63
64        Ok((mode, remaining_bytes))
65    }
66}
67
68#[repr(C)]
69#[derive(
70    Clone, Copy, Debug, PartialEq, Eq, AnchorSerialize, AnchorDeserialize, ZeroCopy, ZeroCopyMut,
71)]
72pub struct Compression {
73    pub mode: CompressionMode,
74    pub amount: u64,
75    pub mint: u8,
76    pub source_or_recipient: u8,
77    pub authority: u8, // Index of owner or delegate account
78    /// pool account index for spl token Compression/Decompression
79    /// rent_sponsor_index for CompressAndClose
80    pub pool_account_index: u8, // This account is not necessary to decompress ctokens because there are no token pools
81    /// pool index for spl token Compression/Decompression
82    /// compressed account index for CompressAndClose
83    pub pool_index: u8, // This account is not necessary to decompress ctokens because there are no token pools
84    pub bump: u8, // This account is not necessary to decompress ctokens because there are no token pools
85    /// decimals for spl token Compression/Decompression (used in transfer_checked)
86    /// rent_sponsor_is_signer flag for CompressAndClose (non-zero = true)
87    pub decimals: u8,
88}
89
90impl ZCompression<'_> {
91    pub fn get_rent_sponsor_index(&self) -> Result<u8, TokenError> {
92        match self.mode {
93            ZCompressionMode::CompressAndClose => Ok(self.pool_account_index),
94            _ => Err(TokenError::InvalidCompressionMode),
95        }
96    }
97    pub fn get_compressed_token_account_index(&self) -> Result<u8, TokenError> {
98        match self.mode {
99            ZCompressionMode::CompressAndClose => Ok(self.pool_index),
100            _ => Err(TokenError::InvalidCompressionMode),
101        }
102    }
103    pub fn get_destination_index(&self) -> Result<u8, TokenError> {
104        match self.mode {
105            ZCompressionMode::CompressAndClose => Ok(self.bump),
106            _ => Err(TokenError::InvalidCompressionMode),
107        }
108    }
109}
110
111impl Compression {
112    #[allow(clippy::too_many_arguments)]
113    pub fn compress_and_close(
114        amount: u64,
115        mint: u8,
116        source: u8,
117        authority: u8,
118        rent_sponsor_index: u8,
119        compressed_account_index: u8,
120        destination_index: u8,
121    ) -> Self {
122        Compression {
123            amount, // the full balance of the token account to be compressed
124            mode: CompressionMode::CompressAndClose,
125            mint,
126            source_or_recipient: source,
127            authority,
128            pool_account_index: rent_sponsor_index,
129            pool_index: compressed_account_index,
130            bump: destination_index,
131            decimals: 0,
132        }
133    }
134
135    #[allow(clippy::too_many_arguments)]
136    pub fn compress_spl(
137        amount: u64,
138        mint: u8,
139        source: u8,
140        authority: u8,
141        pool_account_index: u8,
142        pool_index: u8,
143        bump: u8,
144        decimals: u8,
145    ) -> Self {
146        Compression {
147            amount,
148            mode: CompressionMode::Compress,
149            mint,
150            source_or_recipient: source,
151            authority,
152            pool_account_index,
153            pool_index,
154            bump,
155            decimals,
156        }
157    }
158    pub fn compress(amount: u64, mint: u8, source: u8, authority: u8) -> Self {
159        Compression {
160            amount,
161            mode: CompressionMode::Compress,
162            mint,
163            source_or_recipient: source,
164            authority,
165            pool_account_index: 0,
166            pool_index: 0,
167            bump: 0,
168            decimals: 0,
169        }
170    }
171
172    pub fn decompress_spl(
173        amount: u64,
174        mint: u8,
175        recipient: u8,
176        pool_account_index: u8,
177        pool_index: u8,
178        bump: u8,
179        decimals: u8,
180    ) -> Self {
181        Compression {
182            amount,
183            mode: CompressionMode::Decompress,
184            mint,
185            source_or_recipient: recipient,
186            authority: 0,
187            pool_account_index,
188            pool_index,
189            bump,
190            decimals,
191        }
192    }
193
194    pub fn decompress(amount: u64, mint: u8, recipient: u8) -> Self {
195        Compression {
196            amount,
197            mode: CompressionMode::Decompress,
198            mint,
199            source_or_recipient: recipient,
200            authority: 0,
201            pool_account_index: 0,
202            pool_index: 0,
203            bump: 0,
204            decimals: 0,
205        }
206    }
207}
208
209impl ZCompressionMut<'_> {
210    pub fn mode(&self) -> Result<CompressionMode, TokenError> {
211        match *self.mode {
212            COMPRESS => Ok(CompressionMode::Compress),
213            DECOMPRESS => Ok(CompressionMode::Decompress),
214            COMPRESS_AND_CLOSE => Ok(CompressionMode::CompressAndClose),
215            _ => Err(TokenError::InvalidCompressionMode),
216        }
217    }
218}
219
220impl ZCompression<'_> {
221    pub fn new_balance_compressed_account(&self, current_balance: u64) -> Result<u64, TokenError> {
222        let new_balance = match self.mode {
223            ZCompressionMode::Compress | ZCompressionMode::CompressAndClose => {
224                // Compress: add to balance (tokens are being added to spl token pool)
225                current_balance
226                    .checked_add((*self.amount).into())
227                    .ok_or(TokenError::ArithmeticOverflow)
228            }
229            ZCompressionMode::Decompress => {
230                // Decompress: subtract from balance (tokens are being removed from spl token pool)
231                current_balance
232                    .checked_sub((*self.amount).into())
233                    .ok_or(TokenError::CompressInsufficientFunds)
234            }
235        }?;
236        Ok(new_balance)
237    }
238}