light_program_test/accounts/
state_tree.rs

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