light_system_program/invoke/
verify_state_proof.rs

1use crate::{
2    sdk::{accounts::InvokeAccounts, compressed_account::PackedCompressedAccountWithMerkleContext},
3    NewAddressParamsPacked,
4};
5use account_compression::{
6    utils::check_discrimininator::check_discriminator, AddressMerkleTreeAccount,
7    StateMerkleTreeAccount,
8};
9use anchor_lang::{prelude::*, Bumps};
10use light_concurrent_merkle_tree::zero_copy::ConcurrentMerkleTreeZeroCopy;
11use light_hasher::Poseidon;
12use light_indexed_merkle_tree::zero_copy::IndexedMerkleTreeZeroCopy;
13use light_macros::heap_neutral;
14use light_utils::hash_to_bn254_field_size_be;
15use light_verifier::{
16    verify_create_addresses_and_merkle_proof_zkp, verify_create_addresses_zkp,
17    verify_merkle_proof_zkp, CompressedProof,
18};
19use std::mem;
20
21#[inline(never)]
22#[heap_neutral]
23pub fn fetch_input_compressed_account_roots<
24    'a,
25    'b,
26    'c: 'info,
27    'info,
28    A: InvokeAccounts<'info> + Bumps,
29>(
30    input_compressed_accounts_with_merkle_context: &'a [PackedCompressedAccountWithMerkleContext],
31    ctx: &'a Context<'a, 'b, 'c, 'info, A>,
32    roots: &'a mut [[u8; 32]],
33) -> Result<()> {
34    for (i, input_compressed_account_with_context) in input_compressed_accounts_with_merkle_context
35        .iter()
36        .enumerate()
37    {
38        let merkle_tree = &ctx.remaining_accounts[input_compressed_account_with_context
39            .merkle_context
40            .merkle_tree_pubkey_index as usize];
41        let merkle_tree = merkle_tree.try_borrow_data()?;
42        check_discriminator::<StateMerkleTreeAccount>(&merkle_tree)?;
43        let merkle_tree = ConcurrentMerkleTreeZeroCopy::<Poseidon, 26>::from_bytes_zero_copy(
44            &merkle_tree[8 + mem::size_of::<StateMerkleTreeAccount>()..],
45        )
46        .map_err(ProgramError::from)?;
47        let fetched_roots = &merkle_tree.roots;
48
49        roots[i] = fetched_roots[input_compressed_account_with_context.root_index as usize];
50    }
51    Ok(())
52}
53
54#[inline(never)]
55#[heap_neutral]
56pub fn fetch_roots_address_merkle_tree<
57    'a,
58    'b,
59    'c: 'info,
60    'info,
61    A: InvokeAccounts<'info> + Bumps,
62>(
63    new_address_params: &'a [NewAddressParamsPacked],
64    ctx: &'a Context<'a, 'b, 'c, 'info, A>,
65    roots: &'a mut [[u8; 32]],
66) -> Result<()> {
67    for (i, new_address_param) in new_address_params.iter().enumerate() {
68        let merkle_tree = ctx.remaining_accounts
69            [new_address_param.address_merkle_tree_account_index as usize]
70            .to_account_info();
71        let merkle_tree = merkle_tree.try_borrow_data()?;
72        check_discriminator::<AddressMerkleTreeAccount>(&merkle_tree)?;
73        let merkle_tree =
74            IndexedMerkleTreeZeroCopy::<Poseidon, usize, 26, 16>::from_bytes_zero_copy(
75                &merkle_tree[8 + mem::size_of::<AddressMerkleTreeAccount>()..],
76            )
77            .map_err(ProgramError::from)?;
78        let fetched_roots = &merkle_tree.roots;
79
80        roots[i] = fetched_roots[new_address_param.address_merkle_tree_root_index as usize];
81    }
82    Ok(())
83}
84
85/// Hashes the input compressed accounts and stores the results in the leaves array.
86/// Merkle tree pubkeys are hashed and stored in the hashed_pubkeys array.
87/// Merkle tree pubkeys should be ordered for efficiency.
88#[inline(never)]
89#[heap_neutral]
90#[allow(unused_mut)]
91pub fn hash_input_compressed_accounts<'a, 'b, 'c: 'info, 'info>(
92    remaining_accounts: &'a [AccountInfo<'info>],
93    input_compressed_accounts_with_merkle_context: &'a [PackedCompressedAccountWithMerkleContext],
94    leaves: &'a mut [[u8; 32]],
95    addresses: &'a mut [Option<[u8; 32]>],
96    hashed_pubkeys: &'a mut Vec<(Pubkey, [u8; 32])>,
97) -> Result<()> {
98    let mut owner_pubkey = input_compressed_accounts_with_merkle_context[0]
99        .compressed_account
100        .owner;
101    let mut hashed_owner = hash_to_bn254_field_size_be(&owner_pubkey.to_bytes())
102        .unwrap()
103        .0;
104    hashed_pubkeys.push((owner_pubkey, hashed_owner));
105    #[allow(unused)]
106    let mut current_hashed_mt = [0u8; 32];
107
108    let mut current_mt_index: i16 = -1;
109    for (j, input_compressed_account_with_context) in input_compressed_accounts_with_merkle_context
110        .iter()
111        .enumerate()
112    {
113        // For heap neutrality we cannot allocate new heap memory in this function.
114        match &input_compressed_account_with_context
115            .compressed_account
116            .address
117        {
118            Some(address) => addresses[j] = Some(*address),
119            None => {}
120        };
121        if input_compressed_account_with_context
122            .merkle_context
123            .queue_index
124            .is_some()
125        {
126            unimplemented!("Queue index is not supported.");
127        }
128
129        #[allow(clippy::comparison_chain)]
130        if current_mt_index
131            != input_compressed_account_with_context
132                .merkle_context
133                .merkle_tree_pubkey_index as i16
134        {
135            current_mt_index = input_compressed_account_with_context
136                .merkle_context
137                .merkle_tree_pubkey_index as i16;
138            let merkle_tree_pubkey = remaining_accounts[input_compressed_account_with_context
139                .merkle_context
140                .merkle_tree_pubkey_index
141                as usize]
142                .key();
143            current_hashed_mt = match hashed_pubkeys.iter().find(|x| x.0 == merkle_tree_pubkey) {
144                Some(hashed_merkle_tree_pubkey) => hashed_merkle_tree_pubkey.1,
145                None => {
146                    let hashed_merkle_tree_pubkey =
147                        hash_to_bn254_field_size_be(&merkle_tree_pubkey.to_bytes())
148                            .unwrap()
149                            .0;
150                    hashed_pubkeys.push((merkle_tree_pubkey, hashed_merkle_tree_pubkey));
151                    hashed_merkle_tree_pubkey
152                }
153            };
154        }
155        // Without cpi context all input compressed accounts have the same owner.
156        // With cpi context the owners will be different.
157        if owner_pubkey
158            != input_compressed_account_with_context
159                .compressed_account
160                .owner
161        {
162            owner_pubkey = input_compressed_account_with_context
163                .compressed_account
164                .owner;
165            hashed_owner = match hashed_pubkeys.iter().find(|x| {
166                x.0 == input_compressed_account_with_context
167                    .compressed_account
168                    .owner
169            }) {
170                Some(hashed_owner) => hashed_owner.1,
171                None => {
172                    let hashed_owner = hash_to_bn254_field_size_be(
173                        &input_compressed_account_with_context
174                            .compressed_account
175                            .owner
176                            .to_bytes(),
177                    )
178                    .unwrap()
179                    .0;
180                    hashed_pubkeys.push((
181                        input_compressed_account_with_context
182                            .compressed_account
183                            .owner,
184                        hashed_owner,
185                    ));
186                    hashed_owner
187                }
188            };
189        }
190        leaves[j] = input_compressed_account_with_context
191            .compressed_account
192            .hash_with_hashed_values::<Poseidon>(
193                &hashed_owner,
194                &current_hashed_mt,
195                &input_compressed_account_with_context
196                    .merkle_context
197                    .leaf_index,
198            )?;
199    }
200    Ok(())
201}
202
203#[heap_neutral]
204pub fn verify_state_proof(
205    roots: &[[u8; 32]],
206    leaves: &[[u8; 32]],
207    address_roots: &[[u8; 32]],
208    addresses: &[[u8; 32]],
209    compressed_proof: &CompressedProof,
210) -> anchor_lang::Result<()> {
211    if !addresses.is_empty() && !leaves.is_empty() {
212        verify_create_addresses_and_merkle_proof_zkp(
213            roots,
214            leaves,
215            address_roots,
216            addresses,
217            compressed_proof,
218        )
219        .map_err(ProgramError::from)?;
220    } else if !addresses.is_empty() {
221        verify_create_addresses_zkp(address_roots, addresses, compressed_proof)
222            .map_err(ProgramError::from)?;
223    } else {
224        verify_merkle_proof_zkp(roots, leaves, compressed_proof).map_err(ProgramError::from)?;
225    }
226    Ok(())
227}
light_system_program/invoke/verify_state_proof.rs

light_system_program/invoke/
verify_state_proof.rs
