1use forester_utils::indexer::Indexer;
2use light_hasher::Poseidon;
3use light_system_program::sdk::event::PublicTransactionEvent;
4use light_system_program::{
5 sdk::{
6 address::derive_address,
7 compressed_account::{
8 CompressedAccount, CompressedAccountWithMerkleContext, MerkleContext,
9 },
10 invoke::{create_invoke_instruction, get_sol_pool_pda},
11 },
12 NewAddressParams,
13};
14use solana_sdk::signature::Signature;
15use solana_sdk::{
16 pubkey::Pubkey,
17 signature::{Keypair, Signer},
18};
19
20use crate::assert_compressed_tx::{
21 assert_compressed_transaction, get_merkle_tree_snapshots, AssertCompressedTransactionInputs,
22};
23use light_client::rpc::errors::RpcError;
24use light_client::rpc::RpcConnection;
25use light_client::transaction_params::TransactionParams;
26
27#[allow(clippy::too_many_arguments)]
28pub async fn create_addresses_test<R: RpcConnection, I: Indexer<R>>(
29 rpc: &mut R,
30 test_indexer: &mut I,
31 address_merkle_tree_pubkeys: &[Pubkey],
32 address_merkle_tree_queue_pubkeys: &[Pubkey],
33 mut output_merkle_tree_pubkeys: Vec<Pubkey>,
34 address_seeds: &[[u8; 32]],
35 input_compressed_accounts: &[CompressedAccountWithMerkleContext],
36 create_out_compressed_accounts_for_input_compressed_accounts: bool,
37 transaction_params: Option<TransactionParams>,
38) -> Result<(), RpcError> {
39 if address_merkle_tree_pubkeys.len() != address_seeds.len() {
40 panic!("address_merkle_tree_pubkeys and address_seeds length mismatch for create_addresses_test");
41 }
42 let mut derived_addresses = Vec::new();
43 for (i, address_seed) in address_seeds.iter().enumerate() {
44 let derived_address =
45 derive_address(&address_merkle_tree_pubkeys[i], address_seed).unwrap();
46 println!("derived_address: {:?}", derived_address);
47 derived_addresses.push(derived_address);
48 }
49 let mut address_params = Vec::new();
50
51 for (i, seed) in address_seeds.iter().enumerate() {
52 let new_address_params = NewAddressParams {
53 address_queue_pubkey: address_merkle_tree_queue_pubkeys[i],
54 address_merkle_tree_pubkey: address_merkle_tree_pubkeys[i],
55 seed: *seed,
56 address_merkle_tree_root_index: 0,
57 };
58 address_params.push(new_address_params);
59 }
60
61 let mut output_compressed_accounts = Vec::new();
62 for address in derived_addresses.iter() {
63 output_compressed_accounts.push(CompressedAccount {
64 lamports: 0,
65 owner: rpc.get_payer().pubkey(),
66 data: None,
67 address: Some(*address),
68 });
69 }
70
71 if create_out_compressed_accounts_for_input_compressed_accounts {
72 for compressed_account in input_compressed_accounts.iter() {
73 output_compressed_accounts.push(CompressedAccount {
74 lamports: 0,
75 owner: rpc.get_payer().pubkey(),
76 data: None,
77 address: compressed_account.compressed_account.address,
78 });
79 output_merkle_tree_pubkeys.push(compressed_account.merkle_context.merkle_tree_pubkey);
80 }
81 }
82
83 let payer = rpc.get_payer().insecure_clone();
84
85 let inputs = CompressedTransactionTestInputs {
86 rpc,
87 test_indexer,
88 fee_payer: &payer,
89 authority: &payer,
90 input_compressed_accounts,
91 output_compressed_accounts: output_compressed_accounts.as_slice(),
92 output_merkle_tree_pubkeys: output_merkle_tree_pubkeys.as_slice(),
93 transaction_params,
94 relay_fee: None,
95 compress_or_decompress_lamports: None,
96 is_compress: false,
97 new_address_params: &address_params,
98 sorted_output_accounts: false,
99 created_addresses: Some(derived_addresses.as_slice()),
100 recipient: None,
101 };
102 compressed_transaction_test(inputs).await?;
103 Ok(())
104}
105
106#[allow(clippy::too_many_arguments)]
107pub async fn compress_sol_test<R: RpcConnection, I: Indexer<R>>(
108 rpc: &mut R,
109 test_indexer: &mut I,
110 authority: &Keypair,
111 input_compressed_accounts: &[CompressedAccountWithMerkleContext],
112 create_out_compressed_accounts_for_input_compressed_accounts: bool,
113 compress_amount: u64,
114 output_merkle_tree_pubkey: &Pubkey,
115 transaction_params: Option<TransactionParams>,
116) -> Result<(), RpcError> {
117 let input_lamports = if input_compressed_accounts.is_empty() {
118 0
119 } else {
120 input_compressed_accounts
121 .iter()
122 .map(|x| x.compressed_account.lamports)
123 .sum::<u64>()
124 };
125 let mut output_compressed_accounts = Vec::new();
126 output_compressed_accounts.push(CompressedAccount {
127 lamports: input_lamports + compress_amount,
128 owner: authority.pubkey(),
129 data: None,
130 address: None,
131 });
132 let mut output_merkle_tree_pubkeys = vec![*output_merkle_tree_pubkey];
133 if create_out_compressed_accounts_for_input_compressed_accounts {
134 for compressed_account in input_compressed_accounts.iter() {
135 output_compressed_accounts.push(CompressedAccount {
136 lamports: 0,
137 owner: authority.pubkey(),
138 data: None,
139 address: compressed_account.compressed_account.address,
140 });
141 output_merkle_tree_pubkeys.push(compressed_account.merkle_context.merkle_tree_pubkey);
142 }
143 }
144 let inputs = CompressedTransactionTestInputs {
145 rpc,
146 test_indexer,
147 fee_payer: authority,
148 authority,
149 input_compressed_accounts,
150 output_compressed_accounts: output_compressed_accounts.as_slice(),
151 output_merkle_tree_pubkeys: &[*output_merkle_tree_pubkey],
152 transaction_params,
153 relay_fee: None,
154 compress_or_decompress_lamports: Some(compress_amount),
155 is_compress: true,
156 new_address_params: &[],
157 sorted_output_accounts: false,
158 created_addresses: None,
159 recipient: None,
160 };
161 compressed_transaction_test(inputs).await?;
162 Ok(())
163}
164
165#[allow(clippy::too_many_arguments)]
166pub async fn decompress_sol_test<R: RpcConnection, I: Indexer<R>>(
167 rpc: &mut R,
168 test_indexer: &mut I,
169 authority: &Keypair,
170 input_compressed_accounts: &[CompressedAccountWithMerkleContext],
171 recipient: &Pubkey,
172 decompress_amount: u64,
173 output_merkle_tree_pubkey: &Pubkey,
174 transaction_params: Option<TransactionParams>,
175) -> Result<(), RpcError> {
176 let input_lamports = input_compressed_accounts
177 .iter()
178 .map(|x| x.compressed_account.lamports)
179 .sum::<u64>();
180
181 let output_compressed_accounts = vec![CompressedAccount {
182 lamports: input_lamports - decompress_amount,
183 owner: rpc.get_payer().pubkey(),
184 data: None,
185 address: None,
186 }];
187 let payer = rpc.get_payer().insecure_clone();
188 let inputs = CompressedTransactionTestInputs {
189 rpc,
190 test_indexer,
191 fee_payer: &payer,
192 authority,
193 input_compressed_accounts,
194 output_compressed_accounts: output_compressed_accounts.as_slice(),
195 output_merkle_tree_pubkeys: &[*output_merkle_tree_pubkey],
196 transaction_params,
197 relay_fee: None,
198 compress_or_decompress_lamports: Some(decompress_amount),
199 is_compress: false,
200 new_address_params: &[],
201 sorted_output_accounts: false,
202 created_addresses: None,
203 recipient: Some(*recipient),
204 };
205 compressed_transaction_test(inputs).await?;
206 Ok(())
207}
208
209#[allow(clippy::too_many_arguments)]
210pub async fn transfer_compressed_sol_test<R: RpcConnection, I: Indexer<R>>(
211 rpc: &mut R,
212 test_indexer: &mut I,
213 authority: &Keypair,
214 input_compressed_accounts: &[CompressedAccountWithMerkleContext],
215 recipients: &[Pubkey],
216 output_merkle_tree_pubkeys: &[Pubkey],
217 transaction_params: Option<TransactionParams>,
218) -> Result<Signature, RpcError> {
219 if recipients.len() != output_merkle_tree_pubkeys.len() {
220 panic!("recipients and output_merkle_tree_pubkeys length mismatch for transfer_compressed_sol_test");
221 }
222
223 if input_compressed_accounts.is_empty() {
224 panic!("input_compressed_accounts is empty for transfer_compressed_sol_test");
225 }
226 let input_lamports = input_compressed_accounts
227 .iter()
228 .map(|x| x.compressed_account.lamports)
229 .sum::<u64>();
230 let mut output_compressed_accounts = Vec::new();
231 let mut output_merkle_tree_pubkeys = output_merkle_tree_pubkeys.to_vec();
232 output_merkle_tree_pubkeys.sort();
233 let input_addresses = input_compressed_accounts
234 .iter()
235 .map(|x| x.compressed_account.address)
236 .collect::<Vec<_>>();
237 for (i, _) in output_merkle_tree_pubkeys.iter().enumerate() {
238 let address = if i < input_addresses.len() {
239 input_addresses[i]
240 } else {
241 None
242 };
243 let mut lamports = input_lamports / output_merkle_tree_pubkeys.len() as u64;
244 if i == 0 {
245 lamports += input_lamports % output_merkle_tree_pubkeys.len() as u64;
246 }
247
248 output_compressed_accounts.push(CompressedAccount {
249 lamports,
250 owner: recipients[i],
251 data: None,
252 address,
253 });
254 }
255 let payer = rpc.get_payer().insecure_clone();
256 let inputs = CompressedTransactionTestInputs {
257 rpc,
258 test_indexer,
259 fee_payer: &payer,
260 authority,
261 input_compressed_accounts,
262 output_compressed_accounts: output_compressed_accounts.as_slice(),
263 output_merkle_tree_pubkeys: output_merkle_tree_pubkeys.as_slice(),
264 transaction_params,
265 relay_fee: None,
266 compress_or_decompress_lamports: None,
267 is_compress: false,
268 new_address_params: &[],
269 sorted_output_accounts: false,
270 created_addresses: None,
271 recipient: None,
272 };
273 compressed_transaction_test(inputs).await
274}
275
276pub struct CompressedTransactionTestInputs<'a, R: RpcConnection, I: Indexer<R>> {
277 rpc: &'a mut R,
278 test_indexer: &'a mut I,
279 fee_payer: &'a Keypair,
280 authority: &'a Keypair,
281 input_compressed_accounts: &'a [CompressedAccountWithMerkleContext],
282 output_compressed_accounts: &'a [CompressedAccount],
283 output_merkle_tree_pubkeys: &'a [Pubkey],
284 transaction_params: Option<TransactionParams>,
285 relay_fee: Option<u64>,
286 compress_or_decompress_lamports: Option<u64>,
287 is_compress: bool,
288 new_address_params: &'a [NewAddressParams],
289 sorted_output_accounts: bool,
290 created_addresses: Option<&'a [[u8; 32]]>,
291 recipient: Option<Pubkey>,
292}
293
294#[allow(clippy::too_many_arguments)]
295pub async fn compressed_transaction_test<R: RpcConnection, I: Indexer<R>>(
296 inputs: CompressedTransactionTestInputs<'_, R, I>,
297) -> Result<Signature, RpcError> {
298 let mut compressed_account_hashes = Vec::new();
299
300 let compressed_account_input_hashes = if !inputs.input_compressed_accounts.is_empty() {
301 for compressed_account in inputs.input_compressed_accounts.iter() {
302 compressed_account_hashes.push(
303 compressed_account
304 .compressed_account
305 .hash::<Poseidon>(
306 &compressed_account.merkle_context.merkle_tree_pubkey,
307 &compressed_account.merkle_context.leaf_index,
308 )
309 .unwrap(),
310 );
311 }
312 Some(compressed_account_hashes.as_slice())
313 } else {
314 None
315 };
316 let state_input_merkle_trees = inputs
317 .input_compressed_accounts
318 .iter()
319 .map(|x| x.merkle_context.merkle_tree_pubkey)
320 .collect::<Vec<Pubkey>>();
321 let state_input_merkle_trees = if state_input_merkle_trees.is_empty() {
322 None
323 } else {
324 Some(state_input_merkle_trees.as_slice())
325 };
326 let mut root_indices = Vec::new();
327 let mut proof = None;
328 let mut input_merkle_tree_snapshots = Vec::new();
329 let mut address_params = Vec::new();
330 if !inputs.input_compressed_accounts.is_empty() || !inputs.new_address_params.is_empty() {
331 let address_merkle_tree_pubkeys = if inputs.new_address_params.is_empty() {
332 None
333 } else {
334 Some(
335 inputs
336 .new_address_params
337 .iter()
338 .map(|x| x.address_merkle_tree_pubkey)
339 .collect::<Vec<_>>(),
340 )
341 };
342 let proof_rpc_res = inputs
343 .test_indexer
344 .create_proof_for_compressed_accounts(
345 compressed_account_input_hashes,
346 state_input_merkle_trees,
347 inputs.created_addresses,
348 address_merkle_tree_pubkeys,
349 inputs.rpc,
350 )
351 .await;
352 root_indices = proof_rpc_res.root_indices;
353 proof = Some(proof_rpc_res.proof);
354 let input_merkle_tree_accounts = inputs
355 .test_indexer
356 .get_state_merkle_tree_accounts(state_input_merkle_trees.unwrap_or(&[]));
357 input_merkle_tree_snapshots =
358 get_merkle_tree_snapshots::<R>(inputs.rpc, input_merkle_tree_accounts.as_slice()).await;
359
360 if !inputs.new_address_params.is_empty() {
361 for (i, input_address_params) in inputs.new_address_params.iter().enumerate() {
362 address_params.push(input_address_params.clone());
363 address_params[i].address_merkle_tree_root_index =
364 proof_rpc_res.address_root_indices[i];
365 }
366 }
367 }
368
369 let output_merkle_tree_accounts = inputs
370 .test_indexer
371 .get_state_merkle_tree_accounts(inputs.output_merkle_tree_pubkeys);
372 let output_merkle_tree_snapshots =
373 get_merkle_tree_snapshots::<R>(inputs.rpc, output_merkle_tree_accounts.as_slice()).await;
374 let instruction = create_invoke_instruction(
375 &inputs.fee_payer.pubkey(),
376 &inputs.authority.pubkey().clone(),
377 inputs
378 .input_compressed_accounts
379 .iter()
380 .map(|x| x.compressed_account.clone())
381 .collect::<Vec<CompressedAccount>>()
382 .as_slice(),
383 inputs.output_compressed_accounts,
384 inputs
385 .input_compressed_accounts
386 .iter()
387 .map(|x| x.merkle_context)
388 .collect::<Vec<MerkleContext>>()
389 .as_slice(),
390 inputs.output_merkle_tree_pubkeys,
391 &root_indices,
392 &address_params,
393 proof,
394 inputs.compress_or_decompress_lamports,
395 inputs.is_compress,
396 inputs.recipient,
397 true,
398 );
399 let mut recipient_balance_pre = 0;
400 let mut compressed_sol_pda_balance_pre = 0;
401 if inputs.compress_or_decompress_lamports.is_some() {
402 compressed_sol_pda_balance_pre =
403 match inputs.rpc.get_account(get_sol_pool_pda()).await.unwrap() {
404 Some(account) => account.lamports,
405 None => 0,
406 };
407 }
408 if inputs.recipient.is_some() {
409 recipient_balance_pre = match inputs
411 .rpc
412 .get_account(inputs.recipient.unwrap())
413 .await
414 .unwrap()
415 {
416 Some(account) => account.lamports,
417 None => 0,
418 };
419 }
420 let event = inputs
421 .rpc
422 .create_and_send_transaction_with_event::<PublicTransactionEvent>(
423 &[instruction],
424 &inputs.fee_payer.pubkey(),
425 &[inputs.fee_payer, inputs.authority],
426 inputs.transaction_params,
427 )
428 .await?
429 .unwrap();
430
431 let (created_output_compressed_accounts, _) = inputs
432 .test_indexer
433 .add_event_and_compressed_accounts(&event.0);
434 let input = AssertCompressedTransactionInputs {
435 rpc: inputs.rpc,
436 test_indexer: inputs.test_indexer,
437 output_compressed_accounts: inputs.output_compressed_accounts,
438 created_output_compressed_accounts: created_output_compressed_accounts.as_slice(),
439 event: &event.0,
440 input_merkle_tree_snapshots: input_merkle_tree_snapshots.as_slice(),
441 output_merkle_tree_snapshots: output_merkle_tree_snapshots.as_slice(),
442 recipient_balance_pre,
443 compress_or_decompress_lamports: inputs.compress_or_decompress_lamports,
444 is_compress: inputs.is_compress,
445 compressed_sol_pda_balance_pre,
446 compression_recipient: inputs.recipient,
447 created_addresses: inputs.created_addresses.unwrap_or(&[]),
448 sorted_output_accounts: inputs.sorted_output_accounts,
449 relay_fee: inputs.relay_fee,
450 input_compressed_account_hashes: &compressed_account_hashes,
451 address_queue_pubkeys: &inputs
452 .new_address_params
453 .iter()
454 .map(|x| x.address_queue_pubkey)
455 .collect::<Vec<Pubkey>>(),
456 };
457 assert_compressed_transaction(input).await;
458 Ok(event.1)
459}