light_compressed_account/
compressed_account.rs

1use std::collections::HashMap;
2
3use light_hasher::{Hasher, Poseidon};
4
5use crate::{
6    address::pack_account,
7    hash_to_bn254_field_size_be,
8    instruction_data::{
9        data::OutputCompressedAccountWithPackedContext, zero_copy::ZCompressedAccount,
10    },
11    AnchorDeserialize, AnchorSerialize, CompressedAccountError, Pubkey, TreeType,
12};
13
14#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
15pub struct PackedCompressedAccountWithMerkleContext {
16    pub compressed_account: CompressedAccount,
17    pub merkle_context: PackedMerkleContext,
18    /// Index of root used in inclusion validity proof.
19    pub root_index: u16,
20    pub read_only: bool,
21}
22
23#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
24pub struct InCompressedAccountWithMerkleContext {
25    pub compressed_account: InCompressedAccount,
26    pub merkle_context: MerkleContext,
27}
28
29impl From<CompressedAccount> for InCompressedAccount {
30    fn from(value: CompressedAccount) -> Self {
31        let data = value.data.unwrap_or_default();
32        InCompressedAccount {
33            owner: value.owner,
34            lamports: value.lamports,
35            address: value.address,
36            discriminator: data.discriminator,
37            data_hash: data.data_hash,
38        }
39    }
40}
41
42#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
43pub struct PackedInCompressedAccountWithMerkleContext {
44    pub compressed_account: InCompressedAccount,
45    pub merkle_context: PackedMerkleContext,
46    /// Index of root used in inclusion validity proof.
47    pub root_index: u16,
48}
49
50impl From<PackedCompressedAccountWithMerkleContext> for PackedInCompressedAccountWithMerkleContext {
51    fn from(value: PackedCompressedAccountWithMerkleContext) -> Self {
52        Self {
53            compressed_account: value.compressed_account.into(),
54            merkle_context: value.merkle_context,
55            root_index: value.root_index,
56        }
57    }
58}
59
60impl From<CompressedAccountWithMerkleContext> for InCompressedAccountWithMerkleContext {
61    fn from(value: CompressedAccountWithMerkleContext) -> Self {
62        Self {
63            compressed_account: value.compressed_account.into(),
64            merkle_context: value.merkle_context,
65        }
66    }
67}
68
69#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
70pub struct CompressedAccountWithMerkleContext {
71    pub compressed_account: CompressedAccount,
72    pub merkle_context: MerkleContext,
73}
74
75impl CompressedAccountWithMerkleContext {
76    pub fn hash(&self) -> Result<[u8; 32], CompressedAccountError> {
77        self.compressed_account.hash(
78            &self.merkle_context.merkle_tree_pubkey,
79            &self.merkle_context.leaf_index,
80            self.merkle_context.tree_type == TreeType::StateV2,
81        )
82    }
83}
84
85impl CompressedAccountWithMerkleContext {
86    pub fn into_read_only(
87        &self,
88        root_index: Option<u16>,
89    ) -> Result<ReadOnlyCompressedAccount, CompressedAccountError> {
90        let account_hash = self.hash()?;
91        let merkle_context = if root_index.is_none() {
92            let mut merkle_context = self.merkle_context;
93            merkle_context.prove_by_index = true;
94            merkle_context
95        } else {
96            self.merkle_context
97        };
98        Ok(ReadOnlyCompressedAccount {
99            account_hash,
100            merkle_context,
101            root_index: root_index.unwrap_or_default(),
102        })
103    }
104
105    pub fn pack(
106        &self,
107        root_index: Option<u16>,
108        remaining_accounts: &mut HashMap<Pubkey, usize>,
109    ) -> Result<PackedCompressedAccountWithMerkleContext, CompressedAccountError> {
110        Ok(PackedCompressedAccountWithMerkleContext {
111            compressed_account: self.compressed_account.clone(),
112            merkle_context: PackedMerkleContext {
113                merkle_tree_pubkey_index: pack_account(
114                    &self.merkle_context.merkle_tree_pubkey,
115                    remaining_accounts,
116                ),
117                queue_pubkey_index: pack_account(
118                    &self.merkle_context.queue_pubkey,
119                    remaining_accounts,
120                ),
121                leaf_index: self.merkle_context.leaf_index,
122                prove_by_index: root_index.is_none(),
123            },
124            root_index: root_index.unwrap_or_default(),
125            read_only: false,
126        })
127    }
128}
129#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
130pub struct ReadOnlyCompressedAccount {
131    pub account_hash: [u8; 32],
132    pub merkle_context: MerkleContext,
133    pub root_index: u16,
134}
135
136#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
137pub struct PackedReadOnlyCompressedAccount {
138    pub account_hash: [u8; 32],
139    pub merkle_context: PackedMerkleContext,
140    pub root_index: u16,
141}
142
143#[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Default)]
144pub struct MerkleContext {
145    pub merkle_tree_pubkey: Pubkey,
146    pub queue_pubkey: Pubkey,
147    pub leaf_index: u32,
148    pub prove_by_index: bool,
149    pub tree_type: TreeType,
150}
151
152#[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Default)]
153pub struct PackedMerkleContext {
154    pub merkle_tree_pubkey_index: u8,
155    pub queue_pubkey_index: u8,
156    pub leaf_index: u32,
157    pub prove_by_index: bool,
158}
159
160pub fn pack_compressed_accounts(
161    compressed_accounts: &[CompressedAccountWithMerkleContext],
162    root_indices: &[Option<u16>],
163    remaining_accounts: &mut HashMap<Pubkey, usize>,
164) -> Vec<PackedCompressedAccountWithMerkleContext> {
165    compressed_accounts
166        .iter()
167        .zip(root_indices.iter())
168        .map(|(x, root_index)| {
169            let mut merkle_context = x.merkle_context;
170            let root_index = if let Some(root) = root_index {
171                *root
172            } else {
173                merkle_context.prove_by_index = true;
174                0
175            };
176
177            PackedCompressedAccountWithMerkleContext {
178                compressed_account: x.compressed_account.clone(),
179                merkle_context: pack_merkle_context(&[merkle_context], remaining_accounts)[0],
180                root_index,
181                read_only: false,
182            }
183        })
184        .collect::<Vec<_>>()
185}
186
187pub fn pack_output_compressed_accounts(
188    compressed_accounts: &[CompressedAccount],
189    merkle_trees: &[Pubkey],
190    remaining_accounts: &mut HashMap<Pubkey, usize>,
191) -> Vec<OutputCompressedAccountWithPackedContext> {
192    compressed_accounts
193        .iter()
194        .zip(merkle_trees.iter())
195        .map(|(x, tree)| OutputCompressedAccountWithPackedContext {
196            compressed_account: x.clone(),
197            merkle_tree_index: pack_account(tree, remaining_accounts),
198        })
199        .collect::<Vec<_>>()
200}
201
202pub fn pack_merkle_context(
203    merkle_context: &[MerkleContext],
204    remaining_accounts: &mut HashMap<Pubkey, usize>,
205) -> Vec<PackedMerkleContext> {
206    merkle_context
207        .iter()
208        .map(|merkle_context| PackedMerkleContext {
209            leaf_index: merkle_context.leaf_index,
210            merkle_tree_pubkey_index: pack_account(
211                &merkle_context.merkle_tree_pubkey,
212                remaining_accounts,
213            ),
214            queue_pubkey_index: pack_account(&merkle_context.queue_pubkey, remaining_accounts),
215            prove_by_index: merkle_context.prove_by_index,
216        })
217        .collect::<Vec<_>>()
218}
219
220#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
221pub struct CompressedAccount {
222    pub owner: Pubkey,
223    pub lamports: u64,
224    pub address: Option<[u8; 32]>,
225    pub data: Option<CompressedAccountData>,
226}
227
228#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
229pub struct InCompressedAccount {
230    pub owner: Pubkey,
231    pub lamports: u64,
232    pub discriminator: [u8; 8],
233    pub data_hash: [u8; 32],
234    pub address: Option<[u8; 32]>,
235}
236
237#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
238pub struct CompressedAccountData {
239    pub discriminator: [u8; 8],
240    pub data: Vec<u8>,
241    pub data_hash: [u8; 32],
242}
243
244pub fn hash_with_hashed_values(
245    lamports: &u64,
246    address: Option<&[u8]>,
247    data: Option<(&[u8], &[u8])>,
248    owner_hashed: &[u8; 32],
249    merkle_tree_hashed: &[u8; 32],
250    leaf_index: &u32,
251    is_batched: bool,
252) -> Result<[u8; 32], CompressedAccountError> {
253    // TODO: replace with array
254    let capacity = 3
255        + std::cmp::min(*lamports, 1) as usize
256        + address.is_some() as usize
257        + data.is_some() as usize * 2;
258    let mut vec: Vec<&[u8]> = Vec::with_capacity(capacity);
259    vec.push(owner_hashed.as_slice());
260
261    // leaf index and merkle tree pubkey are used to make every compressed account hash unique
262    let mut leaf_index_bytes = [0u8; 32];
263    if is_batched {
264        leaf_index_bytes[28..].copy_from_slice(&leaf_index.to_be_bytes());
265    } else {
266        leaf_index_bytes[28..].copy_from_slice(&leaf_index.to_le_bytes());
267    };
268    vec.push(leaf_index_bytes.as_slice());
269
270    vec.push(merkle_tree_hashed.as_slice());
271
272    // Lamports are only hashed if non-zero to save CU.
273    // For safety, we prefix lamports with 1 in 1 byte.
274    // Thus, even if the discriminator has the same value as the lamports, the hash will be different.
275    let mut lamports_bytes = [0u8; 32];
276    if *lamports != 0 {
277        if is_batched {
278            lamports_bytes[24..].copy_from_slice(&lamports.to_be_bytes());
279        } else {
280            lamports_bytes[24..].copy_from_slice(&lamports.to_le_bytes());
281        };
282        lamports_bytes[23] = 1;
283
284        vec.push(lamports_bytes.as_slice());
285    }
286
287    if let Some(address) = address {
288        vec.push(address);
289    }
290
291    let mut discriminator_bytes = [0u8; 32];
292    if let Some((discriminator, data_hash)) = data {
293        discriminator_bytes[24..].copy_from_slice(discriminator);
294        discriminator_bytes[23] = 2;
295        vec.push(&discriminator_bytes);
296        vec.push(data_hash);
297    }
298
299    Ok(Poseidon::hashv(&vec)?)
300}
301
302/// Hashing scheme:
303/// H(owner || leaf_index || merkle_tree_pubkey || lamports || address || data.discriminator || data.data_hash)
304impl CompressedAccount {
305    pub fn hash_with_hashed_values(
306        &self,
307        owner_hashed: &[u8; 32],
308        merkle_tree_hashed: &[u8; 32],
309        leaf_index: &u32,
310        is_batched: bool,
311    ) -> Result<[u8; 32], CompressedAccountError> {
312        hash_with_hashed_values(
313            &self.lamports,
314            self.address.as_ref().map(|x| x.as_slice()),
315            self.data
316                .as_ref()
317                .map(|x| (x.discriminator.as_slice(), x.data_hash.as_slice())),
318            owner_hashed,
319            merkle_tree_hashed,
320            leaf_index,
321            is_batched,
322        )
323    }
324
325    pub fn hash(
326        &self,
327        &merkle_tree_pubkey: &Pubkey,
328        leaf_index: &u32,
329        is_batched: bool,
330    ) -> Result<[u8; 32], CompressedAccountError> {
331        let hashed_mt = hash_to_bn254_field_size_be(merkle_tree_pubkey.as_ref());
332        self.hash_with_hashed_values(
333            &hash_to_bn254_field_size_be(self.owner.as_ref()),
334            &hashed_mt,
335            leaf_index,
336            is_batched,
337        )
338    }
339}
340
341/// Hashing scheme:
342/// H(owner || leaf_index || merkle_tree_pubkey || lamports || address || data.discriminator || data.data_hash)
343impl ZCompressedAccount<'_> {
344    pub fn hash_with_hashed_values(
345        &self,
346        owner_hashed: &[u8; 32],
347        merkle_tree_hashed: &[u8; 32],
348        leaf_index: &u32,
349        is_batched: bool,
350    ) -> Result<[u8; 32], CompressedAccountError> {
351        hash_with_hashed_values(
352            &(self.lamports.into()),
353            self.address.as_ref().map(|x| x.as_slice()),
354            self.data
355                .as_ref()
356                .map(|x| (x.discriminator.as_slice(), x.data_hash.as_slice())),
357            owner_hashed,
358            merkle_tree_hashed,
359            leaf_index,
360            is_batched,
361        )
362    }
363    pub fn hash(
364        &self,
365        &merkle_tree_pubkey: &[u8; 32],
366        leaf_index: &u32,
367        is_batched: bool,
368    ) -> Result<[u8; 32], CompressedAccountError> {
369        self.hash_with_hashed_values(
370            &hash_to_bn254_field_size_be(&self.owner.to_bytes()),
371            &hash_to_bn254_field_size_be(merkle_tree_pubkey.as_slice()),
372            leaf_index,
373            is_batched,
374        )
375    }
376}
377
378#[cfg(not(feature = "pinocchio"))]
379#[cfg(test)]
380mod tests {
381    use light_hasher::Poseidon;
382    use light_zero_copy::borsh::Deserialize;
383    use num_bigint::BigUint;
384    use rand::Rng;
385
386    use super::*;
387    /// Tests:
388    /// 1. functional with all inputs set
389    /// 2. no data
390    /// 3. no address
391    /// 4. no address and no lamports
392    /// 5. no address and no data
393    /// 6. no address, no data, no lamports
394    #[test]
395    fn test_compressed_account_hash() {
396        let owner = Pubkey::new_unique();
397        let address = [1u8; 32];
398        let data = CompressedAccountData {
399            discriminator: [1u8; 8],
400            data: vec![2u8; 32],
401            data_hash: [3u8; 32],
402        };
403        let lamports = 100;
404        let compressed_account = CompressedAccount {
405            owner,
406            lamports,
407            address: Some(address),
408            data: Some(data.clone()),
409        };
410        let merkle_tree_pubkey = Pubkey::new_unique();
411        let leaf_index: u32 = 1;
412        let hash = compressed_account
413            .hash(&merkle_tree_pubkey, &leaf_index, false)
414            .unwrap();
415        let hash_manual = Poseidon::hashv(&[
416            hash_to_bn254_field_size_be(&owner.to_bytes()).as_slice(),
417            leaf_index.to_le_bytes().as_slice(),
418            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes()).as_slice(),
419            [&[1u8], lamports.to_le_bytes().as_slice()]
420                .concat()
421                .as_slice(),
422            address.as_slice(),
423            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
424            &data.data_hash,
425        ])
426        .unwrap();
427        assert_eq!(hash, hash_manual);
428        assert_eq!(hash.len(), 32);
429
430        // no data
431        let compressed_account = CompressedAccount {
432            owner,
433            lamports,
434            address: Some(address),
435            data: None,
436        };
437        let no_data_hash = compressed_account
438            .hash(&merkle_tree_pubkey, &leaf_index, false)
439            .unwrap();
440
441        let hash_manual = Poseidon::hashv(&[
442            hash_to_bn254_field_size_be(&owner.to_bytes()).as_slice(),
443            leaf_index.to_le_bytes().as_slice(),
444            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes()).as_slice(),
445            [&[1u8], lamports.to_le_bytes().as_slice()]
446                .concat()
447                .as_slice(),
448            address.as_slice(),
449        ])
450        .unwrap();
451        assert_eq!(no_data_hash, hash_manual);
452        assert_ne!(hash, no_data_hash);
453
454        // no address
455        let compressed_account = CompressedAccount {
456            owner,
457            lamports,
458            address: None,
459            data: Some(data.clone()),
460        };
461        let no_address_hash = compressed_account
462            .hash(&merkle_tree_pubkey, &leaf_index, false)
463            .unwrap();
464        let hash_manual = Poseidon::hashv(&[
465            hash_to_bn254_field_size_be(&owner.to_bytes()).as_slice(),
466            leaf_index.to_le_bytes().as_slice(),
467            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes()).as_slice(),
468            [&[1u8], lamports.to_le_bytes().as_slice()]
469                .concat()
470                .as_slice(),
471            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
472            &data.data_hash,
473        ])
474        .unwrap();
475        assert_eq!(no_address_hash, hash_manual);
476        assert_ne!(hash, no_address_hash);
477        assert_ne!(no_data_hash, no_address_hash);
478
479        // no address no lamports
480        let compressed_account = CompressedAccount {
481            owner,
482            lamports: 0,
483            address: None,
484            data: Some(data.clone()),
485        };
486        let no_address_no_lamports_hash = compressed_account
487            .hash(&merkle_tree_pubkey, &leaf_index, false)
488            .unwrap();
489        let hash_manual = Poseidon::hashv(&[
490            hash_to_bn254_field_size_be(&owner.to_bytes()).as_slice(),
491            leaf_index.to_le_bytes().as_slice(),
492            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes()).as_slice(),
493            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
494            &data.data_hash,
495        ])
496        .unwrap();
497        assert_eq!(no_address_no_lamports_hash, hash_manual);
498        assert_ne!(hash, no_address_no_lamports_hash);
499        assert_ne!(no_data_hash, no_address_no_lamports_hash);
500        assert_ne!(no_address_hash, no_address_no_lamports_hash);
501
502        // no address and no data
503        let compressed_account = CompressedAccount {
504            owner,
505            lamports,
506            address: None,
507            data: None,
508        };
509        let no_address_no_data_hash = compressed_account
510            .hash(&merkle_tree_pubkey, &leaf_index, false)
511            .unwrap();
512        let hash_manual = Poseidon::hashv(&[
513            hash_to_bn254_field_size_be(&owner.to_bytes()).as_slice(),
514            leaf_index.to_le_bytes().as_slice(),
515            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes()).as_slice(),
516            [&[1u8], lamports.to_le_bytes().as_slice()]
517                .concat()
518                .as_slice(),
519        ])
520        .unwrap();
521        assert_eq!(no_address_no_data_hash, hash_manual);
522        assert_ne!(hash, no_address_no_data_hash);
523        assert_ne!(no_data_hash, no_address_no_data_hash);
524        assert_ne!(no_address_hash, no_address_no_data_hash);
525        assert_ne!(no_address_no_lamports_hash, no_address_no_data_hash);
526
527        // no address, no data, no lamports
528        let compressed_account = CompressedAccount {
529            owner,
530            lamports: 0,
531            address: None,
532            data: None,
533        };
534        let no_address_no_data_no_lamports_hash = compressed_account
535            .hash(&merkle_tree_pubkey, &leaf_index, false)
536            .unwrap();
537        let hash_manual = Poseidon::hashv(&[
538            hash_to_bn254_field_size_be(&owner.to_bytes()).as_slice(),
539            leaf_index.to_le_bytes().as_slice(),
540            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes()).as_slice(),
541        ])
542        .unwrap();
543        assert_eq!(no_address_no_data_no_lamports_hash, hash_manual);
544        assert_ne!(no_address_no_data_hash, no_address_no_data_no_lamports_hash);
545        assert_ne!(hash, no_address_no_data_no_lamports_hash);
546        assert_ne!(no_data_hash, no_address_no_data_no_lamports_hash);
547        assert_ne!(no_address_hash, no_address_no_data_no_lamports_hash);
548        assert_ne!(
549            no_address_no_lamports_hash,
550            no_address_no_data_no_lamports_hash
551        );
552    }
553
554    #[test]
555    fn reference() {
556        let owner = Pubkey::new_from_array([
557            0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
558            0, 0, 0,
559        ]);
560        let address = [
561            0, 21, 245, 15, 61, 157, 224, 84, 69, 48, 190, 72, 43, 19, 47, 25, 14, 118, 20, 147,
562            40, 141, 175, 33, 233, 58, 36, 179, 73, 137, 84, 99,
563        ];
564
565        let data = CompressedAccountData {
566            discriminator: [0, 0, 0, 0, 0, 0, 0, 1],
567            data: vec![2u8; 31],
568            data_hash: Poseidon::hash(&[vec![2u8; 31], vec![0u8]].concat()).unwrap(),
569        };
570        let lamports = 100;
571        let compressed_account = CompressedAccount {
572            owner,
573            lamports,
574            address: Some(address),
575            data: Some(data.clone()),
576        };
577        let bytes: Vec<u8> = compressed_account.try_to_vec().unwrap();
578        let merkle_tree_pubkey = Pubkey::new_from_array([
579            0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
580            0, 0, 0,
581        ]);
582
583        let leaf_index: u32 = 1;
584        let hash = compressed_account
585            .hash(&merkle_tree_pubkey, &leaf_index, false)
586            .unwrap();
587        let (z_account, _) = ZCompressedAccount::zero_copy_at(&bytes).unwrap();
588        let z_hash = z_account
589            .hash(&merkle_tree_pubkey.to_bytes(), &leaf_index, false)
590            .unwrap();
591        let manual_hash = {
592            let mut hasher = light_poseidon::Poseidon::<Fr>::new_circom(7).unwrap();
593            use ark_bn254::Fr;
594            use ark_ff::{BigInteger, PrimeField};
595            let hashed_owner = hash_to_bn254_field_size_be(&owner.to_bytes());
596            let owner = Fr::from_be_bytes_mod_order(hashed_owner.as_slice());
597            let leaf_index = Fr::from_be_bytes_mod_order(leaf_index.to_le_bytes().as_ref());
598            let hashed_mt = hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes());
599            let merkle_tree_pubkey = Fr::from_be_bytes_mod_order(hashed_mt.as_slice());
600            let lamports = Fr::from_be_bytes_mod_order(lamports.to_le_bytes().as_ref())
601                + Fr::from_be_bytes_mod_order(&[1u8, 0, 0, 0, 0, 0, 0, 0, 0]);
602            let address = Fr::from_be_bytes_mod_order(address.as_slice());
603            let discriminator = Fr::from_be_bytes_mod_order(data.discriminator.as_ref());
604            let domain_separated_discriminator =
605                Fr::from_be_bytes_mod_order(&[2, 0, 0, 0, 0, 0, 0, 0, 0]);
606            let data_discriminator = discriminator + domain_separated_discriminator;
607            use light_poseidon::PoseidonHasher;
608            let inputs = [
609                owner,
610                leaf_index,
611                merkle_tree_pubkey,
612                lamports,
613                address,
614                data_discriminator,
615                Fr::from_be_bytes_mod_order(data.data_hash.as_ref()),
616            ];
617            hasher.hash(&inputs).unwrap().into_bigint().to_bytes_be()
618        };
619        assert_eq!(hash.to_vec(), manual_hash);
620        assert_eq!(z_hash.to_vec(), manual_hash);
621        assert_eq!(hash.len(), 32);
622
623        let manual_hash_new = {
624            let mut hasher = light_poseidon::Poseidon::<Fr>::new_circom(7).unwrap();
625            use ark_bn254::Fr;
626            use ark_ff::{BigInteger, PrimeField};
627            let hashed_owner = hash_to_bn254_field_size_be(&owner.to_bytes());
628            let owner = Fr::from_be_bytes_mod_order(hashed_owner.as_slice());
629            let leaf_index = Fr::from_be_bytes_mod_order(leaf_index.to_be_bytes().as_ref());
630            let hashed_mt = hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes());
631            let merkle_tree_pubkey = Fr::from_be_bytes_mod_order(hashed_mt.as_slice());
632            let lamports = Fr::from_be_bytes_mod_order(lamports.to_be_bytes().as_ref())
633                + Fr::from_be_bytes_mod_order(&[1u8, 0, 0, 0, 0, 0, 0, 0, 0]);
634            let address = Fr::from_be_bytes_mod_order(address.as_slice());
635            let discriminator = Fr::from_be_bytes_mod_order(data.discriminator.as_ref());
636            let domain_separated_discriminator =
637                Fr::from_be_bytes_mod_order(&[2, 0, 0, 0, 0, 0, 0, 0, 0]);
638            let data_discriminator = discriminator + domain_separated_discriminator;
639            use light_poseidon::PoseidonHasher;
640            let inputs = [
641                owner,
642                leaf_index,
643                merkle_tree_pubkey,
644                lamports,
645                address,
646                data_discriminator,
647                Fr::from_be_bytes_mod_order(data.data_hash.as_ref()),
648            ];
649            hasher.hash(&inputs).unwrap().into_bigint().to_bytes_be()
650        };
651        let hash = compressed_account
652            .hash(&merkle_tree_pubkey, &leaf_index, true)
653            .unwrap();
654        let z_hash = z_account
655            .hash(&merkle_tree_pubkey.to_bytes(), &leaf_index, true)
656            .unwrap();
657        assert_ne!(hash.to_vec(), manual_hash);
658        assert_eq!(hash.to_vec(), manual_hash_new);
659        assert_eq!(z_hash.to_vec(), manual_hash_new);
660        assert_eq!(hash.len(), 32);
661        use std::str::FromStr;
662        let circuit_reference_value = BigUint::from_str(
663            "15638319165413000277907073391141043184436601830909724248083671155000605125280",
664        )
665        .unwrap()
666        .to_bytes_be();
667        println!(
668            "lamports domain: {:?}",
669            BigUint::from_bytes_be(&[1u8, 0, 0, 0, 0, 0, 0, 0, 0]).to_string()
670        );
671        println!(
672            "discriminator domain: {:?}",
673            BigUint::from_bytes_be(&[2u8, 0, 0, 0, 0, 0, 0, 0, 0]).to_string()
674        );
675        assert_eq!(hash.to_vec(), circuit_reference_value);
676    }
677
678    impl CompressedAccount {
679        pub fn legacy_hash_with_values<H: Hasher>(
680            &self,
681            &owner_hashed: &[u8; 32],
682            &merkle_tree_hashed: &[u8; 32],
683            leaf_index: &u32,
684        ) -> Result<[u8; 32], CompressedAccountError> {
685            let capacity = 3
686                + std::cmp::min(self.lamports, 1) as usize
687                + self.address.is_some() as usize
688                + self.data.is_some() as usize * 2;
689            let mut vec: Vec<&[u8]> = Vec::with_capacity(capacity);
690            vec.push(owner_hashed.as_slice());
691
692            // leaf index and merkle tree pubkey are used to make every compressed account hash unique
693            let leaf_index = leaf_index.to_le_bytes();
694            vec.push(leaf_index.as_slice());
695
696            vec.push(merkle_tree_hashed.as_slice());
697
698            // Lamports are only hashed if non-zero to save CU.
699            // For safety, we prefix lamports with 1 in 1 byte.
700            // Thus, even if the discriminator has the same value as the lamports, the hash will be different.
701            let mut lamports_bytes = [1, 0, 0, 0, 0, 0, 0, 0, 0];
702            if self.lamports != 0 {
703                lamports_bytes[1..].copy_from_slice(&self.lamports.to_le_bytes());
704                vec.push(lamports_bytes.as_slice());
705            }
706
707            if self.address.is_some() {
708                vec.push(self.address.as_ref().unwrap().as_slice());
709            }
710
711            let mut discriminator_bytes = [2, 0, 0, 0, 0, 0, 0, 0, 0];
712            if let Some(data) = &self.data {
713                discriminator_bytes[1..].copy_from_slice(&data.discriminator);
714                vec.push(&discriminator_bytes);
715                vec.push(&data.data_hash);
716            }
717            let hash = H::hashv(&vec)?;
718            Ok(hash)
719        }
720
721        pub fn hash_legacy<H: Hasher>(
722            &self,
723            &merkle_tree_pubkey: &Pubkey,
724            leaf_index: &u32,
725        ) -> Result<[u8; 32], CompressedAccountError> {
726            let hashed_mt = hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes());
727            self.legacy_hash_with_values::<H>(
728                &hash_to_bn254_field_size_be(&self.owner.to_bytes()),
729                &hashed_mt,
730                leaf_index,
731            )
732        }
733    }
734
735    fn equivalency_of_hash_functions_rnd_iters<const ITERS: usize>() {
736        let mut rng = rand::thread_rng();
737
738        for _ in 0..ITERS {
739            let account = CompressedAccount {
740                owner: Pubkey::new_unique(),
741                lamports: rng.gen::<u64>(),
742                address: if rng.gen_bool(0.5) {
743                    let mut address = rng.gen::<[u8; 32]>();
744                    address[0] = 0;
745                    Some(address)
746                } else {
747                    None
748                },
749                data: if rng.gen_bool(0.5) {
750                    Some(CompressedAccountData {
751                        discriminator: rng.gen(),
752                        data: Vec::new(), // not used in hash
753                        data_hash: Poseidon::hash(rng.gen::<u64>().to_be_bytes().as_slice())
754                            .unwrap(),
755                    })
756                } else {
757                    None
758                },
759            };
760            let leaf_index = rng.gen::<u32>();
761            let merkle_tree_pubkey = Pubkey::new_unique();
762            let hash_legacy = account
763                .hash_legacy::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
764                .unwrap();
765            let hash = account
766                .hash(&merkle_tree_pubkey, &leaf_index, false)
767                .unwrap();
768            let bytes: Vec<u8> = account.try_to_vec().unwrap();
769            let (z_account, _) = ZCompressedAccount::zero_copy_at(bytes.as_slice()).unwrap();
770            let z_hash = z_account
771                .hash(&merkle_tree_pubkey.to_bytes(), &leaf_index, false)
772                .unwrap();
773            assert_eq!(hash_legacy, hash);
774            assert_eq!(hash, z_hash);
775        }
776    }
777
778    #[test]
779    fn equivalency_of_hash_functions() {
780        equivalency_of_hash_functions_rnd_iters::<100>();
781    }
782}
light_compressed_account/compressed_account.rs

light_compressed_account/
compressed_account.rs
