light_system_program/sdk/
compressed_account.rs

1use std::collections::HashMap;
2
3use anchor_lang::prelude::*;
4use light_hasher::{Hasher, Poseidon};
5use light_utils::hash_to_bn254_field_size_be;
6
7#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
8pub struct PackedCompressedAccountWithMerkleContext {
9    pub compressed_account: CompressedAccount,
10    pub merkle_context: PackedMerkleContext,
11    /// Index of root used in inclusion validity proof.
12    pub root_index: u16,
13    /// Placeholder to mark accounts read-only unimplemented set to false.
14    pub read_only: bool,
15}
16
17#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
18pub struct CompressedAccountWithMerkleContext {
19    pub compressed_account: CompressedAccount,
20    pub merkle_context: MerkleContext,
21}
22impl CompressedAccountWithMerkleContext {
23    pub fn hash(&self) -> Result<[u8; 32]> {
24        self.compressed_account.hash::<Poseidon>(
25            &self.merkle_context.merkle_tree_pubkey,
26            &self.merkle_context.leaf_index,
27        )
28    }
29}
30
31#[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Default)]
32pub struct MerkleContext {
33    pub merkle_tree_pubkey: Pubkey,
34    pub nullifier_queue_pubkey: Pubkey,
35    pub leaf_index: u32,
36    /// Index of leaf in queue. Placeholder of batched Merkle tree updates
37    /// currently unimplemented.
38    pub queue_index: Option<QueueIndex>,
39}
40
41#[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Default)]
42pub struct PackedMerkleContext {
43    pub merkle_tree_pubkey_index: u8,
44    pub nullifier_queue_pubkey_index: u8,
45    pub leaf_index: u32,
46    /// Index of leaf in queue. Placeholder of batched Merkle tree updates
47    /// currently unimplemented.
48    pub queue_index: Option<QueueIndex>,
49}
50
51#[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Default)]
52pub struct QueueIndex {
53    /// Id of queue in queue account.
54    pub queue_id: u8,
55    /// Index of compressed account hash in queue.
56    pub index: u16,
57}
58
59pub fn pack_merkle_context(
60    merkle_context: &[MerkleContext],
61    remaining_accounts: &mut HashMap<Pubkey, usize>,
62) -> Vec<PackedMerkleContext> {
63    let mut merkle_context_packed = merkle_context
64        .iter()
65        .map(|x| PackedMerkleContext {
66            leaf_index: x.leaf_index,
67            merkle_tree_pubkey_index: 0,     // will be assigned later
68            nullifier_queue_pubkey_index: 0, // will be assigned later
69            queue_index: None,
70        })
71        .collect::<Vec<PackedMerkleContext>>();
72    let mut index: usize = remaining_accounts.len();
73    for (i, params) in merkle_context.iter().enumerate() {
74        match remaining_accounts.get(&params.merkle_tree_pubkey) {
75            Some(_) => {}
76            None => {
77                remaining_accounts.insert(params.merkle_tree_pubkey, index);
78                index += 1;
79            }
80        };
81        merkle_context_packed[i].merkle_tree_pubkey_index =
82            *remaining_accounts.get(&params.merkle_tree_pubkey).unwrap() as u8;
83    }
84
85    for (i, params) in merkle_context.iter().enumerate() {
86        match remaining_accounts.get(&params.nullifier_queue_pubkey) {
87            Some(_) => {}
88            None => {
89                remaining_accounts.insert(params.nullifier_queue_pubkey, index);
90                index += 1;
91            }
92        };
93        merkle_context_packed[i].nullifier_queue_pubkey_index = *remaining_accounts
94            .get(&params.nullifier_queue_pubkey)
95            .unwrap() as u8;
96    }
97    merkle_context_packed
98}
99
100#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
101pub struct CompressedAccount {
102    pub owner: Pubkey,
103    pub lamports: u64,
104    pub address: Option<[u8; 32]>,
105    pub data: Option<CompressedAccountData>,
106}
107
108#[derive(Debug, PartialEq, Default, Clone, AnchorSerialize, AnchorDeserialize)]
109pub struct CompressedAccountData {
110    pub discriminator: [u8; 8],
111    pub data: Vec<u8>,
112    pub data_hash: [u8; 32],
113}
114
115/// Hashing scheme:
116/// H(owner || leaf_index || merkle_tree_pubkey || lamports || address || data.discriminator || data.data_hash)
117impl CompressedAccount {
118    pub fn hash_with_hashed_values<H: Hasher>(
119        &self,
120        &owner_hashed: &[u8; 32],
121        &merkle_tree_hashed: &[u8; 32],
122        leaf_index: &u32,
123    ) -> Result<[u8; 32]> {
124        let capacity = 3
125            + std::cmp::min(self.lamports, 1) as usize
126            + self.address.is_some() as usize
127            + self.data.is_some() as usize * 2;
128        let mut vec: Vec<&[u8]> = Vec::with_capacity(capacity);
129        vec.push(owner_hashed.as_slice());
130
131        // leaf index and merkle tree pubkey are used to make every compressed account hash unique
132        let leaf_index = leaf_index.to_le_bytes();
133        vec.push(leaf_index.as_slice());
134
135        vec.push(merkle_tree_hashed.as_slice());
136
137        // Lamports are only hashed if non-zero to safe CU
138        // For safety we prefix the lamports with 1 in 1 byte.
139        // Thus even if the discriminator has the same value as the lamports, the hash will be different.
140        let mut lamports_bytes = [1, 0, 0, 0, 0, 0, 0, 0, 0];
141        if self.lamports != 0 {
142            lamports_bytes[1..].copy_from_slice(&self.lamports.to_le_bytes());
143            vec.push(lamports_bytes.as_slice());
144        }
145
146        if self.address.is_some() {
147            vec.push(self.address.as_ref().unwrap().as_slice());
148        }
149
150        let mut discriminator_bytes = [2, 0, 0, 0, 0, 0, 0, 0, 0];
151        if let Some(data) = &self.data {
152            discriminator_bytes[1..].copy_from_slice(&data.discriminator);
153            vec.push(&discriminator_bytes);
154            vec.push(&data.data_hash);
155        }
156        let hash = H::hashv(&vec).map_err(ProgramError::from)?;
157        Ok(hash)
158    }
159
160    pub fn hash<H: Hasher>(
161        &self,
162        &merkle_tree_pubkey: &Pubkey,
163        leaf_index: &u32,
164    ) -> Result<[u8; 32]> {
165        self.hash_with_hashed_values::<H>(
166            &hash_to_bn254_field_size_be(&self.owner.to_bytes())
167                .unwrap()
168                .0,
169            &hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
170                .unwrap()
171                .0,
172            leaf_index,
173        )
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use light_hasher::Poseidon;
180    use solana_sdk::signature::{Keypair, Signer};
181
182    use super::*;
183    /// Tests:
184    /// 1. functional with all inputs set
185    /// 2. no data
186    /// 3. no address
187    /// 4. no address and no lamports
188    /// 5. no address and no data
189    /// 6. no address, no data, no lamports
190    #[test]
191    fn test_compressed_account_hash() {
192        let owner = Keypair::new().pubkey();
193        let address = [1u8; 32];
194        let data = CompressedAccountData {
195            discriminator: [1u8; 8],
196            data: vec![2u8; 32],
197            data_hash: [3u8; 32],
198        };
199        let lamports = 100;
200        let compressed_account = CompressedAccount {
201            owner,
202            lamports,
203            address: Some(address),
204            data: Some(data.clone()),
205        };
206        let merkle_tree_pubkey = Keypair::new().pubkey();
207        let leaf_index = 1;
208        let hash = compressed_account
209            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
210            .unwrap();
211        let hash_manual = Poseidon::hashv(&[
212            hash_to_bn254_field_size_be(&owner.to_bytes())
213                .unwrap()
214                .0
215                .as_slice(),
216            leaf_index.to_le_bytes().as_slice(),
217            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
218                .unwrap()
219                .0
220                .as_slice(),
221            [&[1u8], lamports.to_le_bytes().as_slice()]
222                .concat()
223                .as_slice(),
224            address.as_slice(),
225            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
226            &data.data_hash,
227        ])
228        .unwrap();
229        assert_eq!(hash, hash_manual);
230        assert_eq!(hash.len(), 32);
231
232        // no data
233        let compressed_account = CompressedAccount {
234            owner,
235            lamports,
236            address: Some(address),
237            data: None,
238        };
239        let no_data_hash = compressed_account
240            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
241            .unwrap();
242
243        let hash_manual = Poseidon::hashv(&[
244            hash_to_bn254_field_size_be(&owner.to_bytes())
245                .unwrap()
246                .0
247                .as_slice(),
248            leaf_index.to_le_bytes().as_slice(),
249            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
250                .unwrap()
251                .0
252                .as_slice(),
253            [&[1u8], lamports.to_le_bytes().as_slice()]
254                .concat()
255                .as_slice(),
256            address.as_slice(),
257        ])
258        .unwrap();
259        assert_eq!(no_data_hash, hash_manual);
260        assert_ne!(hash, no_data_hash);
261
262        // no address
263        let compressed_account = CompressedAccount {
264            owner,
265            lamports,
266            address: None,
267            data: Some(data.clone()),
268        };
269        let no_address_hash = compressed_account
270            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
271            .unwrap();
272        let hash_manual = Poseidon::hashv(&[
273            hash_to_bn254_field_size_be(&owner.to_bytes())
274                .unwrap()
275                .0
276                .as_slice(),
277            leaf_index.to_le_bytes().as_slice(),
278            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
279                .unwrap()
280                .0
281                .as_slice(),
282            [&[1u8], lamports.to_le_bytes().as_slice()]
283                .concat()
284                .as_slice(),
285            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
286            &data.data_hash,
287        ])
288        .unwrap();
289        assert_eq!(no_address_hash, hash_manual);
290        assert_ne!(hash, no_address_hash);
291        assert_ne!(no_data_hash, no_address_hash);
292
293        // no address no lamports
294        let compressed_account = CompressedAccount {
295            owner,
296            lamports: 0,
297            address: None,
298            data: Some(data.clone()),
299        };
300        let no_address_no_lamports_hash = compressed_account
301            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
302            .unwrap();
303        let hash_manual = Poseidon::hashv(&[
304            hash_to_bn254_field_size_be(&owner.to_bytes())
305                .unwrap()
306                .0
307                .as_slice(),
308            leaf_index.to_le_bytes().as_slice(),
309            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
310                .unwrap()
311                .0
312                .as_slice(),
313            [&[2u8], data.discriminator.as_slice()].concat().as_slice(),
314            &data.data_hash,
315        ])
316        .unwrap();
317        assert_eq!(no_address_no_lamports_hash, hash_manual);
318        assert_ne!(hash, no_address_no_lamports_hash);
319        assert_ne!(no_data_hash, no_address_no_lamports_hash);
320        assert_ne!(no_address_hash, no_address_no_lamports_hash);
321
322        // no address and no data
323        let compressed_account = CompressedAccount {
324            owner,
325            lamports,
326            address: None,
327            data: None,
328        };
329        let no_address_no_data_hash = compressed_account
330            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
331            .unwrap();
332        let hash_manual = Poseidon::hashv(&[
333            hash_to_bn254_field_size_be(&owner.to_bytes())
334                .unwrap()
335                .0
336                .as_slice(),
337            leaf_index.to_le_bytes().as_slice(),
338            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
339                .unwrap()
340                .0
341                .as_slice(),
342            [&[1u8], lamports.to_le_bytes().as_slice()]
343                .concat()
344                .as_slice(),
345        ])
346        .unwrap();
347        assert_eq!(no_address_no_data_hash, hash_manual);
348        assert_ne!(hash, no_address_no_data_hash);
349        assert_ne!(no_data_hash, no_address_no_data_hash);
350        assert_ne!(no_address_hash, no_address_no_data_hash);
351        assert_ne!(no_address_no_lamports_hash, no_address_no_data_hash);
352
353        // no address, no data, no lamports
354        let compressed_account = CompressedAccount {
355            owner,
356            lamports: 0,
357            address: None,
358            data: None,
359        };
360        let no_address_no_data_no_lamports_hash = compressed_account
361            .hash::<Poseidon>(&merkle_tree_pubkey, &leaf_index)
362            .unwrap();
363        let hash_manual = Poseidon::hashv(&[
364            hash_to_bn254_field_size_be(&owner.to_bytes())
365                .unwrap()
366                .0
367                .as_slice(),
368            leaf_index.to_le_bytes().as_slice(),
369            hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
370                .unwrap()
371                .0
372                .as_slice(),
373        ])
374        .unwrap();
375        assert_eq!(no_address_no_data_no_lamports_hash, hash_manual);
376        assert_ne!(no_address_no_data_hash, no_address_no_data_no_lamports_hash);
377        assert_ne!(hash, no_address_no_data_no_lamports_hash);
378        assert_ne!(no_data_hash, no_address_no_data_no_lamports_hash);
379        assert_ne!(no_address_hash, no_address_no_data_no_lamports_hash);
380        assert_ne!(
381            no_address_no_lamports_hash,
382            no_address_no_data_no_lamports_hash
383        );
384    }
385}
light_system_program/sdk/compressed_account.rs

light_system_program/sdk/
compressed_account.rs
