1use std::{fmt::Debug, time::Duration};
2
3use account_compression::{
4 AddressMerkleTreeConfig, AddressQueueConfig, NullifierQueueConfig, StateMerkleTreeConfig,
5};
6use async_trait::async_trait;
7use borsh::BorshDeserialize;
8use light_batched_merkle_tree::{
9 constants::{
10 DEFAULT_BATCH_ADDRESS_TREE_HEIGHT, DEFAULT_BATCH_ROOT_HISTORY_LEN,
11 DEFAULT_BATCH_STATE_TREE_HEIGHT,
12 },
13 merkle_tree::BatchedMerkleTreeAccount,
14};
15use light_client::{
16 fee::FeeConfig,
17 indexer::{
18 AccountProofInputs, Address, AddressMerkleTreeAccounts, AddressProofInputs,
19 AddressWithTree, BatchAddressUpdateIndexerResponse, CompressedAccount,
20 CompressedTokenAccount, Context, GetCompressedAccountsByOwnerConfig,
21 GetCompressedTokenAccountsByOwnerOrDelegateOptions, Indexer, IndexerError,
22 IndexerRpcConfig, Items, ItemsWithCursor, MerkleProof, MerkleProofWithContext,
23 NewAddressProofWithContext, OwnerBalance, PaginatedOptions, QueueElementsResult, Response,
24 RetryConfig, RootIndex, SignatureWithMetadata, StateMerkleTreeAccounts, TokenBalance,
25 ValidityProofWithContext,
26 },
27 rpc::{Rpc, RpcError},
28};
29use light_compressed_account::{
30 compressed_account::{CompressedAccountWithMerkleContext, MerkleContext},
31 hash_chain::create_hash_chain_from_slice,
32 indexer_event::event::PublicTransactionEvent,
33 instruction_data::compressed_proof::CompressedProof,
34 tx_hash::create_tx_hash,
35 TreeType,
36};
37use light_hasher::{bigint::bigint_to_be_bytes_array, Poseidon};
38use light_merkle_tree_metadata::QueueType;
39use light_merkle_tree_reference::MerkleTree;
40use light_prover_client::{
41 constants::{PROVE_PATH, SERVER_ADDRESS},
42 helpers::{big_int_to_string, bigint_to_u8_32, string_to_big_int},
43 proof::{compress_proof, deserialize_gnark_proof_json, proof_from_json_struct},
44 proof_type::ProofType,
45 proof_types::{
46 combined::{v1::CombinedJsonStruct as CombinedJsonStructLegacy, v2::CombinedJsonStruct},
47 inclusion::{
48 v1::{
49 BatchInclusionJsonStruct as BatchInclusionJsonStructLegacy,
50 InclusionProofInputs as InclusionProofInputsLegacy,
51 },
52 v2::{BatchInclusionJsonStruct, InclusionMerkleProofInputs, InclusionProofInputs},
53 },
54 non_inclusion::{
55 v1::{
56 BatchNonInclusionJsonStruct as BatchNonInclusionJsonStructLegacy,
57 NonInclusionProofInputs as NonInclusionProofInputsLegacy,
58 },
59 v2::{BatchNonInclusionJsonStruct, NonInclusionProofInputs},
60 },
61 },
62};
63use light_sdk::{
64 light_hasher::Hash,
65 token::{TokenData, TokenDataWithMerkleContext},
66};
67use log::info;
68use num_bigint::{BigInt, BigUint};
69use num_traits::FromBytes;
70use reqwest::Client;
71use solana_sdk::{
72 bs58,
73 pubkey::Pubkey,
74 signature::{Keypair, Signer},
75};
76
77use super::{
78 address_tree::{AddressMerkleTreeBundle, IndexedMerkleTreeVersion},
79 state_tree::{LeafIndexInfo, StateMerkleTreeBundle},
80};
81#[cfg(feature = "devenv")]
82use crate::accounts::{
83 address_tree_v2::create_batch_address_merkle_tree,
84 state_tree_v2::create_batched_state_merkle_tree,
85};
86use crate::{
87 accounts::{
88 address_tree::create_address_merkle_tree_and_queue_account,
89 state_tree::create_state_merkle_tree_and_queue_account, test_accounts::TestAccounts,
90 },
91 indexer::TestIndexerExtensions,
92};
93
94#[derive(Debug)]
95pub struct TestIndexer {
96 pub state_merkle_trees: Vec<StateMerkleTreeBundle>,
97 pub address_merkle_trees: Vec<AddressMerkleTreeBundle>,
98 pub payer: Keypair,
99 pub group_pda: Pubkey,
100 pub compressed_accounts: Vec<CompressedAccountWithMerkleContext>,
101 pub nullified_compressed_accounts: Vec<CompressedAccountWithMerkleContext>,
102 pub token_compressed_accounts: Vec<TokenDataWithMerkleContext>,
103 pub token_nullified_compressed_accounts: Vec<TokenDataWithMerkleContext>,
104 pub events: Vec<PublicTransactionEvent>,
105}
106
107impl Clone for TestIndexer {
108 fn clone(&self) -> Self {
109 Self {
110 state_merkle_trees: self.state_merkle_trees.clone(),
111 address_merkle_trees: self.address_merkle_trees.clone(),
112 payer: self.payer.insecure_clone(),
113 group_pda: self.group_pda,
114 compressed_accounts: self.compressed_accounts.clone(),
115 nullified_compressed_accounts: self.nullified_compressed_accounts.clone(),
116 token_compressed_accounts: self.token_compressed_accounts.clone(),
117 token_nullified_compressed_accounts: self.token_nullified_compressed_accounts.clone(),
118 events: self.events.clone(),
119 }
120 }
121}
122
123#[async_trait]
124impl Indexer for TestIndexer {
125 async fn get_indexer_slot(&self, _config: Option<RetryConfig>) -> Result<u64, IndexerError> {
127 Ok(u64::MAX)
129 }
130
131 async fn get_multiple_compressed_account_proofs(
132 &self,
133 hashes: Vec<[u8; 32]>,
134 _config: Option<IndexerRpcConfig>,
135 ) -> Result<Response<Items<MerkleProof>>, IndexerError> {
136 info!("Getting proofs for {:?}", hashes);
137 let mut proofs: Vec<MerkleProof> = Vec::new();
138 hashes.iter().for_each(|hash| {
139 self.state_merkle_trees.iter().for_each(|tree| {
140 if let Some(leaf_index) = tree.merkle_tree.get_leaf_index(hash) {
141 let proof = tree
142 .merkle_tree
143 .get_proof_of_leaf(leaf_index, false)
144 .unwrap();
145 proofs.push(MerkleProof {
146 hash: *hash,
147 leaf_index: leaf_index as u64,
148 merkle_tree: tree.accounts.merkle_tree,
149 proof: proof.to_vec(),
150 root_seq: tree.merkle_tree.sequence_number as u64,
151 root: *tree.merkle_tree.roots.last().unwrap(),
152 });
153 }
154 })
155 });
156 Ok(Response {
157 context: Context {
158 slot: self.get_current_slot(),
159 },
160 value: Items { items: proofs },
161 })
162 }
163
164 async fn get_compressed_accounts_by_owner(
165 &self,
166 owner: &Pubkey,
167 _options: Option<GetCompressedAccountsByOwnerConfig>,
168 _config: Option<IndexerRpcConfig>,
169 ) -> Result<Response<ItemsWithCursor<CompressedAccount>>, IndexerError> {
170 let accounts_with_context = <TestIndexer as TestIndexerExtensions>::get_compressed_accounts_with_merkle_context_by_owner(self, owner);
171 let accounts: Result<Vec<CompressedAccount>, IndexerError> = accounts_with_context
172 .into_iter()
173 .map(|acc| acc.try_into())
174 .collect();
175
176 Ok(Response {
177 context: Context {
178 slot: self.get_current_slot(),
179 },
180 value: ItemsWithCursor {
181 items: accounts?,
182 cursor: None,
183 },
184 })
185 }
186
187 async fn get_compressed_account(
188 &self,
189 address: Address,
190 _config: Option<IndexerRpcConfig>,
191 ) -> Result<Response<Option<CompressedAccount>>, IndexerError> {
192 let account = self
193 .compressed_accounts
194 .iter()
195 .find(|acc| acc.compressed_account.address == Some(address));
196
197 let account_data = match account {
198 Some(acc) => Some(acc.clone().try_into()?),
199 None => None,
200 };
201
202 Ok(Response {
203 context: Context {
204 slot: self.get_current_slot(),
205 },
206 value: account_data,
207 })
208 }
209
210 async fn get_compressed_account_by_hash(
211 &self,
212 hash: Hash,
213 _config: Option<IndexerRpcConfig>,
214 ) -> Result<Response<Option<CompressedAccount>>, IndexerError> {
215 let res = self
216 .compressed_accounts
217 .iter()
218 .find(|acc| acc.hash() == Ok(hash));
219
220 let account = if res.is_none() {
222 let res = self
223 .token_compressed_accounts
224 .iter()
225 .find(|acc| acc.compressed_account.hash() == Ok(hash));
226 res.map(|x| &x.compressed_account)
227 } else {
228 res
229 };
230
231 let account_data = match account {
232 Some(acc) => Some(acc.clone().try_into()?),
233 None => None,
234 };
235
236 Ok(Response {
237 context: Context {
238 slot: self.get_current_slot(),
239 },
240 value: account_data,
241 })
242 }
243
244 async fn get_compressed_token_accounts_by_owner(
245 &self,
246 owner: &Pubkey,
247 options: Option<GetCompressedTokenAccountsByOwnerOrDelegateOptions>,
248 _config: Option<IndexerRpcConfig>,
249 ) -> Result<Response<ItemsWithCursor<CompressedTokenAccount>>, IndexerError> {
250 let mint = options.as_ref().and_then(|opts| opts.mint);
251 let token_accounts: Result<Vec<CompressedTokenAccount>, IndexerError> = self
252 .token_compressed_accounts
253 .iter()
254 .filter(|acc| {
255 acc.token_data.owner == *owner && mint.is_none_or(|m| acc.token_data.mint == m)
256 })
257 .map(|acc| CompressedTokenAccount::try_from(acc.clone()))
258 .collect();
259 let token_accounts = token_accounts?;
260 let token_accounts = if let Some(options) = options {
261 if let Some(limit) = options.limit {
262 token_accounts.into_iter().take(limit as usize).collect()
263 } else {
264 token_accounts
265 }
266 } else {
267 token_accounts
268 };
269
270 Ok(Response {
271 context: Context {
272 slot: self.get_current_slot(),
273 },
274 value: ItemsWithCursor {
275 items: token_accounts,
276 cursor: None,
277 },
278 })
279 }
280
281 async fn get_compressed_balance(
282 &self,
283 address: Option<Address>,
284 hash: Option<Hash>,
285 _config: Option<IndexerRpcConfig>,
286 ) -> Result<Response<u64>, IndexerError> {
287 let account_response = match (address, hash) {
288 (Some(addr), _) => self.get_compressed_account(addr, None).await?,
289 (_, Some(h)) => self.get_compressed_account_by_hash(h, None).await?,
290 _ => {
291 return Err(IndexerError::InvalidParameters(
292 "Either address or hash must be provided".to_string(),
293 ))
294 }
295 };
296 let account = account_response
297 .value
298 .ok_or(IndexerError::AccountNotFound)?;
299 Ok(Response {
300 context: Context {
301 slot: self.get_current_slot(),
302 },
303 value: account.lamports,
304 })
305 }
306
307 async fn get_compressed_token_account_balance(
308 &self,
309 address: Option<Address>,
310 hash: Option<Hash>,
311 _config: Option<IndexerRpcConfig>,
312 ) -> Result<Response<u64>, IndexerError> {
313 let account = match (address, hash) {
314 (Some(address), _) => self
315 .token_compressed_accounts
316 .iter()
317 .find(|acc| acc.compressed_account.compressed_account.address == Some(address)),
318 (_, Some(hash)) => self
319 .token_compressed_accounts
320 .iter()
321 .find(|acc| acc.compressed_account.hash() == Ok(hash)),
322 (None, None) => {
323 return Err(IndexerError::InvalidParameters(
324 "Either address or hash must be provided".to_string(),
325 ))
326 }
327 };
328
329 let amount = account
330 .map(|acc| acc.token_data.amount)
331 .ok_or(IndexerError::AccountNotFound)?;
332
333 Ok(Response {
334 context: Context {
335 slot: self.get_current_slot(),
336 },
337 value: amount,
338 })
339 }
340
341 async fn get_multiple_compressed_accounts(
342 &self,
343 addresses: Option<Vec<Address>>,
344 hashes: Option<Vec<Hash>>,
345 _config: Option<IndexerRpcConfig>,
346 ) -> Result<Response<Items<Option<CompressedAccount>>>, IndexerError> {
347 match (addresses, hashes) {
348 (Some(addresses), _) => {
349 let accounts: Result<Vec<Option<CompressedAccount>>, IndexerError> = addresses
350 .iter()
351 .map(|addr| {
352 self.compressed_accounts
353 .iter()
354 .find(|acc| acc.compressed_account.address == Some(*addr))
355 .map(|acc| acc.clone().try_into())
356 .transpose()
357 })
358 .collect();
359 Ok(Response {
360 context: Context {
361 slot: self.get_current_slot(),
362 },
363 value: Items { items: accounts? },
364 })
365 }
366 (_, Some(hashes)) => {
367 let accounts: Result<Vec<Option<CompressedAccount>>, IndexerError> = hashes
368 .iter()
369 .map(|hash| {
370 self.compressed_accounts
371 .iter()
372 .find(|acc| acc.hash() == Ok(*hash))
373 .map(|acc| acc.clone().try_into())
374 .transpose()
375 })
376 .collect();
377 Ok(Response {
378 context: Context {
379 slot: self.get_current_slot(),
380 },
381 value: Items { items: accounts? },
382 })
383 }
384 (None, None) => Err(IndexerError::InvalidParameters(
385 "Either addresses or hashes must be provided".to_string(),
386 )),
387 }
388 }
389
390 async fn get_compressed_token_balances_by_owner_v2(
391 &self,
392 owner: &Pubkey,
393 _options: Option<GetCompressedTokenAccountsByOwnerOrDelegateOptions>,
394 _config: Option<IndexerRpcConfig>,
395 ) -> Result<Response<ItemsWithCursor<TokenBalance>>, IndexerError> {
396 let mint = _options.as_ref().and_then(|opts| opts.mint);
397 let balances: Vec<TokenBalance> = self
398 .token_compressed_accounts
399 .iter()
400 .filter(|acc| {
401 acc.token_data.owner == *owner && mint.is_none_or(|m| acc.token_data.mint == m)
402 })
403 .fold(std::collections::HashMap::new(), |mut map, acc| {
404 *map.entry(acc.token_data.mint).or_insert(0) += acc.token_data.amount;
405 map
406 })
407 .into_iter()
408 .map(|(mint, balance)| TokenBalance { balance, mint })
409 .collect();
410
411 Ok(Response {
412 context: Context {
413 slot: self.get_current_slot(),
414 },
415 value: ItemsWithCursor {
416 items: balances,
417 cursor: None,
418 },
419 })
420 }
421
422 async fn get_compression_signatures_for_account(
423 &self,
424 _hash: Hash,
425 _config: Option<IndexerRpcConfig>,
426 ) -> Result<Response<Items<SignatureWithMetadata>>, IndexerError> {
427 todo!()
428 }
429
430 async fn get_multiple_new_address_proofs(
431 &self,
432 merkle_tree_pubkey: [u8; 32],
433 addresses: Vec<[u8; 32]>,
434 _config: Option<IndexerRpcConfig>,
435 ) -> Result<Response<Items<NewAddressProofWithContext>>, IndexerError> {
436 let proofs = self
437 ._get_multiple_new_address_proofs(merkle_tree_pubkey, addresses, false)
438 .await?;
439 Ok(Response {
440 context: Context {
441 slot: self.get_current_slot(),
442 },
443 value: Items { items: proofs },
444 })
445 }
446
447 async fn get_validity_proof(
448 &self,
449 hashes: Vec<[u8; 32]>,
450 new_addresses_with_trees: Vec<AddressWithTree>,
451 _config: Option<IndexerRpcConfig>,
452 ) -> Result<Response<ValidityProofWithContext>, IndexerError> {
453 #[cfg(feature = "v2")]
454 {
455 let mut state_merkle_tree_pubkeys = Vec::new();
457
458 for hash in hashes.iter() {
459 let account = self.get_compressed_account_by_hash(*hash, None).await?;
460 let account_data = account.value.ok_or(IndexerError::AccountNotFound)?;
461 state_merkle_tree_pubkeys.push(account_data.tree_info.tree);
462 }
463 let mut proof_inputs = vec![];
464
465 let mut indices_to_remove = Vec::new();
466 let compressed_accounts = if !hashes.is_empty() && !state_merkle_tree_pubkeys.is_empty()
468 {
469 let zipped_accounts = hashes.iter().zip(state_merkle_tree_pubkeys.iter());
470
471 for (i, (compressed_account, state_merkle_tree_pubkey)) in
472 zipped_accounts.enumerate()
473 {
474 let accounts = self.state_merkle_trees.iter().find(|x| {
475 x.accounts.merkle_tree == *state_merkle_tree_pubkey
476 && x.tree_type == TreeType::StateV2
477 });
478
479 if let Some(accounts) = accounts {
480 let queue_element = accounts
481 .output_queue_elements
482 .iter()
483 .find(|(hash, _)| hash == compressed_account);
484 if let Some((_, index)) = queue_element {
485 if accounts.output_queue_batch_size.is_some()
486 && accounts.leaf_index_in_queue_range(*index as usize)?
487 {
488 use light_client::indexer::RootIndex;
489
490 indices_to_remove.push(i);
491 proof_inputs.push(AccountProofInputs {
492 hash: *compressed_account,
493 root: [0u8; 32],
494 root_index: RootIndex::new_none(),
495 leaf_index: accounts
496 .output_queue_elements
497 .iter()
498 .position(|(x, _)| x == compressed_account)
499 .unwrap()
500 as u64,
501 tree_info: light_client::indexer::TreeInfo {
502 cpi_context: Some(accounts.accounts.cpi_context),
503 tree: accounts.accounts.merkle_tree,
504 queue: accounts.accounts.nullifier_queue,
505 next_tree_info: None,
506 tree_type: accounts.tree_type,
507 },
508 })
509 }
510 }
511 }
512 }
513
514 let compress_accounts = hashes
515 .iter()
516 .enumerate()
517 .filter(|(i, _)| !indices_to_remove.contains(i))
518 .map(|(_, x)| *x)
519 .collect::<Vec<[u8; 32]>>();
520
521 if compress_accounts.is_empty() {
522 None
523 } else {
524 Some(compress_accounts)
525 }
526 } else {
527 None
528 };
529
530 let rpc_result: Option<ValidityProofWithContext> = if (compressed_accounts.is_some()
532 && !compressed_accounts.as_ref().unwrap().is_empty())
533 || !new_addresses_with_trees.is_empty()
534 {
535 Some(
536 self._get_validity_proof_v1_implementation(
537 compressed_accounts.unwrap_or_default(),
538 new_addresses_with_trees,
539 )
540 .await?,
541 )
542 } else {
543 None
544 };
545
546 let addresses = if let Some(rpc_result) = rpc_result.as_ref() {
548 rpc_result.addresses.to_vec()
549 } else {
550 Vec::new()
551 };
552 let accounts = {
553 let mut root_indices = if let Some(rpc_result) = rpc_result.as_ref() {
554 rpc_result.accounts.to_vec()
555 } else {
556 Vec::new()
557 };
558 #[cfg(debug_assertions)]
559 {
560 if std::env::var("RUST_BACKTRACE").is_ok() {
561 println!("get_validit_proof: rpc_result {:?}", rpc_result);
562 }
563 }
564
565 for (proof_input, &index) in proof_inputs.iter().zip(indices_to_remove.iter()) {
567 if root_indices.len() <= index {
568 root_indices.push(proof_input.clone());
569 } else {
570 root_indices.insert(index, proof_input.clone());
571 }
572 }
573 root_indices
574 };
575
576 Ok(Response {
577 context: Context {
578 slot: self.get_current_slot(),
579 },
580 value: ValidityProofWithContext {
581 accounts,
582 addresses,
583 proof: rpc_result
584 .map(|rpc_result| rpc_result.proof.0.unwrap())
585 .into(),
586 },
587 })
588 }
589
590 #[cfg(not(feature = "v2"))]
591 {
592 let result = self
594 ._get_validity_proof_v1_implementation(hashes, new_addresses_with_trees)
595 .await?;
596 Ok(Response {
597 context: Context {
598 slot: self.get_current_slot(),
599 },
600 value: result,
601 })
602 }
603 }
604
605 async fn get_queue_elements(
606 &mut self,
607 _merkle_tree_pubkey: [u8; 32],
608 _queue_type: QueueType,
609 _num_elements: u16,
610 _start_offset: Option<u64>,
611 _config: Option<IndexerRpcConfig>,
612 ) -> Result<Response<QueueElementsResult>, IndexerError> {
613 #[cfg(not(feature = "v2"))]
614 unimplemented!("get_queue_elements");
615 #[cfg(feature = "v2")]
616 {
617 let merkle_tree_pubkey = _merkle_tree_pubkey;
618 let queue_type = _queue_type;
619 let num_elements = _num_elements;
620 let pubkey = Pubkey::new_from_array(merkle_tree_pubkey);
621 let address_tree_bundle = self
622 .address_merkle_trees
623 .iter()
624 .find(|x| x.accounts.merkle_tree == pubkey);
625 if let Some(address_tree_bundle) = address_tree_bundle {
626 let end_offset = std::cmp::min(
627 num_elements as usize,
628 address_tree_bundle.queue_elements.len(),
629 );
630 let queue_elements = address_tree_bundle.queue_elements[0..end_offset].to_vec();
631
632 let merkle_proofs_with_context = queue_elements
633 .iter()
634 .map(|element| MerkleProofWithContext {
635 proof: Vec::new(),
636 leaf: [0u8; 32],
637 leaf_index: 0,
638 merkle_tree: address_tree_bundle.accounts.merkle_tree.to_bytes(),
639 root: address_tree_bundle.root(),
640 tx_hash: None,
641 root_seq: 0,
642 account_hash: *element,
643 })
644 .collect();
645 return Ok(Response {
646 context: Context {
647 slot: self.get_current_slot(),
648 },
649 value: QueueElementsResult {
650 elements: merkle_proofs_with_context,
651 first_value_queue_index: None,
652 },
653 });
654 }
655
656 let state_tree_bundle = self
657 .state_merkle_trees
658 .iter_mut()
659 .find(|x| x.accounts.merkle_tree == pubkey);
660 if queue_type == QueueType::InputStateV2 {
661 if let Some(state_tree_bundle) = state_tree_bundle {
662 let end_offset = std::cmp::min(
663 num_elements as usize,
664 state_tree_bundle.input_leaf_indices.len(),
665 );
666 let queue_elements =
667 state_tree_bundle.input_leaf_indices[0..end_offset].to_vec();
668 let merkle_proofs = queue_elements
669 .iter()
670 .map(|leaf_info| {
671 match state_tree_bundle
672 .merkle_tree
673 .get_proof_of_leaf(leaf_info.leaf_index as usize, true)
674 {
675 Ok(proof) => proof.to_vec(),
676 Err(_) => {
677 let mut next_index =
678 state_tree_bundle.merkle_tree.get_next_index() as u64;
679 while next_index < leaf_info.leaf_index as u64 {
680 state_tree_bundle.merkle_tree.append(&[0u8; 32]).unwrap();
681 next_index =
682 state_tree_bundle.merkle_tree.get_next_index() as u64;
683 }
684 state_tree_bundle
685 .merkle_tree
686 .get_proof_of_leaf(leaf_info.leaf_index as usize, true)
687 .unwrap()
688 .to_vec();
689 Vec::new()
690 }
691 }
692 })
693 .collect::<Vec<_>>();
694 let leaves = queue_elements
695 .iter()
696 .map(|leaf_info| {
697 state_tree_bundle
698 .merkle_tree
699 .get_leaf(leaf_info.leaf_index as usize)
700 .unwrap_or_default()
701 })
702 .collect::<Vec<_>>();
703 let merkle_proofs_with_context = merkle_proofs
704 .iter()
705 .zip(queue_elements.iter())
706 .zip(leaves.iter())
707 .map(|((proof, element), leaf)| MerkleProofWithContext {
708 proof: proof.clone(),
709 leaf: *leaf,
710 leaf_index: element.leaf_index as u64,
711 merkle_tree: state_tree_bundle.accounts.merkle_tree.to_bytes(),
712 root: state_tree_bundle.merkle_tree.root(),
713 tx_hash: Some(element.tx_hash),
714 root_seq: 0,
715 account_hash: element.leaf,
716 })
717 .collect();
718
719 return Ok(Response {
720 context: Context {
721 slot: self.get_current_slot(),
722 },
723 value: QueueElementsResult {
724 elements: merkle_proofs_with_context,
725 first_value_queue_index: None,
726 },
727 });
728 }
729 }
730
731 if queue_type == QueueType::OutputStateV2 {
732 if let Some(state_tree_bundle) = state_tree_bundle {
733 let end_offset = std::cmp::min(
734 num_elements as usize,
735 state_tree_bundle.output_queue_elements.len(),
736 );
737 let queue_elements =
738 state_tree_bundle.output_queue_elements[0..end_offset].to_vec();
739 let indices = queue_elements
740 .iter()
741 .map(|(_, index)| index)
742 .collect::<Vec<_>>();
743 let merkle_proofs = indices
744 .iter()
745 .map(|index| {
746 match state_tree_bundle
747 .merkle_tree
748 .get_proof_of_leaf(**index as usize, true)
749 {
750 Ok(proof) => proof.to_vec(),
751 Err(_) => {
752 let mut next_index =
753 state_tree_bundle.merkle_tree.get_next_index() as u64;
754 while next_index < **index {
755 state_tree_bundle.merkle_tree.append(&[0u8; 32]).unwrap();
756 next_index =
757 state_tree_bundle.merkle_tree.get_next_index() as u64;
758 }
759 state_tree_bundle
760 .merkle_tree
761 .get_proof_of_leaf(**index as usize, true)
762 .unwrap()
763 .to_vec();
764 Vec::new()
765 }
766 }
767 })
768 .collect::<Vec<_>>();
769 let leaves = indices
770 .iter()
771 .map(|index| {
772 state_tree_bundle
773 .merkle_tree
774 .get_leaf(**index as usize)
775 .unwrap_or_default()
776 })
777 .collect::<Vec<_>>();
778 let merkle_proofs_with_context = merkle_proofs
779 .iter()
780 .zip(queue_elements.iter())
781 .zip(leaves.iter())
782 .map(|((proof, (element, index)), leaf)| MerkleProofWithContext {
783 proof: proof.clone(),
784 leaf: *leaf,
785 leaf_index: *index,
786 merkle_tree: state_tree_bundle.accounts.merkle_tree.to_bytes(),
787 root: state_tree_bundle.merkle_tree.root(),
788 tx_hash: None,
789 root_seq: 0,
790 account_hash: *element,
791 })
792 .collect();
793 return Ok(Response {
794 context: Context {
795 slot: self.get_current_slot(),
796 },
797 value: QueueElementsResult {
798 elements: merkle_proofs_with_context,
799 first_value_queue_index: if queue_elements.is_empty() {
800 None
801 } else {
802 Some(queue_elements[0].1)
803 },
804 },
805 });
806 }
807 }
808
809 Err(IndexerError::InvalidParameters(
810 "Merkle tree not found".to_string(),
811 ))
812 }
813 }
814
815 async fn get_subtrees(
816 &self,
817 _merkle_tree_pubkey: [u8; 32],
818 _config: Option<IndexerRpcConfig>,
819 ) -> Result<Response<Items<[u8; 32]>>, IndexerError> {
820 #[cfg(not(feature = "v2"))]
821 unimplemented!("get_subtrees");
822 #[cfg(feature = "v2")]
823 {
824 let merkle_tree_pubkey = Pubkey::new_from_array(_merkle_tree_pubkey);
825 let address_tree_bundle = self
826 .address_merkle_trees
827 .iter()
828 .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey);
829 if let Some(address_tree_bundle) = address_tree_bundle {
830 Ok(Response {
831 context: Context {
832 slot: self.get_current_slot(),
833 },
834 value: Items {
835 items: address_tree_bundle.get_subtrees(),
836 },
837 })
838 } else {
839 let state_tree_bundle = self
840 .state_merkle_trees
841 .iter()
842 .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey);
843 if let Some(state_tree_bundle) = state_tree_bundle {
844 Ok(Response {
845 context: Context {
846 slot: self.get_current_slot(),
847 },
848 value: Items {
849 items: state_tree_bundle.merkle_tree.get_subtrees(),
850 },
851 })
852 } else {
853 Err(IndexerError::InvalidParameters(
854 "Merkle tree not found".to_string(),
855 ))
856 }
857 }
858 }
859 }
860
861 async fn get_address_queue_with_proofs(
862 &mut self,
863 _merkle_tree_pubkey: &Pubkey,
864 _zkp_batch_size: u16,
865 _start_offset: Option<u64>,
866 _config: Option<IndexerRpcConfig>,
867 ) -> Result<Response<BatchAddressUpdateIndexerResponse>, IndexerError> {
868 #[cfg(not(feature = "v2"))]
869 unimplemented!("get_address_queue_with_proofs");
870 #[cfg(feature = "v2")]
871 {
872 use light_client::indexer::AddressQueueIndex;
873 let merkle_tree_pubkey = _merkle_tree_pubkey;
874 let zkp_batch_size = _zkp_batch_size;
875
876 let batch_start_index = self
877 .get_address_merkle_trees()
878 .iter()
879 .find(|x| x.accounts.merkle_tree == *merkle_tree_pubkey)
880 .unwrap()
881 .get_v2_indexed_merkle_tree()
882 .ok_or(IndexerError::Unknown(
883 "Failed to get v2 indexed merkle tree".into(),
884 ))?
885 .merkle_tree
886 .rightmost_index;
887
888 let address_proof_items = self
889 .get_queue_elements(
890 merkle_tree_pubkey.to_bytes(),
891 QueueType::AddressV2,
892 zkp_batch_size,
893 None,
894 None,
895 )
896 .await
897 .map_err(|_| IndexerError::Unknown("Failed to get queue elements".into()))?
898 .value;
899
900 let addresses: Vec<AddressQueueIndex> = address_proof_items
901 .elements
902 .iter()
903 .enumerate()
904 .map(|(i, proof)| AddressQueueIndex {
905 address: proof.account_hash,
906 queue_index: proof.root_seq + i as u64,
907 })
908 .collect();
909 let non_inclusion_proofs = self
910 .get_multiple_new_address_proofs(
911 merkle_tree_pubkey.to_bytes(),
912 address_proof_items
913 .elements
914 .iter()
915 .map(|x| x.account_hash)
916 .collect(),
917 None,
918 )
919 .await
920 .map_err(|_| {
921 IndexerError::Unknown(
922 "Failed to get get_multiple_new_address_proofs_full".into(),
923 )
924 })?
925 .value;
926
927 let subtrees = self
928 .get_subtrees(merkle_tree_pubkey.to_bytes(), None)
929 .await
930 .map_err(|_| IndexerError::Unknown("Failed to get subtrees".into()))?
931 .value;
932
933 Ok(Response {
934 context: Context {
935 slot: self.get_current_slot(),
936 },
937 value: BatchAddressUpdateIndexerResponse {
938 batch_start_index: batch_start_index as u64,
939 addresses,
940 non_inclusion_proofs: non_inclusion_proofs.items,
941 subtrees: subtrees.items,
942 },
943 })
944 }
945 }
946
947 async fn get_compressed_balance_by_owner(
949 &self,
950 _owner: &Pubkey,
951 _config: Option<IndexerRpcConfig>,
952 ) -> Result<Response<u64>, IndexerError> {
953 todo!("get_compressed_balance_by_owner not implemented")
954 }
955
956 async fn get_compressed_mint_token_holders(
957 &self,
958 _mint: &Pubkey,
959 _options: Option<PaginatedOptions>,
960 _config: Option<IndexerRpcConfig>,
961 ) -> Result<Response<ItemsWithCursor<OwnerBalance>>, IndexerError> {
962 todo!("get_compressed_mint_token_holders not implemented")
963 }
964
965 async fn get_compressed_token_accounts_by_delegate(
966 &self,
967 _delegate: &Pubkey,
968 _options: Option<GetCompressedTokenAccountsByOwnerOrDelegateOptions>,
969 _config: Option<IndexerRpcConfig>,
970 ) -> Result<Response<ItemsWithCursor<CompressedTokenAccount>>, IndexerError> {
971 todo!("get_compressed_token_accounts_by_delegate not implemented")
972 }
973
974 async fn get_compression_signatures_for_address(
975 &self,
976 _address: &[u8; 32],
977 _options: Option<PaginatedOptions>,
978 _config: Option<IndexerRpcConfig>,
979 ) -> Result<Response<ItemsWithCursor<SignatureWithMetadata>>, IndexerError> {
980 todo!("get_compression_signatures_for_address not implemented")
981 }
982
983 async fn get_compression_signatures_for_owner(
984 &self,
985 _owner: &Pubkey,
986 _options: Option<PaginatedOptions>,
987 _config: Option<IndexerRpcConfig>,
988 ) -> Result<Response<ItemsWithCursor<SignatureWithMetadata>>, IndexerError> {
989 todo!("get_compression_signatures_for_owner not implemented")
990 }
991
992 async fn get_compression_signatures_for_token_owner(
993 &self,
994 _owner: &Pubkey,
995 _options: Option<PaginatedOptions>,
996 _config: Option<IndexerRpcConfig>,
997 ) -> Result<Response<ItemsWithCursor<SignatureWithMetadata>>, IndexerError> {
998 todo!("get_compression_signatures_for_token_owner not implemented")
999 }
1000
1001 async fn get_indexer_health(&self, _config: Option<RetryConfig>) -> Result<bool, IndexerError> {
1002 todo!("get_indexer_health not implemented")
1003 }
1004}
1005
1006#[async_trait]
1007impl TestIndexerExtensions for TestIndexer {
1008 fn get_address_merkle_trees(&self) -> &Vec<AddressMerkleTreeBundle> {
1009 &self.address_merkle_trees
1010 }
1011
1012 fn get_address_merkle_tree(
1013 &self,
1014 merkle_tree_pubkey: Pubkey,
1015 ) -> Option<&AddressMerkleTreeBundle> {
1016 self.address_merkle_trees
1017 .iter()
1018 .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey)
1019 }
1020
1021 fn add_compressed_accounts_with_token_data(
1028 &mut self,
1029 slot: u64,
1030 event: &PublicTransactionEvent,
1031 ) {
1032 TestIndexer::add_event_and_compressed_accounts(self, slot, event);
1033 }
1034
1035 fn account_nullified(&mut self, merkle_tree_pubkey: Pubkey, account_hash: &str) {
1036 let decoded_hash: [u8; 32] = bs58::decode(account_hash)
1037 .into_vec()
1038 .unwrap()
1039 .as_slice()
1040 .try_into()
1041 .unwrap();
1042
1043 if let Some(state_tree_bundle) = self
1044 .state_merkle_trees
1045 .iter_mut()
1046 .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey)
1047 {
1048 if let Some(leaf_index) = state_tree_bundle.merkle_tree.get_leaf_index(&decoded_hash) {
1049 state_tree_bundle
1050 .merkle_tree
1051 .update(&[0u8; 32], leaf_index)
1052 .unwrap();
1053 }
1054 }
1055 }
1056
1057 fn address_tree_updated(
1058 &mut self,
1059 merkle_tree_pubkey: Pubkey,
1060 context: &NewAddressProofWithContext,
1061 ) {
1062 info!("Updating address tree...");
1063 let pos = self
1064 .address_merkle_trees
1065 .iter()
1066 .position(|x| x.accounts.merkle_tree == merkle_tree_pubkey)
1067 .unwrap();
1068 let new_low_element = context.new_low_element.clone().unwrap();
1069 let new_element = context.new_element.clone().unwrap();
1070 let new_element_next_value = context.new_element_next_value.clone().unwrap();
1071 self.address_merkle_trees[pos]
1073 .get_v1_indexed_merkle_tree_mut()
1074 .expect("Failed to get v1 indexed merkle tree.")
1075 .update(&new_low_element, &new_element, &new_element_next_value)
1076 .unwrap();
1077 self.address_merkle_trees[pos]
1078 .append_with_low_element_index(new_low_element.index, &new_element.value)
1079 .unwrap();
1080 info!("Address tree updated");
1081 }
1082
1083 fn get_state_merkle_tree_accounts(&self, pubkeys: &[Pubkey]) -> Vec<StateMerkleTreeAccounts> {
1084 pubkeys
1085 .iter()
1086 .map(|x| {
1087 self.state_merkle_trees
1088 .iter()
1089 .find(|y| y.accounts.merkle_tree == *x || y.accounts.nullifier_queue == *x)
1090 .unwrap()
1091 .accounts
1092 })
1093 .collect::<Vec<_>>()
1094 }
1095
1096 fn get_state_merkle_trees(&self) -> &Vec<StateMerkleTreeBundle> {
1097 &self.state_merkle_trees
1098 }
1099
1100 fn get_state_merkle_trees_mut(&mut self) -> &mut Vec<StateMerkleTreeBundle> {
1101 &mut self.state_merkle_trees
1102 }
1103
1104 fn get_address_merkle_trees_mut(&mut self) -> &mut Vec<AddressMerkleTreeBundle> {
1105 &mut self.address_merkle_trees
1106 }
1107
1108 fn get_token_compressed_accounts(&self) -> &Vec<TokenDataWithMerkleContext> {
1109 &self.token_compressed_accounts
1110 }
1111
1112 fn get_group_pda(&self) -> &Pubkey {
1113 &self.group_pda
1114 }
1115
1116 fn add_address_merkle_tree_accounts(
1117 &mut self,
1118 merkle_tree_keypair: &Keypair,
1119 queue_keypair: &Keypair,
1120 _owning_program_id: Option<Pubkey>,
1121 ) -> AddressMerkleTreeAccounts {
1122 info!("Adding address merkle tree accounts...");
1123 let address_merkle_tree_accounts = AddressMerkleTreeAccounts {
1124 merkle_tree: merkle_tree_keypair.pubkey(),
1125 queue: queue_keypair.pubkey(),
1126 };
1127 self.address_merkle_trees
1128 .push(Self::add_address_merkle_tree_bundle(address_merkle_tree_accounts).unwrap());
1129 info!(
1130 "Address merkle tree accounts added. Total: {}",
1131 self.address_merkle_trees.len()
1132 );
1133 address_merkle_tree_accounts
1134 }
1135
1136 fn get_compressed_accounts_with_merkle_context_by_owner(
1137 &self,
1138 owner: &Pubkey,
1139 ) -> Vec<CompressedAccountWithMerkleContext> {
1140 self.compressed_accounts
1141 .iter()
1142 .filter(|x| x.compressed_account.owner.to_bytes() == owner.to_bytes())
1143 .cloned()
1144 .collect()
1145 }
1146
1147 fn add_state_bundle(&mut self, state_bundle: StateMerkleTreeBundle) {
1148 Self::get_state_merkle_trees_mut(self).push(state_bundle);
1149 }
1150
1151 fn add_event_and_compressed_accounts(
1152 &mut self,
1153 slot: u64,
1154 event: &PublicTransactionEvent,
1155 ) -> (
1156 Vec<CompressedAccountWithMerkleContext>,
1157 Vec<TokenDataWithMerkleContext>,
1158 ) {
1159 let mut compressed_accounts = Vec::new();
1160 let mut token_compressed_accounts = Vec::new();
1161 let event_inputs_len = event.input_compressed_account_hashes.len();
1162 let event_outputs_len = event.output_compressed_account_hashes.len();
1163 for i in 0..std::cmp::max(event_inputs_len, event_outputs_len) {
1164 self.process_v1_compressed_account(
1165 slot,
1166 event,
1167 i,
1168 &mut token_compressed_accounts,
1169 &mut compressed_accounts,
1170 );
1171 }
1172
1173 self.events.push(event.clone());
1174 (compressed_accounts, token_compressed_accounts)
1175 }
1176
1177 fn get_proof_by_index(&mut self, merkle_tree_pubkey: Pubkey, index: u64) -> MerkleProof {
1178 let bundle = self
1179 .state_merkle_trees
1180 .iter_mut()
1181 .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey)
1182 .unwrap();
1183
1184 while bundle.merkle_tree.leaves().len() <= index as usize {
1185 bundle.merkle_tree.append(&[0u8; 32]).unwrap();
1186 }
1187
1188 let leaf = match bundle.merkle_tree.get_leaf(index as usize) {
1189 Ok(leaf) => leaf,
1190 Err(_) => {
1191 bundle.merkle_tree.append(&[0u8; 32]).unwrap();
1192 bundle.merkle_tree.get_leaf(index as usize).unwrap()
1193 }
1194 };
1195
1196 let proof = bundle
1197 .merkle_tree
1198 .get_proof_of_leaf(index as usize, true)
1199 .unwrap()
1200 .to_vec();
1201
1202 MerkleProof {
1203 hash: leaf,
1204 leaf_index: index,
1205 merkle_tree: merkle_tree_pubkey,
1206 proof,
1207 root_seq: bundle.merkle_tree.sequence_number as u64,
1208 root: bundle.merkle_tree.root(),
1209 }
1210 }
1211
1212 async fn finalize_batched_address_tree_update(
1213 &mut self,
1214 merkle_tree_pubkey: Pubkey,
1215 account_data: &mut [u8],
1216 ) {
1217 let onchain_account =
1218 BatchedMerkleTreeAccount::address_from_bytes(account_data, &merkle_tree_pubkey.into())
1219 .unwrap();
1220 let address_tree = self
1221 .address_merkle_trees
1222 .iter_mut()
1223 .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey)
1224 .unwrap();
1225 let address_tree_index = address_tree.right_most_index();
1226 let onchain_next_index = onchain_account.next_index;
1227 let diff_onchain_indexer = onchain_next_index - address_tree_index as u64;
1228 let addresses = address_tree.queue_elements[0..diff_onchain_indexer as usize].to_vec();
1229 for _ in 0..diff_onchain_indexer {
1230 address_tree.queue_elements.remove(0);
1231 }
1232 for new_element_value in &addresses {
1233 address_tree
1234 .append(&BigUint::from_bytes_be(new_element_value))
1235 .unwrap();
1236 }
1237 match &mut address_tree.merkle_tree {
1238 IndexedMerkleTreeVersion::V2(tree) => tree.merkle_tree.num_root_updates += 1,
1239 IndexedMerkleTreeVersion::V1(_) => {
1240 unimplemented!("finalize_batched_address_tree_update not implemented for v1 trees.")
1241 }
1242 }
1243 let onchain_root = onchain_account.root_history.last().unwrap();
1244 let new_root = address_tree.root();
1245 assert_eq!(*onchain_root, new_root);
1246 }
1247}
1248
1249impl TestIndexer {
1250 fn get_current_slot(&self) -> u64 {
1251 u64::MAX
1253 }
1254
1255 pub async fn init_from_acounts(
1256 payer: &Keypair,
1257 env: &TestAccounts,
1258 output_queue_batch_size: usize,
1259 ) -> Self {
1260 let mut state_merkle_tree_accounts = env.v1_state_trees.clone();
1262
1263 for v2_state_tree in &env.v2_state_trees {
1265 state_merkle_tree_accounts.push(StateMerkleTreeAccounts {
1266 merkle_tree: v2_state_tree.merkle_tree,
1267 nullifier_queue: v2_state_tree.output_queue,
1268 cpi_context: v2_state_tree.cpi_context,
1269 tree_type: TreeType::StateV2,
1270 });
1271 }
1272
1273 let mut address_merkle_tree_accounts = env.v1_address_trees.clone();
1275
1276 for &v2_address_tree in &env.v2_address_trees {
1278 address_merkle_tree_accounts.push(AddressMerkleTreeAccounts {
1279 merkle_tree: v2_address_tree,
1280 queue: v2_address_tree,
1281 });
1282 }
1283
1284 Self::new(
1285 state_merkle_tree_accounts,
1286 address_merkle_tree_accounts,
1287 payer.insecure_clone(),
1288 env.protocol.group_pda,
1289 output_queue_batch_size,
1290 )
1291 .await
1292 }
1293
1294 pub async fn new(
1295 state_merkle_tree_accounts: Vec<StateMerkleTreeAccounts>,
1296 address_merkle_tree_accounts: Vec<AddressMerkleTreeAccounts>,
1297 payer: Keypair,
1298 group_pda: Pubkey,
1299 output_queue_batch_size: usize,
1300 ) -> Self {
1301 let mut state_merkle_trees = Vec::new();
1302 for state_merkle_tree_account in state_merkle_tree_accounts.iter() {
1303 let (tree_type, merkle_tree, output_queue_batch_size) =
1304 if state_merkle_tree_account.tree_type == TreeType::StateV2 {
1305 let merkle_tree = Box::new(MerkleTree::<Poseidon>::new_with_history(
1306 DEFAULT_BATCH_STATE_TREE_HEIGHT as usize,
1307 0,
1308 0,
1309 DEFAULT_BATCH_ROOT_HISTORY_LEN as usize,
1310 ));
1311 (
1312 TreeType::StateV2,
1313 merkle_tree,
1314 Some(output_queue_batch_size),
1315 )
1316 } else {
1317 let merkle_tree = Box::new(MerkleTree::<Poseidon>::new_with_history(
1318 account_compression::utils::constants::STATE_MERKLE_TREE_HEIGHT as usize,
1319 account_compression::utils::constants::STATE_MERKLE_TREE_CANOPY_DEPTH
1320 as usize,
1321 0,
1322 account_compression::utils::constants::STATE_MERKLE_TREE_ROOTS as usize,
1323 ));
1324 (TreeType::StateV1, merkle_tree, None)
1325 };
1326
1327 state_merkle_trees.push(StateMerkleTreeBundle {
1328 accounts: *state_merkle_tree_account,
1329 merkle_tree,
1330 rollover_fee: FeeConfig::default().state_merkle_tree_rollover as i64,
1331 tree_type,
1332 output_queue_elements: vec![],
1333 input_leaf_indices: vec![],
1334 output_queue_batch_size,
1335 num_inserted_batches: 0,
1336 });
1337 }
1338
1339 let mut address_merkle_trees = Vec::new();
1340 for address_merkle_tree_account in address_merkle_tree_accounts {
1341 address_merkle_trees
1342 .push(Self::add_address_merkle_tree_bundle(address_merkle_tree_account).unwrap());
1343 }
1344
1345 Self {
1346 state_merkle_trees,
1347 address_merkle_trees,
1348 payer,
1349 compressed_accounts: vec![],
1350 nullified_compressed_accounts: vec![],
1351 events: vec![],
1352 token_compressed_accounts: vec![],
1353 token_nullified_compressed_accounts: vec![],
1354 group_pda,
1355 }
1356 }
1357
1358 pub fn add_address_merkle_tree_bundle(
1359 address_merkle_tree_accounts: AddressMerkleTreeAccounts,
1360 ) -> Result<AddressMerkleTreeBundle, IndexerError> {
1362 if address_merkle_tree_accounts.merkle_tree == address_merkle_tree_accounts.queue {
1363 AddressMerkleTreeBundle::new_v2(address_merkle_tree_accounts)
1364 } else {
1365 AddressMerkleTreeBundle::new_v1(address_merkle_tree_accounts)
1366 }
1367 }
1368
1369 async fn add_address_merkle_tree_v1<R: Rpc>(
1370 &mut self,
1371 rpc: &mut R,
1372 merkle_tree_keypair: &Keypair,
1373 queue_keypair: &Keypair,
1374 owning_program_id: Option<Pubkey>,
1375 ) -> Result<AddressMerkleTreeAccounts, RpcError> {
1376 create_address_merkle_tree_and_queue_account(
1377 &self.payer,
1378 true,
1379 rpc,
1380 merkle_tree_keypair,
1381 queue_keypair,
1382 owning_program_id,
1383 None,
1384 &AddressMerkleTreeConfig::default(),
1385 &AddressQueueConfig::default(),
1386 0,
1387 )
1388 .await?;
1389
1390 let accounts = <TestIndexer as TestIndexerExtensions>::add_address_merkle_tree_accounts(
1391 self,
1392 merkle_tree_keypair,
1393 queue_keypair,
1394 owning_program_id,
1395 );
1396 Ok(accounts)
1397 }
1398
1399 #[cfg(feature = "devenv")]
1400 async fn add_address_merkle_tree_v2<R: Rpc>(
1401 &mut self,
1402 rpc: &mut R,
1403 merkle_tree_keypair: &Keypair,
1404 queue_keypair: &Keypair,
1405 _owning_program_id: Option<Pubkey>,
1406 ) -> Result<AddressMerkleTreeAccounts, RpcError> {
1407 info!(
1408 "Adding address merkle tree accounts v2 {:?}",
1409 merkle_tree_keypair.pubkey()
1410 );
1411
1412 let params = light_batched_merkle_tree::initialize_address_tree::InitAddressTreeAccountsInstructionData::test_default();
1413
1414 info!(
1415 "Creating batched address merkle tree {:?}",
1416 merkle_tree_keypair.pubkey()
1417 );
1418 create_batch_address_merkle_tree(rpc, &self.payer, merkle_tree_keypair, params).await?;
1419 info!(
1420 "Batched address merkle tree created {:?}",
1421 merkle_tree_keypair.pubkey()
1422 );
1423
1424 let accounts = self.add_address_merkle_tree_accounts(
1425 merkle_tree_keypair,
1426 queue_keypair,
1427 _owning_program_id,
1428 );
1429 Ok(accounts)
1430 }
1431
1432 pub async fn add_address_merkle_tree<R: Rpc>(
1433 &mut self,
1434 rpc: &mut R,
1435 merkle_tree_keypair: &Keypair,
1436 queue_keypair: &Keypair,
1437 owning_program_id: Option<Pubkey>,
1438 tree_type: TreeType,
1439 ) -> Result<AddressMerkleTreeAccounts, RpcError> {
1440 if tree_type == TreeType::AddressV1 {
1441 self.add_address_merkle_tree_v1(
1442 rpc,
1443 merkle_tree_keypair,
1444 queue_keypair,
1445 owning_program_id,
1446 )
1447 .await
1448 } else if tree_type == TreeType::AddressV2 {
1449 #[cfg(not(feature = "devenv"))]
1450 panic!("Batched address merkle trees require the 'devenv' feature to be enabled");
1451 #[cfg(feature = "devenv")]
1452 self.add_address_merkle_tree_v2(
1453 rpc,
1454 merkle_tree_keypair,
1455 queue_keypair,
1456 owning_program_id,
1457 )
1458 .await
1459 } else {
1460 Err(RpcError::CustomError(format!(
1461 "add_address_merkle_tree: Version not supported, {}. Versions: AddressV1, AddressV2",
1462 tree_type
1463 )))
1464 }
1465 }
1466
1467 #[allow(clippy::too_many_arguments)]
1468 pub async fn add_state_merkle_tree<R: Rpc>(
1469 &mut self,
1470 rpc: &mut R,
1471 merkle_tree_keypair: &Keypair,
1472 queue_keypair: &Keypair,
1473 cpi_context_keypair: &Keypair,
1474 owning_program_id: Option<Pubkey>,
1475 forester: Option<Pubkey>,
1476 tree_type: TreeType,
1477 ) {
1478 let (rollover_fee, merkle_tree, output_queue_batch_size) = match tree_type {
1479 TreeType::StateV1 => {
1480 create_state_merkle_tree_and_queue_account(
1481 &self.payer,
1482 true,
1483 rpc,
1484 merkle_tree_keypair,
1485 queue_keypair,
1486 Some(cpi_context_keypair),
1487 owning_program_id,
1488 forester,
1489 self.state_merkle_trees.len() as u64,
1490 &StateMerkleTreeConfig::default(),
1491 &NullifierQueueConfig::default(),
1492 )
1493 .await
1494 .unwrap();
1495 let merkle_tree = Box::new(MerkleTree::<Poseidon>::new_with_history(
1496 account_compression::utils::constants::STATE_MERKLE_TREE_HEIGHT as usize,
1497 account_compression::utils::constants::STATE_MERKLE_TREE_CANOPY_DEPTH as usize,
1498 0,
1499 account_compression::utils::constants::STATE_MERKLE_TREE_ROOTS as usize,
1500
1501 ));
1502 (FeeConfig::default().state_merkle_tree_rollover as i64,merkle_tree, None)
1503 }
1504 TreeType::StateV2 => {
1505 #[cfg(feature = "devenv")]
1506 {
1507 let params = light_batched_merkle_tree::initialize_state_tree::InitStateTreeAccountsInstructionData::test_default();
1508
1509 create_batched_state_merkle_tree(
1510 &self.payer,
1511 true,
1512 rpc,
1513 merkle_tree_keypair,
1514 queue_keypair,
1515 cpi_context_keypair,
1516 params,
1517 ).await.unwrap();
1518 let merkle_tree = Box::new(MerkleTree::<Poseidon>::new_with_history(
1519 DEFAULT_BATCH_STATE_TREE_HEIGHT as usize,
1520 0,
1521 0,
1522 DEFAULT_BATCH_ROOT_HISTORY_LEN as usize,
1523
1524 ));
1525 (FeeConfig::test_batched().state_merkle_tree_rollover as i64,merkle_tree, Some(params.output_queue_batch_size as usize))
1526 }
1527
1528 #[cfg(not(feature = "devenv"))]
1529 panic!("Batched state merkle trees require the 'devenv' feature to be enabled")
1530 }
1531 _ => panic!(
1532 "add_state_merkle_tree: tree_type not supported, {}. tree_type: 1 concurrent, 2 batched",
1533 tree_type
1534 ),
1535 };
1536 let state_merkle_tree_account = StateMerkleTreeAccounts {
1537 merkle_tree: merkle_tree_keypair.pubkey(),
1538 nullifier_queue: queue_keypair.pubkey(),
1539 cpi_context: cpi_context_keypair.pubkey(),
1540 tree_type,
1541 };
1542
1543 self.state_merkle_trees.push(StateMerkleTreeBundle {
1544 merkle_tree,
1545 accounts: state_merkle_tree_account,
1546 rollover_fee,
1547 tree_type,
1548 output_queue_elements: vec![],
1549 input_leaf_indices: vec![],
1550 num_inserted_batches: 0,
1551 output_queue_batch_size,
1552 });
1553 println!(
1554 "creating Merkle tree bundle {:?}",
1555 self.state_merkle_trees
1556 .iter()
1557 .map(|x| x.accounts.merkle_tree)
1558 .collect::<Vec<_>>()
1559 );
1560 }
1561
1562 pub fn add_lamport_compressed_accounts(&mut self, slot: u64, event_bytes: Vec<u8>) {
1567 let event_bytes = event_bytes.clone();
1568 let event = PublicTransactionEvent::deserialize(&mut event_bytes.as_slice()).unwrap();
1569 <TestIndexer as TestIndexerExtensions>::add_event_and_compressed_accounts(
1571 self, slot, &event,
1572 );
1573 }
1574
1575 pub fn get_compressed_balance(&self, owner: &Pubkey) -> u64 {
1577 self.compressed_accounts
1578 .iter()
1579 .filter(|x| x.compressed_account.owner.to_bytes() == owner.to_bytes())
1580 .map(|x| x.compressed_account.lamports)
1581 .sum()
1582 }
1583
1584 pub fn get_compressed_token_balance(&self, owner: &Pubkey, mint: &Pubkey) -> u64 {
1586 self.token_compressed_accounts
1587 .iter()
1588 .filter(|x| {
1589 x.compressed_account.compressed_account.owner.to_bytes() == owner.to_bytes()
1590 && x.token_data.mint == *mint
1591 })
1592 .map(|x| x.token_data.amount)
1593 .sum()
1594 }
1595
1596 fn process_v1_compressed_account(
1597 &mut self,
1598 slot: u64,
1599 event: &PublicTransactionEvent,
1600 i: usize,
1601 token_compressed_accounts: &mut Vec<TokenDataWithMerkleContext>,
1602 compressed_accounts: &mut Vec<CompressedAccountWithMerkleContext>,
1603 ) {
1604 let mut input_addresses = vec![];
1605 let mut new_addresses = vec![];
1606 if event.output_compressed_accounts.len() > i {
1607 let compressed_account = &event.output_compressed_accounts[i];
1608 if let Some(address) = compressed_account.compressed_account.address {
1609 if !input_addresses.iter().any(|x| x == &address) {
1610 new_addresses.push(address);
1611 }
1612 }
1613 let merkle_tree = self.state_merkle_trees.iter().find(|x| {
1614 x.accounts.merkle_tree
1615 == solana_pubkey::Pubkey::from(
1616 event.pubkey_array
1617 [event.output_compressed_accounts[i].merkle_tree_index as usize]
1618 .to_bytes(),
1619 )
1620 });
1621 let merkle_tree = if let Some(merkle_tree) = merkle_tree {
1623 merkle_tree
1624 } else {
1625 self.state_merkle_trees
1626 .iter()
1627 .find(|x| {
1628 x.accounts.nullifier_queue
1629 == solana_pubkey::Pubkey::from(
1630 event.pubkey_array[event.output_compressed_accounts[i]
1631 .merkle_tree_index
1632 as usize]
1633 .to_bytes(),
1634 )
1635 })
1636 .unwrap()
1637 };
1638 let nullifier_queue_pubkey = merkle_tree.accounts.nullifier_queue;
1639 let merkle_tree_pubkey = merkle_tree.accounts.merkle_tree;
1640 match compressed_account.compressed_account.data.as_ref() {
1644 Some(data) => {
1645 let is_v1_token = data.discriminator == light_compressed_token::constants::TOKEN_COMPRESSED_ACCOUNT_DISCRIMINATOR; let is_v2_token = data.discriminator == [0, 0, 0, 0, 0, 0, 0, 3]; if compressed_account.compressed_account.owner
1650 == light_compressed_token::ID.to_bytes()
1651 && (is_v1_token || is_v2_token)
1652 {
1653 if let Ok(token_data) = TokenData::deserialize(&mut data.data.as_slice()) {
1654 let token_account = TokenDataWithMerkleContext {
1655 token_data,
1656 compressed_account: CompressedAccountWithMerkleContext {
1657 compressed_account: compressed_account
1658 .compressed_account
1659 .clone(),
1660 merkle_context: MerkleContext {
1661 leaf_index: event.output_leaf_indices[i],
1662 merkle_tree_pubkey: merkle_tree_pubkey.into(),
1663 queue_pubkey: nullifier_queue_pubkey.into(),
1664 prove_by_index: false,
1665 tree_type: merkle_tree.tree_type,
1666 },
1667 },
1668 };
1669 token_compressed_accounts.push(token_account.clone());
1670 self.token_compressed_accounts.insert(0, token_account);
1671 }
1672 } else {
1673 let compressed_account = CompressedAccountWithMerkleContext {
1674 compressed_account: compressed_account.compressed_account.clone(),
1675 merkle_context: MerkleContext {
1676 leaf_index: event.output_leaf_indices[i],
1677 merkle_tree_pubkey: merkle_tree_pubkey.into(),
1678 queue_pubkey: nullifier_queue_pubkey.into(),
1679 prove_by_index: false,
1680 tree_type: merkle_tree.tree_type,
1681 },
1682 };
1683 compressed_accounts.push(compressed_account.clone());
1684 self.compressed_accounts.insert(0, compressed_account);
1685 }
1686 }
1687 None => {
1688 let compressed_account = CompressedAccountWithMerkleContext {
1689 compressed_account: compressed_account.compressed_account.clone(),
1690 merkle_context: MerkleContext {
1691 leaf_index: event.output_leaf_indices[i],
1692 merkle_tree_pubkey: merkle_tree_pubkey.into(),
1693 queue_pubkey: nullifier_queue_pubkey.into(),
1694 prove_by_index: false,
1695 tree_type: merkle_tree.tree_type,
1696 },
1697 };
1698 compressed_accounts.push(compressed_account.clone());
1699 self.compressed_accounts.insert(0, compressed_account);
1700 }
1701 };
1702 let merkle_tree = &mut self.state_merkle_trees.iter_mut().find(|x| {
1703 x.accounts.merkle_tree
1704 == solana_pubkey::Pubkey::from(
1705 event.pubkey_array
1706 [event.output_compressed_accounts[i].merkle_tree_index as usize]
1707 .to_bytes(),
1708 )
1709 });
1710 if merkle_tree.is_some() {
1711 let merkle_tree = merkle_tree.as_mut().unwrap();
1712 let leaf_hash = compressed_account
1713 .compressed_account
1714 .hash(
1715 &event.pubkey_array
1716 [event.output_compressed_accounts[i].merkle_tree_index as usize],
1717 &event.output_leaf_indices[i],
1718 false,
1719 )
1720 .unwrap();
1721 merkle_tree
1722 .merkle_tree
1723 .append(&leaf_hash)
1724 .expect("insert failed");
1725 } else {
1726 let merkle_tree = &mut self
1727 .state_merkle_trees
1728 .iter_mut()
1729 .find(|x| {
1730 x.accounts.nullifier_queue
1731 == solana_pubkey::Pubkey::from(
1732 event.pubkey_array[event.output_compressed_accounts[i]
1733 .merkle_tree_index
1734 as usize]
1735 .to_bytes(),
1736 )
1737 })
1738 .unwrap();
1739
1740 merkle_tree.output_queue_elements.push((
1741 event.output_compressed_account_hashes[i],
1742 event.output_leaf_indices[i].into(),
1743 ));
1744 }
1745 }
1746 if event.input_compressed_account_hashes.len() > i {
1747 let tx_hash: [u8; 32] = create_tx_hash(
1748 &event.input_compressed_account_hashes,
1749 &event.output_compressed_account_hashes,
1750 slot,
1751 )
1752 .unwrap();
1753 let hash = event.input_compressed_account_hashes[i];
1754 let index = self
1755 .compressed_accounts
1756 .iter()
1757 .position(|x| x.hash().unwrap() == hash);
1758 let (leaf_index, merkle_tree_pubkey) = if let Some(index) = index {
1759 self.nullified_compressed_accounts
1760 .push(self.compressed_accounts[index].clone());
1761 let leaf_index = self.compressed_accounts[index].merkle_context.leaf_index;
1762 let merkle_tree_pubkey = self.compressed_accounts[index]
1763 .merkle_context
1764 .merkle_tree_pubkey;
1765 if let Some(address) = self.compressed_accounts[index].compressed_account.address {
1766 input_addresses.push(address);
1767 }
1768 self.compressed_accounts.remove(index);
1769 (Some(leaf_index), Some(merkle_tree_pubkey))
1770 } else if let Some(index) = self
1771 .token_compressed_accounts
1772 .iter()
1773 .position(|x| x.compressed_account.hash().unwrap() == hash)
1774 {
1775 self.token_nullified_compressed_accounts
1776 .push(self.token_compressed_accounts[index].clone());
1777 let leaf_index = self.token_compressed_accounts[index]
1778 .compressed_account
1779 .merkle_context
1780 .leaf_index;
1781 let merkle_tree_pubkey = self.token_compressed_accounts[index]
1782 .compressed_account
1783 .merkle_context
1784 .merkle_tree_pubkey;
1785 self.token_compressed_accounts.remove(index);
1786 (Some(leaf_index), Some(merkle_tree_pubkey))
1787 } else {
1788 (None, None)
1789 };
1790 if let Some(leaf_index) = leaf_index {
1791 let merkle_tree_pubkey = merkle_tree_pubkey.unwrap();
1792 let bundle =
1793 &mut <TestIndexer as TestIndexerExtensions>::get_state_merkle_trees_mut(self)
1794 .iter_mut()
1795 .find(|x| {
1796 x.accounts.merkle_tree
1797 == solana_pubkey::Pubkey::from(merkle_tree_pubkey.to_bytes())
1798 })
1799 .unwrap();
1800 if bundle.tree_type == TreeType::StateV2 {
1802 let leaf_hash = event.input_compressed_account_hashes[i];
1803 bundle.input_leaf_indices.push(LeafIndexInfo {
1804 leaf_index,
1805 leaf: leaf_hash,
1806 tx_hash,
1807 });
1808 }
1809 } else {
1810 println!("Test indexer didn't find input compressed accounts to nullify");
1811 }
1812 }
1813 if !new_addresses.is_empty() {
1820 for pubkey in event.pubkey_array.iter() {
1821 if let Some((_, address_merkle_tree)) = self
1822 .address_merkle_trees
1823 .iter_mut()
1824 .enumerate()
1825 .find(|(_, x)| {
1826 x.accounts.merkle_tree == solana_pubkey::Pubkey::from(pubkey.to_bytes())
1827 })
1828 {
1829 address_merkle_tree
1830 .queue_elements
1831 .append(&mut new_addresses);
1832 }
1833 }
1834 }
1835 }
1836
1837 async fn _get_multiple_new_address_proofs(
1838 &self,
1839 merkle_tree_pubkey: [u8; 32],
1840 addresses: Vec<[u8; 32]>,
1841 full: bool,
1842 ) -> Result<Vec<NewAddressProofWithContext>, IndexerError> {
1843 let mut proofs: Vec<NewAddressProofWithContext> = Vec::new();
1844
1845 for address in addresses.iter() {
1846 info!("Getting new address proof for {:?}", address);
1847 let pubkey = Pubkey::from(merkle_tree_pubkey);
1848 let address_tree_bundle = self
1849 .address_merkle_trees
1850 .iter()
1851 .find(|x| x.accounts.merkle_tree == pubkey)
1852 .unwrap();
1853
1854 let address_biguint = BigUint::from_bytes_be(address.as_slice());
1855 let (old_low_address, _old_low_address_next_value) =
1856 address_tree_bundle.find_low_element_for_nonexistent(&address_biguint)?;
1857 let address_bundle = address_tree_bundle
1858 .new_element_with_low_element_index(old_low_address.index, &address_biguint)?;
1859
1860 let (old_low_address, old_low_address_next_value) =
1861 address_tree_bundle.find_low_element_for_nonexistent(&address_biguint)?;
1862
1863 let low_address_proof =
1865 address_tree_bundle.get_proof_of_leaf(old_low_address.index, full)?;
1866
1867 let low_address_index: u64 = old_low_address.index as u64;
1868 let low_address_value: [u8; 32] =
1869 bigint_to_be_bytes_array(&old_low_address.value).unwrap();
1870 let low_address_next_index: u64 = old_low_address.next_index as u64;
1871 let low_address_next_value: [u8; 32] =
1872 bigint_to_be_bytes_array(&old_low_address_next_value).unwrap();
1873 let proof = NewAddressProofWithContext {
1874 merkle_tree: Pubkey::new_from_array(merkle_tree_pubkey),
1875 low_address_index,
1876 low_address_value,
1877 low_address_next_index,
1878 low_address_next_value,
1879 low_address_proof,
1880 root: address_tree_bundle.root(),
1881 root_seq: address_tree_bundle.sequence_number(),
1882 new_low_element: Some(address_bundle.new_low_element),
1883 new_element: Some(address_bundle.new_element),
1884 new_element_next_value: Some(address_bundle.new_element_next_value),
1885 };
1886 proofs.push(proof);
1887 }
1888 Ok(proofs)
1889 }
1890}
1891
1892impl TestIndexer {
1893 async fn process_inclusion_proofs(
1894 &self,
1895 merkle_tree_pubkeys: &[Pubkey],
1896 accounts: &[[u8; 32]],
1897 ) -> Result<
1898 (
1899 Option<BatchInclusionJsonStruct>,
1900 Option<BatchInclusionJsonStructLegacy>,
1901 Vec<AccountProofInputs>,
1902 ),
1903 IndexerError,
1904 > {
1905 let mut inclusion_proofs = Vec::new();
1906 let mut account_proof_inputs = Vec::new();
1907 let mut height = 0;
1908 let mut queues = vec![];
1909 let mut cpi_contextes = vec![];
1910 let mut tree_types = vec![];
1911 let proof_data: Vec<_> = accounts
1913 .iter()
1914 .zip(merkle_tree_pubkeys.iter())
1915 .map(|(account, &pubkey)| {
1916 let bundle = self
1917 .state_merkle_trees
1918 .iter()
1919 .find(|x| {
1920 x.accounts.merkle_tree == pubkey || x.accounts.nullifier_queue == pubkey
1921 })
1922 .unwrap();
1923 let merkle_tree = &bundle.merkle_tree;
1924 queues.push(bundle.accounts.nullifier_queue);
1925 cpi_contextes.push(bundle.accounts.cpi_context);
1926 tree_types.push(bundle.tree_type);
1927 let leaf_index = merkle_tree.get_leaf_index(account).unwrap();
1928 let proof = merkle_tree.get_proof_of_leaf(leaf_index, true).unwrap();
1929
1930 let proof: Vec<BigInt> = proof.iter().map(|x| BigInt::from_be_bytes(x)).collect();
1932
1933 if height == 0 {
1934 height = merkle_tree.height;
1935 } else {
1936 assert_eq!(height, merkle_tree.height);
1937 }
1938 let root_index = if bundle.tree_type == TreeType::StateV1 {
1939 merkle_tree.get_history_root_index().unwrap()
1940 } else {
1941 merkle_tree.get_history_root_index_v2().unwrap()
1942 };
1943
1944 Ok((leaf_index, proof, merkle_tree.root(), root_index))
1945 })
1946 .collect::<Result<_, IndexerError>>()?;
1947
1948 for (i, (leaf_index, proof, merkle_root, root_index)) in proof_data.into_iter().enumerate()
1950 {
1951 inclusion_proofs.push(InclusionMerkleProofInputs {
1952 root: BigInt::from_be_bytes(merkle_root.as_slice()),
1953 leaf: BigInt::from_be_bytes(&accounts[i]),
1954 path_index: BigInt::from_be_bytes(leaf_index.to_be_bytes().as_slice()),
1955 path_elements: proof,
1956 });
1957
1958 account_proof_inputs.push(AccountProofInputs {
1959 root_index: RootIndex::new_some(root_index),
1960 root: merkle_root,
1961 leaf_index: leaf_index as u64,
1962 hash: accounts[i],
1963 tree_info: light_client::indexer::TreeInfo {
1964 cpi_context: Some(cpi_contextes[i]),
1965 next_tree_info: None,
1966 queue: queues[i],
1967 tree: merkle_tree_pubkeys[i],
1968 tree_type: tree_types[i],
1969 },
1970 });
1971 }
1972
1973 let (batch_inclusion_proof_inputs, legacy) = if height
1974 == DEFAULT_BATCH_STATE_TREE_HEIGHT as usize
1975 {
1976 let inclusion_proof_inputs =
1977 InclusionProofInputs::new(inclusion_proofs.as_slice()).unwrap();
1978 (
1979 Some(BatchInclusionJsonStruct::from_inclusion_proof_inputs(
1980 &inclusion_proof_inputs,
1981 )),
1982 None,
1983 )
1984 } else if height == account_compression::utils::constants::STATE_MERKLE_TREE_HEIGHT as usize
1985 {
1986 let inclusion_proof_inputs = InclusionProofInputsLegacy(inclusion_proofs.as_slice());
1987 (
1988 None,
1989 Some(BatchInclusionJsonStructLegacy::from_inclusion_proof_inputs(
1990 &inclusion_proof_inputs,
1991 )),
1992 )
1993 } else {
1994 return Err(IndexerError::CustomError(
1995 "Unsupported tree height".to_string(),
1996 ));
1997 };
1998
1999 Ok((batch_inclusion_proof_inputs, legacy, account_proof_inputs))
2000 }
2001
2002 async fn process_non_inclusion_proofs(
2003 &self,
2004 address_merkle_tree_pubkeys: &[Pubkey],
2005 addresses: Vec<[u8; 32]>,
2006 ) -> Result<
2007 (
2008 Option<BatchNonInclusionJsonStruct>,
2009 Option<BatchNonInclusionJsonStructLegacy>,
2010 Vec<AddressProofInputs>,
2011 ),
2012 IndexerError,
2013 > {
2014 let mut non_inclusion_proofs = Vec::new();
2015 let mut address_root_indices = Vec::new();
2016 let mut tree_heights = Vec::new();
2017 for (i, address) in addresses.iter().enumerate() {
2018 let address_tree = self
2019 .address_merkle_trees
2020 .iter()
2021 .find(|x| x.accounts.merkle_tree == address_merkle_tree_pubkeys[i])
2022 .unwrap();
2023 tree_heights.push(address_tree.height());
2024
2025 let proof_inputs = address_tree.get_non_inclusion_proof_inputs(address)?;
2026 non_inclusion_proofs.push(proof_inputs);
2027
2028 let (root_index, root, tree_type) = match &address_tree.merkle_tree {
2029 super::address_tree::IndexedMerkleTreeVersion::V1(tree) => (
2030 tree.merkle_tree.get_history_root_index().unwrap() + 1,
2031 tree.merkle_tree.root(),
2032 TreeType::AddressV1,
2033 ),
2034 super::address_tree::IndexedMerkleTreeVersion::V2(tree) => (
2035 tree.merkle_tree.get_history_root_index_v2().unwrap(),
2036 tree.merkle_tree.root(),
2037 TreeType::AddressV2,
2038 ),
2039 };
2040 address_root_indices.push(AddressProofInputs {
2041 root_index,
2042 root,
2043 address: *address,
2044 tree_info: light_client::indexer::TreeInfo {
2045 cpi_context: None,
2046 next_tree_info: None,
2047 queue: address_tree.accounts.queue,
2048 tree: address_tree.accounts.merkle_tree,
2049 tree_type,
2050 },
2051 });
2052 }
2053 if tree_heights.iter().any(|&x| x != tree_heights[0]) {
2055 return Err(IndexerError::CustomError(format!(
2056 "All address merkle trees must have the same height {:?}",
2057 tree_heights
2058 )));
2059 }
2060 let (batch_non_inclusion_proof_inputs, batch_non_inclusion_proof_inputs_legacy) =
2061 if tree_heights[0] == 26 {
2062 let non_inclusion_proof_inputs =
2063 NonInclusionProofInputsLegacy::new(non_inclusion_proofs.as_slice());
2064 (
2065 None,
2066 Some(
2067 BatchNonInclusionJsonStructLegacy::from_non_inclusion_proof_inputs(
2068 &non_inclusion_proof_inputs,
2069 ),
2070 ),
2071 )
2072 } else if tree_heights[0] == 40 {
2073 let non_inclusion_proof_inputs =
2074 NonInclusionProofInputs::new(non_inclusion_proofs.as_slice()).unwrap();
2075 (
2076 Some(
2077 BatchNonInclusionJsonStruct::from_non_inclusion_proof_inputs(
2078 &non_inclusion_proof_inputs,
2079 ),
2080 ),
2081 None,
2082 )
2083 } else {
2084 return Err(IndexerError::CustomError(
2085 "Unsupported tree height".to_string(),
2086 ));
2087 };
2088 Ok((
2089 batch_non_inclusion_proof_inputs,
2090 batch_non_inclusion_proof_inputs_legacy,
2091 address_root_indices,
2092 ))
2093 }
2094}
2095
2096impl TestIndexer {
2097 async fn _get_validity_proof_v1_implementation(
2098 &self,
2099 hashes: Vec<[u8; 32]>,
2100 new_addresses_with_trees: Vec<AddressWithTree>,
2101 ) -> Result<ValidityProofWithContext, IndexerError> {
2102 let mut state_merkle_tree_pubkeys = Vec::new();
2103
2104 for hash in hashes.iter() {
2105 let account = self.get_compressed_account_by_hash(*hash, None).await?;
2106 let account_data = account.value.ok_or(IndexerError::AccountNotFound)?;
2107 state_merkle_tree_pubkeys.push(account_data.tree_info.tree);
2108 }
2109
2110 let state_merkle_tree_pubkeys = if state_merkle_tree_pubkeys.is_empty() {
2111 None
2112 } else {
2113 Some(state_merkle_tree_pubkeys)
2114 };
2115 let hashes = if hashes.is_empty() {
2116 None
2117 } else {
2118 Some(hashes)
2119 };
2120 let new_addresses = if new_addresses_with_trees.is_empty() {
2121 None
2122 } else {
2123 Some(
2124 new_addresses_with_trees
2125 .iter()
2126 .map(|x| x.address)
2127 .collect::<Vec<[u8; 32]>>(),
2128 )
2129 };
2130 let address_merkle_tree_pubkeys = if new_addresses_with_trees.is_empty() {
2131 None
2132 } else {
2133 Some(
2134 new_addresses_with_trees
2135 .iter()
2136 .map(|x| x.tree)
2137 .collect::<Vec<Pubkey>>(),
2138 )
2139 };
2140
2141 {
2142 let compressed_accounts = hashes;
2143 if compressed_accounts.is_some()
2144 && ![1usize, 2usize, 3usize, 4usize, 8usize]
2145 .contains(&compressed_accounts.as_ref().unwrap().len())
2146 {
2147 return Err(IndexerError::CustomError(format!(
2148 "compressed_accounts must be of length 1, 2, 3, 4 or 8 != {}",
2149 compressed_accounts.unwrap().len()
2150 )));
2151 }
2152 if new_addresses.is_some()
2153 && ![1usize, 2usize, 3usize, 4usize, 8usize]
2154 .contains(&new_addresses.as_ref().unwrap().len())
2155 {
2156 return Err(IndexerError::CustomError(format!(
2157 "new_addresses must be of length 1, 2, 3, 4 or 8 != {}",
2158 new_addresses.unwrap().len()
2159 )));
2160 }
2161 let client = Client::new();
2162 let (account_proof_inputs, address_proof_inputs, json_payload) =
2163 match (compressed_accounts, new_addresses) {
2164 (Some(accounts), None) => {
2165 let (payload, payload_legacy, indices) = self
2166 .process_inclusion_proofs(
2167 &state_merkle_tree_pubkeys.unwrap(),
2168 &accounts,
2169 )
2170 .await?;
2171 if let Some(payload) = payload {
2172 (indices, Vec::new(), payload.to_string())
2173 } else {
2174 (indices, Vec::new(), payload_legacy.unwrap().to_string())
2175 }
2176 }
2177 (None, Some(addresses)) => {
2178 let (payload, payload_legacy, indices) = self
2179 .process_non_inclusion_proofs(
2180 address_merkle_tree_pubkeys.unwrap().as_slice(),
2181 addresses,
2182 )
2183 .await?;
2184 let payload_string = if let Some(payload) = payload {
2185 payload.to_string()
2186 } else {
2187 payload_legacy.unwrap().to_string()
2188 };
2189 (Vec::new(), indices, payload_string)
2190 }
2191 (Some(accounts), Some(addresses)) => {
2192 let (inclusion_payload, inclusion_payload_legacy, inclusion_indices) = self
2193 .process_inclusion_proofs(
2194 &state_merkle_tree_pubkeys.unwrap(),
2195 &accounts,
2196 )
2197 .await?;
2198
2199 let (
2200 non_inclusion_payload,
2201 non_inclusion_payload_legacy,
2202 non_inclusion_indices,
2203 ) = self
2204 .process_non_inclusion_proofs(
2205 address_merkle_tree_pubkeys.unwrap().as_slice(),
2206 addresses,
2207 )
2208 .await?;
2209 let json_payload = if let Some(non_inclusion_payload) =
2210 non_inclusion_payload
2211 {
2212 let public_input_hash = BigInt::from_bytes_be(
2213 num_bigint::Sign::Plus,
2214 &create_hash_chain_from_slice(&[
2215 bigint_to_u8_32(
2216 &string_to_big_int(
2217 &inclusion_payload.as_ref().unwrap().public_input_hash,
2218 )
2219 .unwrap(),
2220 )
2221 .unwrap(),
2222 bigint_to_u8_32(
2223 &string_to_big_int(
2224 &non_inclusion_payload.public_input_hash,
2225 )
2226 .unwrap(),
2227 )
2228 .unwrap(),
2229 ])
2230 .unwrap(),
2231 );
2232
2233 CombinedJsonStruct {
2234 circuit_type: ProofType::Combined.to_string(),
2235 state_tree_height: DEFAULT_BATCH_STATE_TREE_HEIGHT,
2236 address_tree_height: DEFAULT_BATCH_ADDRESS_TREE_HEIGHT,
2237 public_input_hash: big_int_to_string(&public_input_hash),
2238 inclusion: inclusion_payload.unwrap().inputs,
2239 non_inclusion: non_inclusion_payload.inputs,
2240 }
2241 .to_string()
2242 } else if let Some(non_inclusion_payload) = non_inclusion_payload_legacy {
2243 CombinedJsonStructLegacy {
2244 circuit_type: ProofType::Combined.to_string(),
2245 state_tree_height: 26,
2246 address_tree_height: 26,
2247 inclusion: inclusion_payload_legacy.unwrap().inputs,
2248 non_inclusion: non_inclusion_payload.inputs,
2249 }
2250 .to_string()
2251 } else {
2252 panic!("Unsupported tree height")
2253 };
2254 (inclusion_indices, non_inclusion_indices, json_payload)
2255 }
2256 _ => {
2257 panic!(
2258 "At least one of compressed_accounts or new_addresses must be provided"
2259 )
2260 }
2261 };
2262
2263 let mut retries = 3;
2264 while retries > 0 {
2265 let response_result = client
2266 .post(format!("{}{}", SERVER_ADDRESS, PROVE_PATH))
2267 .header("Content-Type", "text/plain; charset=utf-8")
2268 .body(json_payload.clone())
2269 .send()
2270 .await;
2271 if let Ok(response_result) = response_result {
2272 if response_result.status().is_success() {
2273 let body = response_result.text().await.unwrap();
2274 let proof_json = deserialize_gnark_proof_json(&body).unwrap();
2275 let (proof_a, proof_b, proof_c) = proof_from_json_struct(proof_json);
2276 let (proof_a, proof_b, proof_c) =
2277 compress_proof(&proof_a, &proof_b, &proof_c);
2278 return Ok(ValidityProofWithContext {
2279 accounts: account_proof_inputs,
2280 addresses: address_proof_inputs,
2281 proof: CompressedProof {
2282 a: proof_a,
2283 b: proof_b,
2284 c: proof_c,
2285 }
2286 .into(),
2287 });
2288 }
2289 } else {
2290 println!("Error: {:#?}", response_result);
2291 tokio::time::sleep(Duration::from_secs(5)).await;
2292 retries -= 1;
2293 }
2294 }
2295 Err(IndexerError::CustomError(
2296 "Failed to get proof from server".to_string(),
2297 ))
2298 }
2299 }
2300}