light_program_test/accounts/
state_tree.rs

1use account_compression::{
2    instruction::InitializeStateMerkleTreeAndNullifierQueue, NullifierQueueConfig,
3    StateMerkleTreeConfig,
4};
5use anchor_lang::{InstructionData, ToAccountMetas};
6use light_client::rpc::{errors::RpcError, Rpc};
7use light_compressed_account::instruction_data::insert_into_queues::InsertIntoQueuesInstructionDataMut;
8use light_registry::protocol_config::state::ProtocolConfig;
9use solana_sdk::{
10    instruction::{AccountMeta, Instruction},
11    pubkey::Pubkey,
12    signature::{Keypair, Signature, Signer},
13    transaction::Transaction,
14};
15
16use crate::utils::create_account::create_account_instruction;
17
18#[allow(clippy::too_many_arguments)]
19pub fn create_initialize_merkle_tree_instruction(
20    payer: Pubkey,
21    registered_program_pda: Option<Pubkey>,
22    merkle_tree_pubkey: Pubkey,
23    nullifier_queue_pubkey: Pubkey,
24    state_merkle_tree_config: StateMerkleTreeConfig,
25    nullifier_queue_config: NullifierQueueConfig,
26    program_owner: Option<Pubkey>,
27    forester: Option<Pubkey>,
28    index: u64,
29) -> Instruction {
30    let instruction_data = InitializeStateMerkleTreeAndNullifierQueue {
31        index,
32        program_owner,
33        forester,
34        state_merkle_tree_config,
35        nullifier_queue_config,
36        additional_bytes: 0,
37    };
38    let registered_program = match registered_program_pda {
39        Some(registered_program_pda) => AccountMeta::new(registered_program_pda, false),
40        None => AccountMeta::new(account_compression::ID, false),
41    };
42    Instruction {
43        program_id: account_compression::ID,
44        accounts: vec![
45            AccountMeta::new(payer, true),
46            AccountMeta::new(merkle_tree_pubkey, false),
47            AccountMeta::new(nullifier_queue_pubkey, false),
48            registered_program,
49        ],
50        data: instruction_data.data(),
51    }
52}
53
54pub fn create_insert_leaves_instruction(
55    leaves: Vec<(u8, [u8; 32])>,
56    authority: Pubkey,
57    merkle_tree_pubkeys: Vec<Pubkey>,
58) -> Instruction {
59    let mut bytes = vec![
60        0u8;
61        InsertIntoQueuesInstructionDataMut::required_size_for_capacity(
62            leaves.len() as u8,
63            0,
64            0,
65            merkle_tree_pubkeys.len() as u8,
66            0,
67            0,
68        )
69    ];
70    let (mut ix_data, _) = InsertIntoQueuesInstructionDataMut::new_at(
71        &mut bytes,
72        leaves.len() as u8,
73        0,
74        0,
75        merkle_tree_pubkeys.len() as u8,
76        0,
77        0,
78    )
79    .unwrap();
80    ix_data.num_output_queues = merkle_tree_pubkeys.len() as u8;
81    for (i, (index, leaf)) in leaves.iter().enumerate() {
82        ix_data.leaves[i].leaf = *leaf;
83        ix_data.leaves[i].account_index = *index;
84    }
85
86    let instruction_data = account_compression::instruction::InsertIntoQueues { bytes };
87
88    let accounts = account_compression::accounts::GenericInstruction { authority };
89    let merkle_tree_account_metas = merkle_tree_pubkeys
90        .iter()
91        .map(|pubkey| AccountMeta::new(*pubkey, false))
92        .collect::<Vec<AccountMeta>>();
93
94    Instruction {
95        program_id: account_compression::ID,
96        accounts: [
97            accounts.to_account_metas(Some(true)),
98            merkle_tree_account_metas,
99        ]
100        .concat(),
101        data: instruction_data.data(),
102    }
103}
104
105#[allow(clippy::too_many_arguments)]
106pub async fn create_state_merkle_tree_and_queue_account<R: Rpc>(
107    payer: &Keypair,
108    registry: bool,
109    rpc: &mut R,
110    merkle_tree_keypair: &Keypair,
111    nullifier_queue_keypair: &Keypair,
112    cpi_context_keypair: Option<&Keypair>,
113    program_owner: Option<Pubkey>,
114    forester: Option<Pubkey>,
115    index: u64,
116    merkle_tree_config: &StateMerkleTreeConfig,
117    queue_config: &NullifierQueueConfig,
118) -> Result<Signature, RpcError> {
119    use light_registry::account_compression_cpi::sdk::create_initialize_merkle_tree_instruction as create_initialize_merkle_tree_instruction_registry;
120    let size = account_compression::state::StateMerkleTreeAccount::size(
121        merkle_tree_config.height as usize,
122        merkle_tree_config.changelog_size as usize,
123        merkle_tree_config.roots_size as usize,
124        merkle_tree_config.canopy_depth as usize,
125    );
126
127    let merkle_tree_account_create_ix = create_account_instruction(
128        &payer.pubkey(),
129        size,
130        rpc.get_minimum_balance_for_rent_exemption(size)
131            .await
132            .unwrap(),
133        &account_compression::ID,
134        Some(merkle_tree_keypair),
135    );
136    let size =
137        account_compression::state::queue::QueueAccount::size(queue_config.capacity as usize)
138            .unwrap();
139    let nullifier_queue_account_create_ix = create_account_instruction(
140        &payer.pubkey(),
141        size,
142        rpc.get_minimum_balance_for_rent_exemption(size)
143            .await
144            .unwrap(),
145        &account_compression::ID,
146        Some(nullifier_queue_keypair),
147    );
148
149    let transaction = if registry {
150        let cpi_context_keypair = cpi_context_keypair.unwrap();
151        let rent_cpi_config = rpc
152            .get_minimum_balance_for_rent_exemption(
153                ProtocolConfig::default().cpi_context_size as usize,
154            )
155            .await
156            .unwrap();
157        let create_cpi_context_instruction = create_account_instruction(
158            &payer.pubkey(),
159            ProtocolConfig::default().cpi_context_size as usize,
160            rent_cpi_config,
161            &Pubkey::from(light_sdk::constants::LIGHT_SYSTEM_PROGRAM_ID),
162            Some(cpi_context_keypair),
163        );
164
165        let instruction = create_initialize_merkle_tree_instruction_registry(
166            payer.pubkey(),
167            merkle_tree_keypair.pubkey(),
168            nullifier_queue_keypair.pubkey(),
169            cpi_context_keypair.pubkey(),
170            merkle_tree_config.clone(),
171            queue_config.clone(),
172            program_owner,
173            forester,
174        );
175        Transaction::new_signed_with_payer(
176            &[
177                create_cpi_context_instruction,
178                merkle_tree_account_create_ix,
179                nullifier_queue_account_create_ix,
180                instruction,
181            ],
182            Some(&payer.pubkey()),
183            &vec![
184                payer,
185                merkle_tree_keypair,
186                nullifier_queue_keypair,
187                cpi_context_keypair,
188            ],
189            rpc.get_latest_blockhash().await?.0,
190        )
191    } else {
192        let instruction = create_initialize_merkle_tree_instruction(
193            payer.pubkey(),
194            None,
195            merkle_tree_keypair.pubkey(),
196            nullifier_queue_keypair.pubkey(),
197            merkle_tree_config.clone(),
198            queue_config.clone(),
199            program_owner,
200            forester,
201            index,
202        );
203        Transaction::new_signed_with_payer(
204            &[
205                merkle_tree_account_create_ix,
206                nullifier_queue_account_create_ix,
207                instruction,
208            ],
209            Some(&payer.pubkey()),
210            &vec![payer, merkle_tree_keypair, nullifier_queue_keypair],
211            rpc.get_latest_blockhash().await?.0,
212        )
213    };
214
215    rpc.process_transaction(transaction).await
216}