1use account_compression::{
2 AddressMerkleTreeConfig, AddressQueueConfig, NullifierQueueConfig, QueueAccount,
3 StateMerkleTreeConfig,
4};
5use light_client::{
6 indexer::{AddressMerkleTreeAccounts, StateMerkleTreeAccounts},
7 rpc::{Rpc, RpcError},
8};
9use light_registry::{
10 account_compression_cpi::sdk::{
11 create_rollover_state_merkle_tree_instruction, CreateRolloverMerkleTreeInstructionInputs,
12 },
13 protocol_config::state::ProtocolConfig,
14 sdk::{create_register_forester_instruction, create_update_forester_pda_instruction},
15 utils::get_forester_pda,
16 ForesterConfig, ForesterPda,
17};
18use light_sdk;
19use solana_sdk::{
20 instruction::Instruction,
21 pubkey::Pubkey,
22 signature::{Keypair, Signer},
23};
24
25use crate::{
26 address_merkle_tree_config::{get_address_bundle_config, get_state_bundle_config},
27 instructions::create_account::create_account_instruction,
28};
29
30pub async fn register_test_forester<R: Rpc>(
32 rpc: &mut R,
33 governance_authority: &Keypair,
34 forester_authority: &Pubkey,
35 config: ForesterConfig,
36) -> Result<(), RpcError> {
37 let ix = create_register_forester_instruction(
38 &governance_authority.pubkey(),
39 &governance_authority.pubkey(),
40 forester_authority,
41 config,
42 );
43 rpc.create_and_send_transaction(
44 &[ix],
45 &governance_authority.pubkey(),
46 &[governance_authority],
47 )
48 .await?;
49 assert_registered_forester(
50 rpc,
51 forester_authority,
52 ForesterPda {
53 authority: *forester_authority,
54 config,
55 active_weight: 1,
56 ..Default::default()
57 },
58 )
59 .await
60}
61
62pub async fn update_test_forester<R: Rpc>(
63 rpc: &mut R,
64 forester_authority: &Keypair,
65 derivation_key: &Pubkey,
66 new_forester_authority: Option<&Keypair>,
67 config: ForesterConfig,
68) -> Result<(), RpcError> {
69 let mut pre_account_state = rpc
70 .get_anchor_account::<ForesterPda>(&get_forester_pda(derivation_key).0)
71 .await?
72 .unwrap();
73 let (signers, new_forester_authority) = if let Some(new_authority) = new_forester_authority {
74 pre_account_state.authority = new_authority.pubkey();
75
76 (
77 vec![forester_authority, &new_authority],
78 Some(new_authority.pubkey()),
79 )
80 } else {
81 (vec![forester_authority], None)
82 };
83 let ix = create_update_forester_pda_instruction(
84 &forester_authority.pubkey(),
85 derivation_key,
86 new_forester_authority,
87 Some(config),
88 );
89
90 rpc.create_and_send_transaction(&[ix], &forester_authority.pubkey(), &signers)
91 .await?;
92
93 pre_account_state.config = config;
94 assert_registered_forester(rpc, derivation_key, pre_account_state).await
95}
96
97pub async fn assert_registered_forester<R: Rpc>(
98 rpc: &mut R,
99 forester: &Pubkey,
100 expected_account: ForesterPda,
101) -> Result<(), RpcError> {
102 let pda = get_forester_pda(forester).0;
103 let account_data = rpc.get_anchor_account::<ForesterPda>(&pda).await?.unwrap();
104 if account_data != expected_account {
105 return Err(RpcError::AssertRpcError(format!(
106 "Expected account data: {:?}, got: {:?}",
107 expected_account, account_data
108 )));
109 }
110 Ok(())
111}
112
113pub struct RentExemption {
114 pub size: usize,
115 pub lamports: u64,
116}
117
118pub async fn get_rent_exemption_for_address_merkle_tree_and_queue<R: Rpc>(
119 rpc: &mut R,
120 address_merkle_tree_config: &AddressMerkleTreeConfig,
121 address_queue_config: &AddressQueueConfig,
122) -> (RentExemption, RentExemption) {
123 let queue_size = QueueAccount::size(address_queue_config.capacity as usize).unwrap();
124
125 let queue_rent_exempt_lamports = rpc
126 .get_minimum_balance_for_rent_exemption(queue_size)
127 .await
128 .unwrap();
129 let tree_size = account_compression::state::AddressMerkleTreeAccount::size(
130 address_merkle_tree_config.height as usize,
131 address_merkle_tree_config.changelog_size as usize,
132 address_merkle_tree_config.roots_size as usize,
133 address_merkle_tree_config.canopy_depth as usize,
134 address_merkle_tree_config.address_changelog_size as usize,
135 );
136 let merkle_tree_rent_exempt_lamports = rpc
137 .get_minimum_balance_for_rent_exemption(tree_size)
138 .await
139 .unwrap();
140 (
141 RentExemption {
142 lamports: merkle_tree_rent_exempt_lamports,
143 size: tree_size,
144 },
145 RentExemption {
146 lamports: queue_rent_exempt_lamports,
147 size: queue_size,
148 },
149 )
150}
151
152pub async fn get_rent_exemption_for_state_merkle_tree_and_queue<R: Rpc>(
153 rpc: &mut R,
154 merkle_tree_config: &StateMerkleTreeConfig,
155 queue_config: &NullifierQueueConfig,
156) -> (RentExemption, RentExemption) {
157 let queue_size = QueueAccount::size(queue_config.capacity as usize).unwrap();
158
159 let queue_rent_exempt_lamports = rpc
160 .get_minimum_balance_for_rent_exemption(queue_size)
161 .await
162 .unwrap();
163 let tree_size = account_compression::state::StateMerkleTreeAccount::size(
164 merkle_tree_config.height as usize,
165 merkle_tree_config.changelog_size as usize,
166 merkle_tree_config.roots_size as usize,
167 merkle_tree_config.canopy_depth as usize,
168 );
169 let merkle_tree_rent_exempt_lamports = rpc
170 .get_minimum_balance_for_rent_exemption(tree_size)
171 .await
172 .unwrap();
173 (
174 RentExemption {
175 lamports: merkle_tree_rent_exempt_lamports,
176 size: tree_size,
177 },
178 RentExemption {
179 lamports: queue_rent_exempt_lamports,
180 size: queue_size,
181 },
182 )
183}
184
185#[allow(clippy::too_many_arguments)]
186pub async fn create_rollover_address_merkle_tree_instructions<R: Rpc>(
187 rpc: &mut R,
188 authority: &Pubkey,
189 derivation: &Pubkey,
190 new_nullifier_queue_keypair: &Keypair,
191 new_address_merkle_tree_keypair: &Keypair,
192 merkle_tree_pubkey: &Pubkey,
193 nullifier_queue_pubkey: &Pubkey,
194 epoch: u64,
195 is_metadata_forester: bool,
196) -> Vec<Instruction> {
197 let (merkle_tree_config, queue_config) = get_address_bundle_config(
198 rpc,
199 AddressMerkleTreeAccounts {
200 merkle_tree: *merkle_tree_pubkey,
201 queue: *nullifier_queue_pubkey,
202 },
203 )
204 .await;
205 let (merkle_tree_rent_exemption, queue_rent_exemption) =
206 get_rent_exemption_for_address_merkle_tree_and_queue(
207 rpc,
208 &merkle_tree_config,
209 &queue_config,
210 )
211 .await;
212 let create_nullifier_queue_instruction = create_account_instruction(
213 authority,
214 queue_rent_exemption.size,
215 queue_rent_exemption.lamports,
216 &account_compression::ID,
217 Some(new_nullifier_queue_keypair),
218 );
219 let create_state_merkle_tree_instruction = create_account_instruction(
220 authority,
221 merkle_tree_rent_exemption.size,
222 merkle_tree_rent_exemption.lamports,
223 &account_compression::ID,
224 Some(new_address_merkle_tree_keypair),
225 );
226 let instruction = light_registry::account_compression_cpi::sdk::create_rollover_address_merkle_tree_instruction(
227 CreateRolloverMerkleTreeInstructionInputs {
228 authority: *authority,
229 derivation: *derivation,
230 new_queue: new_nullifier_queue_keypair.pubkey(),
231 new_merkle_tree: new_address_merkle_tree_keypair.pubkey(),
232 old_queue: *nullifier_queue_pubkey,
233 old_merkle_tree: *merkle_tree_pubkey,
234 cpi_context_account: None,
235 is_metadata_forester,
236 },epoch
237 );
238 vec![
239 create_nullifier_queue_instruction,
240 create_state_merkle_tree_instruction,
241 instruction,
242 ]
243}
244
245#[allow(clippy::too_many_arguments)]
246pub async fn perform_state_merkle_tree_roll_over<R: Rpc>(
247 rpc: &mut R,
248 authority: &Keypair,
249 derivation: &Pubkey,
250 new_nullifier_queue_keypair: &Keypair,
251 new_state_merkle_tree_keypair: &Keypair,
252 merkle_tree_pubkey: &Pubkey,
253 nullifier_queue_pubkey: &Pubkey,
254 epoch: u64,
255 is_metadata_forester: bool,
256) -> Result<(), RpcError> {
257 let instructions = create_rollover_address_merkle_tree_instructions(
258 rpc,
259 &authority.pubkey(),
260 derivation,
261 new_nullifier_queue_keypair,
262 new_state_merkle_tree_keypair,
263 merkle_tree_pubkey,
264 nullifier_queue_pubkey,
265 epoch,
266 is_metadata_forester,
267 )
268 .await;
269 rpc.create_and_send_transaction(
270 &instructions,
271 &authority.pubkey(),
272 &[
273 authority,
274 new_nullifier_queue_keypair,
275 new_state_merkle_tree_keypair,
276 ],
277 )
278 .await?;
279 Ok(())
280}
281#[allow(clippy::too_many_arguments)]
282pub async fn create_rollover_state_merkle_tree_instructions<R: Rpc>(
283 rpc: &mut R,
284 authority: &Pubkey,
285 derivation: &Pubkey,
286 new_nullifier_queue_keypair: &Keypair,
287 new_state_merkle_tree_keypair: &Keypair,
288 new_cpi_context_keypair: &Keypair,
289 merkle_tree_pubkey: &Pubkey,
290 nullifier_queue_pubkey: &Pubkey,
291 epoch: u64,
292 is_metadata_forester: bool,
293) -> Vec<Instruction> {
294 let (merkle_tree_config, queue_config) = get_state_bundle_config(
295 rpc,
296 StateMerkleTreeAccounts {
297 merkle_tree: *merkle_tree_pubkey,
298 nullifier_queue: *nullifier_queue_pubkey,
299 cpi_context: new_cpi_context_keypair.pubkey(),
300 },
301 )
302 .await;
303 let (state_merkle_tree_rent_exemption, queue_rent_exemption) =
304 get_rent_exemption_for_state_merkle_tree_and_queue(rpc, &merkle_tree_config, &queue_config)
305 .await;
306 let create_nullifier_queue_instruction = create_account_instruction(
307 authority,
308 queue_rent_exemption.size,
309 queue_rent_exemption.lamports,
310 &account_compression::ID,
311 Some(new_nullifier_queue_keypair),
312 );
313 let create_state_merkle_tree_instruction = create_account_instruction(
314 authority,
315 state_merkle_tree_rent_exemption.size,
316 state_merkle_tree_rent_exemption.lamports,
317 &account_compression::ID,
318 Some(new_state_merkle_tree_keypair),
319 );
320 let account_size: usize = ProtocolConfig::default().cpi_context_size as usize;
321 let create_cpi_context_account_instruction = create_account_instruction(
322 authority,
323 account_size,
324 rpc.get_minimum_balance_for_rent_exemption(account_size)
325 .await
326 .unwrap(),
327 &Pubkey::from(light_sdk::constants::LIGHT_SYSTEM_PROGRAM_ID),
328 Some(new_cpi_context_keypair),
329 );
330 let instruction = create_rollover_state_merkle_tree_instruction(
331 CreateRolloverMerkleTreeInstructionInputs {
332 authority: *authority,
333 derivation: *derivation,
334 new_queue: new_nullifier_queue_keypair.pubkey(),
335 new_merkle_tree: new_state_merkle_tree_keypair.pubkey(),
336 old_queue: *nullifier_queue_pubkey,
337 old_merkle_tree: *merkle_tree_pubkey,
338 cpi_context_account: Some(new_cpi_context_keypair.pubkey()),
339 is_metadata_forester,
340 },
341 epoch,
342 );
343 vec![
344 create_nullifier_queue_instruction,
345 create_state_merkle_tree_instruction,
346 create_cpi_context_account_instruction,
347 instruction,
348 ]
349}