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