light_sdk/
compressed_account.rs

1use std::ops::{Deref, DerefMut};
2
3use anchor_lang::prelude::{
4    AccountInfo, AnchorDeserialize, AnchorSerialize, ProgramError, Pubkey, Result,
5};
6use light_hasher::{DataHasher, Discriminator, Hasher, Poseidon};
7use light_utils::hash_to_bn254_field_size_be;
8
9use crate::{
10    address::{derive_address, NewAddressParamsPacked},
11    merkle_context::{
12        pack_merkle_context, MerkleContext, PackedAddressMerkleContext, PackedMerkleContext,
13        RemainingAccounts,
14    },
15    program_merkle_context::unpack_address_merkle_context,
16};
17
18pub trait LightAccounts: Sized {
19    fn try_light_accounts(
20        inputs: Vec<Vec<u8>>,
21        merkle_context: PackedMerkleContext,
22        merkle_tree_root_index: u16,
23        address_merkle_context: PackedAddressMerkleContext,
24        address_merkle_tree_root_index: u16,
25        remaining_accounts: &[AccountInfo],
26    ) -> Result<Self>;
27    fn new_address_params(&self) -> Vec<NewAddressParamsPacked>;
28    fn input_accounts(
29        &self,
30        remaining_accounts: &[AccountInfo],
31    ) -> Result<Vec<PackedCompressedAccountWithMerkleContext>>;
32    fn output_accounts(
33        &self,
34        remaining_accounts: &[AccountInfo],
35    ) -> Result<Vec<OutputCompressedAccountWithPackedContext>>;
36}
37
38/// A wrapper which abstracts away the UTXO model.
39pub enum LightAccount<T>
40where
41    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
42{
43    Init(LightInitAccount<T>),
44    Mut(LightMutAccount<T>),
45    Close(LightCloseAccount<T>),
46}
47
48impl<T> LightAccount<T>
49where
50    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Default + Discriminator,
51{
52    pub fn new_init(
53        merkle_context: &PackedMerkleContext,
54        address_merkle_context: &PackedAddressMerkleContext,
55        address_merkle_tree_root_index: u16,
56    ) -> Self {
57        Self::Init(LightInitAccount::new(
58            merkle_context,
59            address_merkle_context,
60            address_merkle_tree_root_index,
61        ))
62    }
63
64    pub fn try_from_slice_mut(
65        v: &[u8],
66        merkle_context: &PackedMerkleContext,
67        merkle_tree_root_index: u16,
68        address_merkle_context: &PackedAddressMerkleContext,
69    ) -> Result<Self> {
70        Ok(Self::Mut(LightMutAccount::try_from_slice(
71            v,
72            merkle_context,
73            merkle_tree_root_index,
74            address_merkle_context,
75        )?))
76    }
77
78    pub fn try_from_slice_close(
79        v: &[u8],
80        merkle_context: &PackedMerkleContext,
81        merkle_tree_root_index: u16,
82        address_merkle_context: &PackedAddressMerkleContext,
83    ) -> Result<Self> {
84        Ok(Self::Close(LightCloseAccount::try_from_slice(
85            v,
86            merkle_context,
87            merkle_tree_root_index,
88            address_merkle_context,
89        )?))
90    }
91
92    pub fn set_address_seed(&mut self, address_seed: [u8; 32]) {
93        match self {
94            Self::Init(light_init_account) => light_init_account.set_address_seed(address_seed),
95            Self::Mut(light_mut_account) => light_mut_account.set_address_seed(address_seed),
96            Self::Close(light_close_account) => light_close_account.set_address_seed(address_seed),
97        }
98    }
99
100    pub fn new_address_params(&self) -> Option<NewAddressParamsPacked> {
101        match self {
102            Self::Init(self_init) => Some(self_init.new_address_params()),
103            Self::Mut(_) => None,
104            Self::Close(_) => None,
105        }
106    }
107
108    pub fn input_compressed_account(
109        &self,
110        program_id: &Pubkey,
111        remaining_accounts: &[AccountInfo],
112    ) -> Result<Option<PackedCompressedAccountWithMerkleContext>> {
113        match self {
114            Self::Init(_) => Ok(None),
115            Self::Mut(light_mut_account) => {
116                let account =
117                    light_mut_account.input_compressed_account(program_id, remaining_accounts)?;
118                Ok(Some(account))
119            }
120            Self::Close(light_close_account) => {
121                let account =
122                    light_close_account.input_compressed_account(program_id, remaining_accounts)?;
123                Ok(Some(account))
124            }
125        }
126    }
127
128    pub fn output_compressed_account(
129        &self,
130        program_id: &Pubkey,
131        remaining_accounts: &[AccountInfo],
132    ) -> Result<Option<OutputCompressedAccountWithPackedContext>> {
133        match self {
134            Self::Init(light_init_account) => {
135                let account =
136                    light_init_account.output_compressed_account(program_id, remaining_accounts)?;
137                Ok(Some(account))
138            }
139            Self::Mut(light_mut_account) => {
140                let account =
141                    light_mut_account.output_compressed_account(program_id, remaining_accounts)?;
142                Ok(Some(account))
143            }
144            Self::Close(_) => Ok(None),
145        }
146    }
147}
148
149impl<T> Deref for LightAccount<T>
150where
151    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
152{
153    type Target = T;
154
155    fn deref(&self) -> &Self::Target {
156        match self {
157            Self::Init(light_init_account) => &light_init_account.output_account,
158            Self::Mut(light_mut_account) => &light_mut_account.output_account,
159            Self::Close(light_close_account) => &light_close_account.input_account,
160        }
161    }
162}
163
164impl<T> DerefMut for LightAccount<T>
165where
166    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
167{
168    fn deref_mut(&mut self) -> &mut Self::Target {
169        match self {
170            Self::Init(light_init_account) => &mut light_init_account.output_account,
171            Self::Mut(light_mut_account) => &mut light_mut_account.output_account,
172            Self::Close(light_close_account) => &mut light_close_account.input_account,
173        }
174    }
175}
176
177pub struct LightInitAccount<T>
178where
179    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
180{
181    output_account: T,
182    address_seed: Option<[u8; 32]>,
183    merkle_context: PackedMerkleContext,
184    address_merkle_context: PackedAddressMerkleContext,
185    address_merkle_tree_root_index: u16,
186}
187
188impl<T> LightInitAccount<T>
189where
190    T: AnchorDeserialize + AnchorSerialize + Clone + Default + DataHasher + Discriminator,
191{
192    pub fn new(
193        merkle_context: &PackedMerkleContext,
194        address_merkle_context: &PackedAddressMerkleContext,
195        address_merkle_tree_root_index: u16,
196    ) -> Self {
197        let output_account = T::default();
198
199        Self {
200            output_account,
201            address_seed: None,
202            merkle_context: *merkle_context,
203            address_merkle_context: *address_merkle_context,
204            address_merkle_tree_root_index,
205        }
206    }
207
208    pub fn set_address_seed(&mut self, address_seed: [u8; 32]) {
209        self.address_seed = Some(address_seed);
210    }
211
212    pub fn new_address_params(&self) -> NewAddressParamsPacked {
213        NewAddressParamsPacked {
214            seed: self.address_seed.unwrap(),
215            address_merkle_tree_account_index: self
216                .address_merkle_context
217                .address_merkle_tree_pubkey_index,
218            address_queue_account_index: self.address_merkle_context.address_queue_pubkey_index,
219            address_merkle_tree_root_index: self.address_merkle_tree_root_index,
220        }
221    }
222
223    pub fn output_compressed_account(
224        &self,
225        program_id: &Pubkey,
226        remaining_accounts: &[AccountInfo],
227    ) -> Result<OutputCompressedAccountWithPackedContext> {
228        output_compressed_account(
229            &self.output_account,
230            &self.address_seed.unwrap(),
231            program_id,
232            &self.merkle_context,
233            &self.address_merkle_context,
234            remaining_accounts,
235        )
236    }
237}
238
239impl<T> Deref for LightInitAccount<T>
240where
241    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
242{
243    type Target = T;
244
245    fn deref(&self) -> &Self::Target {
246        &self.output_account
247    }
248}
249
250impl<T> DerefMut for LightInitAccount<T>
251where
252    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
253{
254    fn deref_mut(&mut self) -> &mut Self::Target {
255        &mut self.output_account
256    }
257}
258
259pub struct LightMutAccount<T>
260where
261    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
262{
263    input_account: T,
264    output_account: T,
265    address_seed: Option<[u8; 32]>,
266    merkle_context: PackedMerkleContext,
267    merkle_tree_root_index: u16,
268    address_merkle_context: PackedAddressMerkleContext,
269}
270
271impl<T> LightMutAccount<T>
272where
273    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
274{
275    pub fn try_from_slice(
276        v: &[u8],
277        merkle_context: &PackedMerkleContext,
278        merkle_tree_root_index: u16,
279        address_merkle_context: &PackedAddressMerkleContext,
280    ) -> Result<Self> {
281        let account = T::try_from_slice(v)?;
282
283        Ok(Self {
284            input_account: account.clone(),
285            output_account: account,
286            address_seed: None,
287            merkle_context: *merkle_context,
288            merkle_tree_root_index,
289            address_merkle_context: *address_merkle_context,
290        })
291    }
292
293    pub fn set_address_seed(&mut self, address_seed: [u8; 32]) {
294        self.address_seed = Some(address_seed);
295    }
296
297    pub fn input_compressed_account(
298        &self,
299        program_id: &Pubkey,
300        remaining_accounts: &[AccountInfo],
301    ) -> Result<PackedCompressedAccountWithMerkleContext> {
302        input_compressed_account(
303            &self.input_account,
304            &self.address_seed.unwrap(),
305            program_id,
306            &self.merkle_context,
307            self.merkle_tree_root_index,
308            &self.address_merkle_context,
309            remaining_accounts,
310        )
311    }
312
313    pub fn output_compressed_account(
314        &self,
315        program_id: &Pubkey,
316        remaining_accounts: &[AccountInfo],
317    ) -> Result<OutputCompressedAccountWithPackedContext> {
318        output_compressed_account(
319            &self.output_account,
320            &self.address_seed.unwrap(),
321            program_id,
322            &self.merkle_context,
323            &self.address_merkle_context,
324            remaining_accounts,
325        )
326    }
327}
328
329impl<T> Deref for LightMutAccount<T>
330where
331    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
332{
333    type Target = T;
334
335    fn deref(&self) -> &Self::Target {
336        &self.output_account
337    }
338}
339
340impl<T> DerefMut for LightMutAccount<T>
341where
342    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
343{
344    fn deref_mut(&mut self) -> &mut Self::Target {
345        &mut self.output_account
346    }
347}
348
349pub struct LightCloseAccount<T>
350where
351    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
352{
353    input_account: T,
354    address_seed: Option<[u8; 32]>,
355    merkle_context: PackedMerkleContext,
356    merkle_tree_root_index: u16,
357    address_merkle_context: PackedAddressMerkleContext,
358}
359
360impl<T> LightCloseAccount<T>
361where
362    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
363{
364    pub fn try_from_slice(
365        v: &[u8],
366        merkle_context: &PackedMerkleContext,
367        merkle_tree_root_index: u16,
368        address_merkle_context: &PackedAddressMerkleContext,
369    ) -> Result<Self> {
370        let input_account = T::try_from_slice(v)?;
371
372        Ok(Self {
373            input_account,
374            address_seed: None,
375            merkle_context: *merkle_context,
376            merkle_tree_root_index,
377            address_merkle_context: *address_merkle_context,
378        })
379    }
380
381    pub fn set_address_seed(&mut self, address_seed: [u8; 32]) {
382        self.address_seed = Some(address_seed);
383    }
384
385    pub fn input_compressed_account(
386        &self,
387        program_id: &Pubkey,
388        remaining_accounts: &[AccountInfo],
389    ) -> Result<PackedCompressedAccountWithMerkleContext> {
390        input_compressed_account(
391            &self.input_account,
392            &self.address_seed.unwrap(),
393            program_id,
394            &self.merkle_context,
395            self.merkle_tree_root_index,
396            &self.address_merkle_context,
397            remaining_accounts,
398        )
399    }
400}
401
402impl<T> Deref for LightCloseAccount<T>
403where
404    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
405{
406    type Target = T;
407
408    fn deref(&self) -> &Self::Target {
409        &self.input_account
410    }
411}
412
413impl<T> DerefMut for LightCloseAccount<T>
414where
415    T: AnchorDeserialize + AnchorSerialize + Clone + DataHasher + Discriminator,
416{
417    fn deref_mut(&mut self) -> &mut Self::Target {
418        &mut self.input_account
419    }
420}
421
422#[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
423pub struct CompressedAccount {
424    pub owner: Pubkey,
425    pub lamports: u64,
426    pub address: Option<[u8; 32]>,
427    pub data: Option<CompressedAccountData>,
428}
429
430/// Hashing scheme:
431/// H(owner || leaf_index || merkle_tree_pubkey || lamports || address || data.discriminator || data.data_hash)
432impl CompressedAccount {
433    pub fn hash_with_hashed_values<H: Hasher>(
434        &self,
435        &owner_hashed: &[u8; 32],
436        &merkle_tree_hashed: &[u8; 32],
437        leaf_index: &u32,
438    ) -> Result<[u8; 32]> {
439        let capacity = 3
440            + std::cmp::min(self.lamports, 1) as usize
441            + self.address.is_some() as usize
442            + self.data.is_some() as usize * 2;
443        let mut vec: Vec<&[u8]> = Vec::with_capacity(capacity);
444        vec.push(owner_hashed.as_slice());
445
446        // leaf index and merkle tree pubkey are used to make every compressed account hash unique
447        let leaf_index = leaf_index.to_le_bytes();
448        vec.push(leaf_index.as_slice());
449
450        vec.push(merkle_tree_hashed.as_slice());
451
452        // Lamports are only hashed if non-zero to safe CU
453        // For safety we prefix the lamports with 1 in 1 byte.
454        // Thus even if the discriminator has the same value as the lamports, the hash will be different.
455        let mut lamports_bytes = [1, 0, 0, 0, 0, 0, 0, 0, 0];
456        if self.lamports != 0 {
457            lamports_bytes[1..].copy_from_slice(&self.lamports.to_le_bytes());
458            vec.push(lamports_bytes.as_slice());
459        }
460
461        if self.address.is_some() {
462            vec.push(self.address.as_ref().unwrap().as_slice());
463        }
464
465        let mut discriminator_bytes = [2, 0, 0, 0, 0, 0, 0, 0, 0];
466        if let Some(data) = &self.data {
467            discriminator_bytes[1..].copy_from_slice(&data.discriminator);
468            vec.push(&discriminator_bytes);
469            vec.push(&data.data_hash);
470        }
471        let hash = H::hashv(&vec).map_err(ProgramError::from)?;
472        Ok(hash)
473    }
474
475    pub fn hash<H: Hasher>(
476        &self,
477        &merkle_tree_pubkey: &Pubkey,
478        leaf_index: &u32,
479    ) -> Result<[u8; 32]> {
480        self.hash_with_hashed_values::<H>(
481            &hash_to_bn254_field_size_be(&self.owner.to_bytes())
482                .unwrap()
483                .0,
484            &hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
485                .unwrap()
486                .0,
487            leaf_index,
488        )
489    }
490}
491
492#[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
493pub struct CompressedAccountData {
494    pub discriminator: [u8; 8],
495    pub data: Vec<u8>,
496    pub data_hash: [u8; 32],
497}
498
499#[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
500pub struct CompressedAccountWithMerkleContext {
501    pub compressed_account: CompressedAccount,
502    pub merkle_context: MerkleContext,
503}
504
505#[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
506pub struct PackedCompressedAccountWithMerkleContext {
507    pub compressed_account: CompressedAccount,
508    pub merkle_context: PackedMerkleContext,
509    /// Index of root used in inclusion validity proof.
510    pub root_index: u16,
511    /// Placeholder to mark accounts read-only unimplemented set to false.
512    pub read_only: bool,
513}
514
515impl CompressedAccountWithMerkleContext {
516    pub fn hash(&self) -> Result<[u8; 32]> {
517        self.compressed_account.hash::<Poseidon>(
518            &self.merkle_context.merkle_tree_pubkey,
519            &self.merkle_context.leaf_index,
520        )
521    }
522}
523
524#[derive(Debug, PartialEq, Default, Clone, AnchorDeserialize, AnchorSerialize)]
525pub struct OutputCompressedAccountWithPackedContext {
526    pub compressed_account: CompressedAccount,
527    pub merkle_tree_index: u8,
528}
529
530pub fn serialize_and_hash_account<T>(
531    account: &T,
532    address_seed: &[u8; 32],
533    program_id: &Pubkey,
534    address_merkle_context: &PackedAddressMerkleContext,
535    remaining_accounts: &[AccountInfo],
536) -> Result<CompressedAccount>
537where
538    T: AnchorSerialize + DataHasher + Discriminator,
539{
540    let data = account.try_to_vec()?;
541    let data_hash = account.hash::<Poseidon>().map_err(ProgramError::from)?;
542    let compressed_account_data = CompressedAccountData {
543        discriminator: T::discriminator(),
544        data,
545        data_hash,
546    };
547
548    let address_merkle_context =
549        unpack_address_merkle_context(*address_merkle_context, remaining_accounts);
550    let address = derive_address(address_seed, &address_merkle_context);
551
552    let compressed_account = CompressedAccount {
553        owner: *program_id,
554        lamports: 0,
555        address: Some(address),
556        data: Some(compressed_account_data),
557    };
558
559    Ok(compressed_account)
560}
561
562pub fn input_compressed_account<T>(
563    account: &T,
564    address_seed: &[u8; 32],
565    program_id: &Pubkey,
566    merkle_context: &PackedMerkleContext,
567    merkle_tree_root_index: u16,
568    address_merkle_context: &PackedAddressMerkleContext,
569    remaining_accounts: &[AccountInfo],
570) -> Result<PackedCompressedAccountWithMerkleContext>
571where
572    T: AnchorSerialize + DataHasher + Discriminator,
573{
574    let compressed_account = serialize_and_hash_account(
575        account,
576        address_seed,
577        program_id,
578        address_merkle_context,
579        remaining_accounts,
580    )?;
581
582    Ok(PackedCompressedAccountWithMerkleContext {
583        compressed_account,
584        merkle_context: *merkle_context,
585        root_index: merkle_tree_root_index,
586        read_only: false,
587    })
588}
589
590pub fn output_compressed_account<T>(
591    account: &T,
592    address_seed: &[u8; 32],
593    program_id: &Pubkey,
594    merkle_context: &PackedMerkleContext,
595    address_merkle_context: &PackedAddressMerkleContext,
596    remaining_accounts: &[AccountInfo],
597) -> Result<OutputCompressedAccountWithPackedContext>
598where
599    T: AnchorSerialize + DataHasher + Discriminator,
600{
601    let compressed_account = serialize_and_hash_account(
602        account,
603        address_seed,
604        program_id,
605        address_merkle_context,
606        remaining_accounts,
607    )?;
608
609    Ok(OutputCompressedAccountWithPackedContext {
610        compressed_account,
611        merkle_tree_index: merkle_context.merkle_tree_pubkey_index,
612    })
613}
614
615pub fn pack_compressed_accounts(
616    compressed_accounts: &[CompressedAccountWithMerkleContext],
617    root_indices: &[u16],
618    remaining_accounts: &mut RemainingAccounts,
619) -> Vec<PackedCompressedAccountWithMerkleContext> {
620    compressed_accounts
621        .iter()
622        .zip(root_indices.iter())
623        .map(|(x, root_index)| PackedCompressedAccountWithMerkleContext {
624            compressed_account: x.compressed_account.clone(),
625            merkle_context: pack_merkle_context(x.merkle_context, remaining_accounts),
626            root_index: *root_index,
627            read_only: false,
628        })
629        .collect::<Vec<_>>()
630}
631
632pub fn pack_compressed_account(
633    compressed_account: CompressedAccountWithMerkleContext,
634    root_index: u16,
635    remaining_accounts: &mut RemainingAccounts,
636) -> PackedCompressedAccountWithMerkleContext {
637    pack_compressed_accounts(&[compressed_account], &[root_index], remaining_accounts)[0].clone()
638}
light_sdk/compressed_account.rs

light_sdk/
compressed_account.rs
