light_test_utils/
e2e_test_env.rs

1// Flow:
2// init indexer
3// init first keypair
4// init crank
5// vec of public Merkle tree NF queue pairs
6// vec of public address Mt and queue pairs
7// for i in rounds
8//   randomly add new keypair
9// for every keypair randomly select whether it does an action
10
11// Architecture:
12// - bundle trees, indexer etc in a E2ETestEnv struct
13// - methods:
14// 	// bundles all general actions
15//   - activate general actions
16//   // bundles all keypair actions
17//   - activate keypair actions
18// 	// calls general and keypair actions
19//   - execute round
20//   // every action takes a probability as input
21//   // if you want to execute the action on purpose pass 1
22//   - method for every action
23//  - add action activation config with default configs
24//    - all enabled
25//    - only spl, only sol, etc
26//  Forester struct
27//  - payer keypair, authority keypair
28//  -methods
29//   - empty nullifier queue
30//   - empty address queue
31//   - rollover Merkle tree
32//   - rollover address Merkle tree
33
34// keypair actions:
35// safeguard every action in case of no balance
36// 1. compress sol
37// 2. decompress sol
38// 2. transfer sol
39// 3. compress spl
40// 4. decompress spl
41// 5. mint spl
42// 6. transfer spl
43
44// general actions:
45// add keypair
46// create new state Mt
47// create new address Mt
48
49// extension:
50// keypair actions:
51// - create pda
52// - escrow tokens
53// - delegate, revoke, delegated transaction
54
55// general actions:
56// - create new program owned state Merkle tree and queue
57// - create new program owned address Merkle tree and queue
58
59// minimal start
60// struct with env and test-indexer
61// only spl transactions
62
63// second pr
64// refactor sol tests to functions that can be reused
65
66// TODO: implement traits for context object and indexer that we can implement with an rpc as well
67// context trait: send_transaction -> return transaction result, get_account_info -> return account info
68// indexer trait: get_compressed_accounts_by_owner -> return compressed accounts,
69// refactor all tests to work with that so that we can run all tests with a test validator and concurrency
70
71use light_compressed_token::token_data::AccountState;
72use light_registry::protocol_config::state::{ProtocolConfig, ProtocolConfigPda};
73use light_registry::sdk::create_finalize_registration_instruction;
74use light_registry::utils::get_protocol_config_pda_address;
75use light_registry::ForesterConfig;
76use log::info;
77use num_bigint::{BigUint, RandBigInt};
78use num_traits::Num;
79use rand::distributions::uniform::{SampleRange, SampleUniform};
80use rand::prelude::SliceRandom;
81use rand::rngs::{StdRng, ThreadRng};
82use rand::{Rng, RngCore, SeedableRng};
83use solana_sdk::pubkey::Pubkey;
84use solana_sdk::signature::Keypair;
85use solana_sdk::signature::Signature;
86use solana_sdk::signer::{SeedDerivable, Signer};
87use spl_token::solana_program::native_token::LAMPORTS_PER_SOL;
88
89use crate::address_tree_rollover::{
90    assert_rolled_over_address_merkle_tree_and_queue,
91    perform_address_merkle_tree_roll_over_forester, perform_state_merkle_tree_roll_over_forester,
92};
93use crate::assert_epoch::{
94    assert_finalized_epoch_registration, assert_report_work, fetch_epoch_and_forester_pdas,
95};
96use crate::rpc::ProgramTestRpcConnection;
97use crate::spl::{
98    approve_test, burn_test, compress_test, compressed_transfer_test, create_mint_helper,
99    create_token_account, decompress_test, freeze_test, mint_tokens_helper, revoke_test, thaw_test,
100};
101use crate::state_tree_rollover::assert_rolled_over_pair;
102use crate::system_program::{
103    compress_sol_test, create_addresses_test, decompress_sol_test, transfer_compressed_sol_test,
104};
105use crate::test_env::{
106    create_address_merkle_tree_and_queue_account, create_state_merkle_tree_and_queue_account,
107    EnvAccounts,
108};
109use crate::test_forester::{empty_address_queue_test, nullify_compressed_accounts};
110use account_compression::utils::constants::{
111    STATE_MERKLE_TREE_CANOPY_DEPTH, STATE_MERKLE_TREE_HEIGHT,
112};
113use account_compression::{
114    AddressMerkleTreeConfig, AddressQueueConfig, NullifierQueueConfig, StateMerkleTreeConfig,
115    SAFETY_MARGIN,
116};
117use forester_utils::address_merkle_tree_config::{
118    address_tree_ready_for_rollover, state_tree_ready_for_rollover,
119};
120use forester_utils::forester_epoch::{Epoch, Forester, TreeAccounts, TreeType};
121use forester_utils::indexer::{
122    AddressMerkleTreeAccounts, AddressMerkleTreeBundle, Indexer, StateMerkleTreeAccounts,
123    StateMerkleTreeBundle, TokenDataWithContext,
124};
125use forester_utils::registry::register_test_forester;
126use forester_utils::{airdrop_lamports, AccountZeroCopy};
127use light_hasher::Poseidon;
128use light_indexed_merkle_tree::HIGHEST_ADDRESS_PLUS_ONE;
129use light_indexed_merkle_tree::{array::IndexedArray, reference::IndexedMerkleTree};
130use light_system_program::sdk::compressed_account::CompressedAccountWithMerkleContext;
131use light_utils::bigint::bigint_to_be_bytes_array;
132use light_utils::rand::gen_prime;
133
134use crate::indexer::TestIndexer;
135use light_client::rpc::errors::RpcError;
136use light_client::rpc::RpcConnection;
137use light_client::transaction_params::{FeeConfig, TransactionParams};
138
139pub struct User {
140    pub keypair: Keypair,
141    // Vector of (mint, token account)
142    pub token_accounts: Vec<(Pubkey, Pubkey)>,
143}
144
145#[derive(Debug, Default)]
146pub struct Stats {
147    pub spl_transfers: u64,
148    pub mints: u64,
149    pub spl_decompress: u64,
150    pub spl_compress: u64,
151    pub sol_transfers: u64,
152    pub sol_decompress: u64,
153    pub sol_compress: u64,
154    pub create_address: u64,
155    pub create_pda: u64,
156    pub create_state_mt: u64,
157    pub create_address_mt: u64,
158    pub rolledover_state_trees: u64,
159    pub rolledover_address_trees: u64,
160    pub spl_approved: u64,
161    pub spl_revoked: u64,
162    pub spl_burned: u64,
163    pub spl_frozen: u64,
164    pub spl_thawed: u64,
165    pub registered_foresters: u64,
166    pub created_foresters: u64,
167    pub work_reported: u64,
168    pub finalized_registrations: u64,
169}
170
171impl Stats {
172    pub fn print(&self, users: u64) {
173        println!("Stats:");
174        println!("Users {}", users);
175        println!("Mints {}", self.mints);
176        println!("Spl transfers {}", self.spl_transfers);
177        println!("Spl decompress {}", self.spl_decompress);
178        println!("Spl compress {}", self.spl_compress);
179        println!("Sol transfers {}", self.sol_transfers);
180        println!("Sol decompress {}", self.sol_decompress);
181        println!("Sol compress {}", self.sol_compress);
182        println!("Create address {}", self.create_address);
183        println!("Create pda {}", self.create_pda);
184        println!("Create state mt {}", self.create_state_mt);
185        println!("Create address mt {}", self.create_address_mt);
186        println!("Rolled over state trees {}", self.rolledover_state_trees);
187        println!(
188            "Rolled over address trees {}",
189            self.rolledover_address_trees
190        );
191        println!("Spl approved {}", self.spl_approved);
192        println!("Spl revoked {}", self.spl_revoked);
193        println!("Spl burned {}", self.spl_burned);
194        println!("Spl frozen {}", self.spl_frozen);
195        println!("Spl thawed {}", self.spl_thawed);
196        println!("Registered foresters {}", self.registered_foresters);
197        println!("Created foresters {}", self.created_foresters);
198        println!("Work reported {}", self.work_reported);
199        println!("Finalized registrations {}", self.finalized_registrations);
200    }
201}
202pub async fn init_program_test_env(
203    rpc: ProgramTestRpcConnection,
204    env_accounts: &EnvAccounts,
205) -> E2ETestEnv<ProgramTestRpcConnection, TestIndexer<ProgramTestRpcConnection>> {
206    let indexer: TestIndexer<ProgramTestRpcConnection> = TestIndexer::init_from_env(
207        &env_accounts.forester.insecure_clone(),
208        env_accounts,
209        KeypairActionConfig::all_default().inclusion(),
210        KeypairActionConfig::all_default().non_inclusion(),
211    )
212    .await;
213
214    E2ETestEnv::<ProgramTestRpcConnection, TestIndexer<ProgramTestRpcConnection>>::new(
215        rpc,
216        indexer,
217        env_accounts,
218        KeypairActionConfig::all_default(),
219        GeneralActionConfig::default(),
220        10,
221        None,
222    )
223    .await
224}
225
226#[derive(Debug, PartialEq)]
227pub struct TestForester {
228    keypair: Keypair,
229    forester: Forester,
230    is_registered: Option<u64>,
231}
232
233pub struct E2ETestEnv<R: RpcConnection, I: Indexer<R>> {
234    pub payer: Keypair,
235    pub governance_keypair: Keypair,
236    pub indexer: I,
237    pub users: Vec<User>,
238    pub mints: Vec<Pubkey>,
239    pub foresters: Vec<TestForester>,
240    pub rpc: R,
241    pub keypair_action_config: KeypairActionConfig,
242    pub general_action_config: GeneralActionConfig,
243    pub round: u64,
244    pub rounds: u64,
245    pub rng: StdRng,
246    pub stats: Stats,
247    pub epoch: u64,
248    pub slot: u64,
249    /// Forester struct is reused but not used for foresting here
250    /// Epoch config keeps track of the ongong epochs.
251    pub epoch_config: Forester,
252    pub protocol_config: ProtocolConfig,
253    pub registration_epoch: u64,
254}
255
256impl<R: RpcConnection, I: Indexer<R>> E2ETestEnv<R, I>
257where
258    R: RpcConnection,
259    I: Indexer<R>,
260{
261    pub async fn new(
262        mut rpc: R,
263        mut indexer: I,
264        env_accounts: &EnvAccounts,
265        keypair_action_config: KeypairActionConfig,
266        general_action_config: GeneralActionConfig,
267        rounds: u64,
268        seed: Option<u64>,
269    ) -> Self {
270        let payer = rpc.get_payer().insecure_clone();
271
272        airdrop_lamports(&mut rpc, &payer.pubkey(), 1_000_000_000_000)
273            .await
274            .unwrap();
275
276        airdrop_lamports(&mut rpc, &env_accounts.forester.pubkey(), 1_000_000_000_000)
277            .await
278            .unwrap();
279        let mut thread_rng = ThreadRng::default();
280        let random_seed = thread_rng.next_u64();
281        let seed: u64 = seed.unwrap_or(random_seed);
282        // Keep this print so that in case the test fails
283        // we can use the seed to reproduce the error.
284        println!("\n\ne2e test seed {}\n\n", seed);
285        let mut rng = StdRng::seed_from_u64(seed);
286        let user = Self::create_user(&mut rng, &mut rpc).await;
287        let mint = create_mint_helper(&mut rpc, &payer).await;
288        mint_tokens_helper(
289            &mut rpc,
290            &mut indexer,
291            &env_accounts.merkle_tree_pubkey,
292            &payer,
293            &mint,
294            vec![100_000_000; 1],
295            vec![user.keypair.pubkey()],
296        )
297        .await;
298        let protocol_config_pda_address = get_protocol_config_pda_address().0;
299        println!("here");
300        let protocol_config = rpc
301            .get_anchor_account::<ProtocolConfigPda>(&protocol_config_pda_address)
302            .await
303            .unwrap()
304            .unwrap()
305            .config;
306        // TODO: add clear test env enum
307        // register foresters is only compatible with ProgramTest environment
308        let (foresters, epoch_config) =
309            if let Some(registered_epoch) = env_accounts.forester_epoch.as_ref() {
310                let _forester = Forester {
311                    registration: registered_epoch.clone(),
312                    active: registered_epoch.clone(),
313                    ..Default::default()
314                };
315                // Forester epoch account is assumed to exist (is inited with test program deployment)
316                let forester = TestForester {
317                    keypair: env_accounts.forester.insecure_clone(),
318                    forester: _forester.clone(),
319                    is_registered: Some(0),
320                };
321                (vec![forester], _forester)
322            } else {
323                (Vec::<TestForester>::new(), Forester::default())
324            };
325        Self {
326            payer,
327            indexer,
328            users: vec![user],
329            rpc,
330            keypair_action_config,
331            general_action_config,
332            round: 0,
333            rounds,
334            rng,
335            mints: vec![],
336            stats: Stats::default(),
337            foresters,
338            registration_epoch: 0,
339            epoch: 0,
340            slot: 0,
341            epoch_config,
342            protocol_config,
343            governance_keypair: env_accounts.governance_authority.insecure_clone(),
344        }
345    }
346
347    /// Creates a new user with a random keypair and 100 sol
348    pub async fn create_user(rng: &mut StdRng, rpc: &mut R) -> User {
349        let keypair: Keypair = Keypair::from_seed(&[rng.gen_range(0..255); 32]).unwrap();
350
351        rpc.airdrop_lamports(&keypair.pubkey(), LAMPORTS_PER_SOL * 5000)
352            .await
353            .unwrap();
354        User {
355            keypair,
356            token_accounts: vec![],
357        }
358    }
359
360    pub async fn get_balance(&mut self, pubkey: &Pubkey) -> u64 {
361        self.rpc.get_balance(pubkey).await.unwrap()
362    }
363
364    pub async fn execute_rounds(&mut self) {
365        for _ in 0..=self.rounds {
366            self.execute_round().await;
367        }
368    }
369
370    pub async fn execute_round(&mut self) {
371        println!("\n------------------------------------------------------\n");
372        println!("Round: {}", self.round);
373        self.stats.print(self.users.len() as u64);
374
375        // TODO: check at the beginning of the round that the Merkle trees are in sync
376        let len = self.users.len();
377        for i in 0..len {
378            self.activate_keypair_actions(&self.users[i].keypair.pubkey())
379                .await;
380        }
381        self.activate_general_actions().await;
382        self.round += 1;
383    }
384
385    /// 1. Add a new keypair
386    /// 2. Create a new state Merkle tree
387    pub async fn activate_general_actions(&mut self) {
388        // If we want to test rollovers we set the threshold to 0 for all newly created trees
389        let rollover_threshold = if self.general_action_config.rollover.is_some() {
390            Some(0)
391        } else {
392            None
393        };
394        if self
395            .rng
396            .gen_bool(self.general_action_config.add_keypair.unwrap_or_default())
397        {
398            let user = Self::create_user(&mut self.rng, &mut self.rpc).await;
399            self.users.push(user);
400        }
401
402        if self.rng.gen_bool(
403            self.general_action_config
404                .create_state_mt
405                .unwrap_or_default(),
406        ) {
407            self.create_state_tree(rollover_threshold).await;
408            self.stats.create_state_mt += 1;
409        }
410
411        if self.rng.gen_bool(
412            self.general_action_config
413                .create_address_mt
414                .unwrap_or_default(),
415        ) {
416            self.create_address_tree(rollover_threshold).await;
417            self.stats.create_address_mt += 1;
418        }
419
420        if self.rng.gen_bool(
421            self.general_action_config
422                .nullify_compressed_accounts
423                .unwrap_or_default(),
424        ) {
425            for state_tree_bundle in self.indexer.get_state_merkle_trees_mut().iter_mut() {
426                println!("\n --------------------------------------------------\n\t\t NULLIFYING LEAVES\n --------------------------------------------------");
427                // find forester which is eligible this slot for this tree
428                if let Some(payer) = Self::get_eligible_forester_for_queue(
429                    &state_tree_bundle.accounts.nullifier_queue,
430                    &self.foresters,
431                    self.slot,
432                ) {
433                    // TODO: add newly addeded trees to foresters
434                    nullify_compressed_accounts(
435                        &mut self.rpc,
436                        &payer,
437                        state_tree_bundle,
438                        self.epoch,
439                        false,
440                    )
441                    .await
442                    .unwrap();
443                } else {
444                    println!("No forester found for nullifier queue");
445                };
446            }
447        }
448
449        if self.rng.gen_bool(
450            self.general_action_config
451                .empty_address_queue
452                .unwrap_or_default(),
453        ) {
454            for address_merkle_tree_bundle in self.indexer.get_address_merkle_trees_mut().iter_mut()
455            {
456                // find forester which is eligible this slot for this tree
457                if let Some(payer) = Self::get_eligible_forester_for_queue(
458                    &address_merkle_tree_bundle.accounts.queue,
459                    &self.foresters,
460                    self.slot,
461                ) {
462                    println!("\n --------------------------------------------------\n\t\t Empty Address Queue\n --------------------------------------------------");
463                    println!("epoch {}", self.epoch);
464                    println!("forester {}", payer.pubkey());
465                    // TODO: add newly addeded trees to foresters
466                    empty_address_queue_test(
467                        &payer,
468                        &mut self.rpc,
469                        address_merkle_tree_bundle,
470                        false,
471                        self.epoch,
472                        false,
473                    )
474                    .await
475                    .unwrap();
476                } else {
477                    println!("No forester found for address queue");
478                };
479            }
480        }
481
482        for index in 0..self.indexer.get_state_merkle_trees().len() {
483            let is_read_for_rollover = state_tree_ready_for_rollover(
484                &mut self.rpc,
485                self.indexer.get_state_merkle_trees()[index]
486                    .accounts
487                    .merkle_tree,
488            )
489            .await;
490            if self
491                .rng
492                .gen_bool(self.general_action_config.rollover.unwrap_or_default())
493                && is_read_for_rollover
494            {
495                println!("\n --------------------------------------------------\n\t\t Rollover State Merkle Tree\n --------------------------------------------------");
496                // find forester which is eligible this slot for this tree
497                if let Some(payer) = Self::get_eligible_forester_for_queue(
498                    &self.indexer.get_state_merkle_trees()[index]
499                        .accounts
500                        .nullifier_queue,
501                    &self.foresters,
502                    self.slot,
503                ) {
504                    self.rollover_state_merkle_tree_and_queue(index, &payer, self.epoch)
505                        .await
506                        .unwrap();
507                    self.stats.rolledover_state_trees += 1;
508                }
509            }
510        }
511
512        for index in 0..self.indexer.get_address_merkle_trees().len() {
513            let is_read_for_rollover = address_tree_ready_for_rollover(
514                &mut self.rpc,
515                self.indexer.get_address_merkle_trees()[index]
516                    .accounts
517                    .merkle_tree,
518            )
519            .await;
520            if self
521                .rng
522                .gen_bool(self.general_action_config.rollover.unwrap_or_default())
523                && is_read_for_rollover
524            {
525                // find forester which is eligible this slot for this tree
526                if let Some(payer) = Self::get_eligible_forester_for_queue(
527                    &self.indexer.get_address_merkle_trees()[index]
528                        .accounts
529                        .queue,
530                    &self.foresters,
531                    self.slot,
532                ) {
533                    println!("\n --------------------------------------------------\n\t\t Rollover Address Merkle Tree\n --------------------------------------------------");
534                    self.rollover_address_merkle_tree_and_queue(index, &payer, self.epoch)
535                        .await
536                        .unwrap();
537                    self.stats.rolledover_address_trees += 1;
538                }
539            }
540        }
541
542        if self
543            .rng
544            .gen_bool(self.general_action_config.add_forester.unwrap_or_default())
545        {
546            println!("\n --------------------------------------------------\n\t\t Add Forester\n --------------------------------------------------");
547            let forester = TestForester {
548                keypair: Keypair::new(),
549                forester: Forester::default(),
550                is_registered: None,
551            };
552            let forester_config = ForesterConfig {
553                fee: self.rng.gen_range(0..=100),
554            };
555            register_test_forester(
556                &mut self.rpc,
557                &self.governance_keypair,
558                &forester.keypair.pubkey(),
559                forester_config,
560            )
561            .await
562            .unwrap();
563            self.foresters.push(forester);
564            self.stats.created_foresters += 1;
565        }
566
567        // advance to next light slot and perform forester epoch actions
568        if !self.general_action_config.disable_epochs {
569            println!("\n --------------------------------------------------\n\t\t Start Epoch Actions \n --------------------------------------------------");
570
571            let current_solana_slot = self.rpc.get_slot().await.unwrap();
572            let current_light_slot = self
573                .protocol_config
574                .get_current_active_epoch_progress(current_solana_slot)
575                / self.protocol_config.slot_length;
576            // If slot didn't change, advance to next slot
577            // if current_light_slot != self.slot {
578            let new_slot = current_solana_slot + self.protocol_config.slot_length;
579            println!("advanced slot from {} to {}", self.slot, current_light_slot);
580            println!("solana slot from {} to {}", current_solana_slot, new_slot);
581            self.rpc.warp_to_slot(new_slot).await.unwrap();
582
583            self.slot = current_light_slot + 1;
584
585            let current_solana_slot = self.rpc.get_slot().await.unwrap();
586            // need to detect whether new registration phase started
587            let current_registration_epoch = self
588                .protocol_config
589                .get_latest_register_epoch(current_solana_slot)
590                .unwrap();
591            // If reached new registration phase register all foresters
592            if current_registration_epoch != self.registration_epoch {
593                println!("\n --------------------------------------------------\n\t\t Register Foresters for new Epoch \n --------------------------------------------------");
594
595                self.registration_epoch = current_registration_epoch;
596                println!("new register epoch {}", self.registration_epoch);
597                println!("num foresters {}", self.foresters.len());
598                for forester in self.foresters.iter_mut() {
599                    println!(
600                        "registered forester {} for epoch {}",
601                        forester.keypair.pubkey(),
602                        self.registration_epoch
603                    );
604
605                    let registered_epoch =
606                        Epoch::register(&mut self.rpc, &self.protocol_config, &forester.keypair)
607                            .await
608                            .unwrap()
609                            .unwrap();
610                    println!("registered_epoch {:?}", registered_epoch.phases);
611                    forester.forester.registration = registered_epoch;
612                    if forester.is_registered.is_none() {
613                        forester.is_registered = Some(self.registration_epoch);
614                    }
615                    self.stats.registered_foresters += 1;
616                }
617            }
618
619            let current_active_epoch = self
620                .protocol_config
621                .get_current_active_epoch(current_solana_slot)
622                .unwrap();
623            // If reached new active epoch
624            // 1. move epoch in every forester to report work epoch
625            // 2. report work for every forester
626            // 3. finalize registration for every forester
627            #[allow(clippy::comparison_chain)]
628            if current_active_epoch > self.epoch {
629                self.slot = current_light_slot;
630                self.epoch = current_active_epoch;
631                // 1. move epoch in every forester to report work epoch
632                for forester in self.foresters.iter_mut() {
633                    if forester.is_registered.is_none() {
634                        continue;
635                    }
636                    forester.forester.switch_to_report_work();
637                }
638                println!("\n --------------------------------------------------\n\t\t Report Work \n --------------------------------------------------");
639
640                // 2. report work for every forester
641                for forester in self.foresters.iter_mut() {
642                    if forester.is_registered.is_none() {
643                        continue;
644                    }
645                    println!("report work for forester {}", forester.keypair.pubkey());
646                    println!(
647                        "forester.forester.report_work.forester_epoch_pda {}",
648                        forester.forester.report_work.forester_epoch_pda
649                    );
650                    println!(
651                        "forester.forester.report_work.epoch_pda {}",
652                        forester.forester.report_work.epoch_pda
653                    );
654
655                    let (pre_forester_epoch_pda, pre_epoch_pda) = fetch_epoch_and_forester_pdas(
656                        &mut self.rpc,
657                        &forester.forester.report_work.forester_epoch_pda,
658                        &forester.forester.report_work.epoch_pda,
659                    )
660                    .await;
661                    forester
662                        .forester
663                        .report_work(&mut self.rpc, &forester.keypair)
664                        .await
665                        .unwrap();
666                    println!("reported work");
667                    assert_report_work(
668                        &mut self.rpc,
669                        &forester.forester.report_work.forester_epoch_pda,
670                        &forester.forester.report_work.epoch_pda,
671                        pre_forester_epoch_pda,
672                        pre_epoch_pda,
673                    )
674                    .await;
675                    self.stats.work_reported += 1;
676                }
677
678                // 3. finalize registration for every forester
679                println!("\n --------------------------------------------------\n\t\t Finalize Registration \n --------------------------------------------------");
680
681                // 3.1 get tree accounts
682                // TODO: use TreeAccounts in TestIndexer
683                let mut tree_accounts = self
684                    .indexer
685                    .get_state_merkle_trees()
686                    .iter()
687                    .map(|state_merkle_tree_bundle| TreeAccounts {
688                        tree_type: TreeType::State,
689                        merkle_tree: state_merkle_tree_bundle.accounts.merkle_tree,
690                        queue: state_merkle_tree_bundle.accounts.nullifier_queue,
691                        is_rolledover: false,
692                    })
693                    .collect::<Vec<TreeAccounts>>();
694                self.indexer.get_address_merkle_trees().iter().for_each(
695                    |address_merkle_tree_bundle| {
696                        tree_accounts.push(TreeAccounts {
697                            tree_type: TreeType::Address,
698                            merkle_tree: address_merkle_tree_bundle.accounts.merkle_tree,
699                            queue: address_merkle_tree_bundle.accounts.queue,
700                            is_rolledover: false,
701                        });
702                    },
703                );
704                // 3.2 finalize registration for every forester
705                for forester in self.foresters.iter_mut() {
706                    if forester.is_registered.is_none() {
707                        continue;
708                    }
709                    println!(
710                        "registered forester {} for epoch {}",
711                        forester.keypair.pubkey(),
712                        self.epoch
713                    );
714                    println!(
715                        "forester.forester registration epoch {:?}",
716                        forester.forester.registration.epoch
717                    );
718                    println!(
719                        "forester.forester active epoch {:?}",
720                        forester.forester.active.epoch
721                    );
722                    println!(
723                        "forester.forester report_work epoch {:?}",
724                        forester.forester.report_work.epoch
725                    );
726
727                    forester
728                        .forester
729                        .active
730                        .fetch_account_and_add_trees_with_schedule(&mut self.rpc, &tree_accounts)
731                        .await
732                        .unwrap();
733                    let ix = create_finalize_registration_instruction(
734                        &forester.keypair.pubkey(),
735                        forester.forester.active.epoch,
736                    );
737                    self.rpc
738                        .create_and_send_transaction(
739                            &[ix],
740                            &forester.keypair.pubkey(),
741                            &[&forester.keypair],
742                        )
743                        .await
744                        .unwrap();
745                    assert_finalized_epoch_registration(
746                        &mut self.rpc,
747                        &forester.forester.active.forester_epoch_pda,
748                        &forester.forester.active.epoch_pda,
749                    )
750                    .await;
751                    self.stats.finalized_registrations += 1;
752                }
753            } else if current_active_epoch < self.epoch {
754                panic!(
755                    "current_active_epoch {} is less than self.epoch {}",
756                    current_active_epoch, self.epoch
757                );
758            }
759        }
760    }
761
762    pub async fn create_state_tree(&mut self, rollover_threshold: Option<u64>) {
763        let merkle_tree_keypair = Keypair::new(); //from_seed(&[self.rng.gen_range(0..255); 32]).unwrap();
764        let nullifier_queue_keypair = Keypair::new(); //from_seed(&[self.rng.gen_range(0..255); 32]).unwrap();
765        let cpi_context_keypair = Keypair::new();
766        let rollover_threshold = if let Some(rollover_threshold) = rollover_threshold {
767            Some(rollover_threshold)
768        } else if self.rng.gen_bool(0.5) && !self.keypair_action_config.fee_assert {
769            Some(self.rng.gen_range(1..100))
770        } else {
771            None
772        };
773        let merkle_tree_config = if !self.keypair_action_config.fee_assert {
774            StateMerkleTreeConfig {
775                height: 26,
776                changelog_size: self.rng.gen_range(1..5000),
777                roots_size: self.rng.gen_range(1..10000),
778                canopy_depth: 10,
779                network_fee: Some(5000),
780                close_threshold: None,
781                rollover_threshold,
782            }
783        } else {
784            StateMerkleTreeConfig::default()
785        };
786        println!("merkle tree config: {:?}", merkle_tree_config);
787        let queue_config = if !self.keypair_action_config.fee_assert {
788            let capacity: u32 = gen_prime(&mut self.rng, 1..10000).unwrap();
789            NullifierQueueConfig {
790                capacity: capacity as u16,
791                sequence_threshold: merkle_tree_config.roots_size + SAFETY_MARGIN,
792                network_fee: None,
793            }
794        } else if rollover_threshold.is_some() {
795            panic!("rollover_threshold should not be set when fee_assert is set (keypair_action_config.fee_assert)");
796        } else {
797            NullifierQueueConfig::default()
798        };
799        let forester = Pubkey::new_unique();
800        println!("queue config: {:?}", queue_config);
801        create_state_merkle_tree_and_queue_account(
802            &self.payer,
803            true,
804            &mut self.rpc,
805            &merkle_tree_keypair,
806            &nullifier_queue_keypair,
807            Some(&cpi_context_keypair),
808            None,
809            Some(forester),
810            1,
811            &merkle_tree_config,
812            &queue_config,
813        )
814        .await
815        .unwrap();
816        let merkle_tree = Box::new(light_merkle_tree_reference::MerkleTree::<Poseidon>::new(
817            STATE_MERKLE_TREE_HEIGHT as usize,
818            STATE_MERKLE_TREE_CANOPY_DEPTH as usize,
819        ));
820        let state_tree_account =
821            AccountZeroCopy::<account_compression::StateMerkleTreeAccount>::new(
822                &mut self.rpc,
823                nullifier_queue_keypair.pubkey(),
824            )
825            .await;
826        self.indexer
827            .get_state_merkle_trees_mut()
828            .push(StateMerkleTreeBundle {
829                rollover_fee: state_tree_account
830                    .deserialized()
831                    .metadata
832                    .rollover_metadata
833                    .rollover_fee as i64,
834                accounts: StateMerkleTreeAccounts {
835                    merkle_tree: merkle_tree_keypair.pubkey(),
836                    nullifier_queue: nullifier_queue_keypair.pubkey(),
837                    cpi_context: cpi_context_keypair.pubkey(),
838                },
839                merkle_tree,
840            });
841        // TODO: Add assert
842    }
843
844    pub async fn create_address_tree(&mut self, rollover_threshold: Option<u64>) {
845        let merkle_tree_keypair = Keypair::new();
846        let nullifier_queue_keypair = Keypair::new();
847        let rollover_threshold = if let Some(rollover_threshold) = rollover_threshold {
848            Some(rollover_threshold)
849        } else if self.rng.gen_bool(0.5) && !self.keypair_action_config.fee_assert {
850            Some(self.rng.gen_range(1..100))
851        } else {
852            None
853        };
854
855        let (config, address_config) = if !self.keypair_action_config.fee_assert {
856            let root_history = self.rng.gen_range(1..10000);
857            (
858                AddressMerkleTreeConfig {
859                    height: 26,
860                    changelog_size: self.rng.gen_range(1..5000),
861                    roots_size: root_history,
862                    canopy_depth: 10,
863                    address_changelog_size: self.rng.gen_range(1..5000),
864                    rollover_threshold,
865                    network_fee: Some(5000),
866                    close_threshold: None,
867                    // TODO: double check that close threshold cannot be set
868                },
869                AddressQueueConfig {
870                    sequence_threshold: root_history + SAFETY_MARGIN,
871                    ..Default::default()
872                },
873            )
874        } else if rollover_threshold.is_some() {
875            panic!("rollover_threshold should not be set when fee_assert is set (keypair_action_config.fee_assert)");
876        } else {
877            (
878                AddressMerkleTreeConfig::default(),
879                AddressQueueConfig::default(),
880            )
881        };
882
883        create_address_merkle_tree_and_queue_account(
884            &self.payer,
885            true,
886            &mut self.rpc,
887            &merkle_tree_keypair,
888            &nullifier_queue_keypair,
889            None,
890            None,
891            &config,
892            &address_config,
893            0,
894        )
895        .await
896        .unwrap();
897        let init_value = BigUint::from_str_radix(HIGHEST_ADDRESS_PLUS_ONE, 10).unwrap();
898        let mut merkle_tree = Box::new(
899            IndexedMerkleTree::<Poseidon, usize>::new(
900                STATE_MERKLE_TREE_HEIGHT as usize,
901                STATE_MERKLE_TREE_CANOPY_DEPTH as usize,
902            )
903            .unwrap(),
904        );
905        let mut indexed_array = Box::<IndexedArray<Poseidon, usize>>::default();
906        merkle_tree.append(&init_value, &mut indexed_array).unwrap();
907
908        let queue_account = AccountZeroCopy::<account_compression::QueueAccount>::new(
909            &mut self.rpc,
910            nullifier_queue_keypair.pubkey(),
911        )
912        .await;
913        self.indexer
914            .get_address_merkle_trees_mut()
915            .push(AddressMerkleTreeBundle {
916                rollover_fee: queue_account
917                    .deserialized()
918                    .metadata
919                    .rollover_metadata
920                    .rollover_fee as i64,
921                accounts: AddressMerkleTreeAccounts {
922                    merkle_tree: merkle_tree_keypair.pubkey(),
923                    queue: nullifier_queue_keypair.pubkey(),
924                },
925                merkle_tree,
926                indexed_array,
927            });
928        // TODO: Add assert
929    }
930
931    pub fn safe_gen_range<T, RR>(rng: &mut StdRng, range: RR, empty_fallback: T) -> T
932    where
933        T: SampleUniform + Copy,
934        RR: SampleRange<T> + Sized,
935    {
936        if range.is_empty() {
937            return empty_fallback;
938        }
939        rng.gen_range(range)
940    }
941
942    /// 1. Transfer spl tokens between random users
943    pub async fn activate_keypair_actions(&mut self, user: &Pubkey) {
944        let user_index = self
945            .users
946            .iter()
947            .position(|u| &u.keypair.pubkey() == user)
948            .unwrap();
949        // compress spl
950        // check sufficient spl balance
951        if self
952            .rng
953            .gen_bool(self.keypair_action_config.compress_spl.unwrap_or(0.0))
954            && self.users[user_index].token_accounts.is_empty()
955        // TODO: enable compress spl test
956        {
957            self.compress_spl(user_index).await;
958        }
959        // decompress spl
960        // check sufficient compressed spl balance
961        if self
962            .rng
963            .gen_bool(self.keypair_action_config.decompress_spl.unwrap_or(0.0))
964        {
965            self.decompress_spl(user_index).await;
966        }
967
968        // transfer spl
969        // check sufficient compressed spl balance
970        if self
971            .rng
972            .gen_bool(self.keypair_action_config.transfer_spl.unwrap_or(0.0))
973        {
974            self.transfer_spl(user_index).await;
975        }
976        // create address
977        if self
978            .rng
979            .gen_bool(self.keypair_action_config.create_address.unwrap_or(0.0))
980        {
981            self.create_address(None, None).await;
982        }
983
984        // compress sol
985        // check sufficient sol balance
986        let balance = self
987            .rpc
988            .get_balance(&self.users[user_index].keypair.pubkey())
989            .await
990            .unwrap();
991        if self
992            .rng
993            .gen_bool(self.keypair_action_config.compress_sol.unwrap_or(0.0))
994            && balance > 1000
995        {
996            self.compress_sol(user_index, balance).await;
997        } else {
998            println!("Not enough balance to compress sol. Balance: {}", balance);
999        }
1000
1001        // decompress sol
1002        // check sufficient compressed sol balance
1003        if self
1004            .rng
1005            .gen_bool(self.keypair_action_config.decompress_sol.unwrap_or(0.0))
1006        {
1007            self.decompress_sol(user_index).await;
1008        }
1009
1010        // transfer sol
1011        if self
1012            .rng
1013            .gen_bool(self.keypair_action_config.transfer_sol.unwrap_or(0.0))
1014        {
1015            self.transfer_sol(user_index).await;
1016        }
1017        // approve spl
1018        if self
1019            .rng
1020            .gen_bool(self.keypair_action_config.approve_spl.unwrap_or(0.0))
1021            && !self.users[user_index].token_accounts.is_empty()
1022        {
1023            self.approve_spl(user_index).await;
1024        }
1025        // revoke spl
1026        if self
1027            .rng
1028            .gen_bool(self.keypair_action_config.revoke_spl.unwrap_or(0.0))
1029            && !self.users[user_index].token_accounts.is_empty()
1030        {
1031            self.revoke_spl(user_index).await;
1032        }
1033        // burn spl
1034        if self
1035            .rng
1036            .gen_bool(self.keypair_action_config.burn_spl.unwrap_or(0.0))
1037            && !self.users[user_index].token_accounts.is_empty()
1038        {
1039            self.burn_spl(user_index).await;
1040        }
1041        // freeze spl
1042        if self
1043            .rng
1044            .gen_bool(self.keypair_action_config.freeze_spl.unwrap_or(0.0))
1045            && !self.users[user_index].token_accounts.is_empty()
1046        {
1047            self.freeze_spl(user_index).await;
1048        }
1049        // thaw spl
1050        if self
1051            .rng
1052            .gen_bool(self.keypair_action_config.thaw_spl.unwrap_or(0.0))
1053            && !self.users[user_index].token_accounts.is_empty()
1054        {
1055            self.thaw_spl(user_index).await;
1056        }
1057    }
1058
1059    pub fn get_eligible_forester_for_queue(
1060        queue_pubkey: &Pubkey,
1061        foresters: &[TestForester],
1062        light_slot: u64,
1063    ) -> Option<Keypair> {
1064        for f in foresters.iter() {
1065            let tree = f
1066                .forester
1067                .active
1068                .merkle_trees
1069                .iter()
1070                .find(|mt| mt.tree_accounts.queue == *queue_pubkey);
1071            if let Some(tree) = tree {
1072                if tree.is_eligible(light_slot) {
1073                    return Some(f.keypair.insecure_clone());
1074                }
1075            }
1076        }
1077        None
1078    }
1079    pub async fn transfer_sol_deterministic(
1080        &mut self,
1081        from: &Keypair,
1082        to: &Pubkey,
1083        tree_index: Option<usize>,
1084    ) -> Result<Signature, RpcError> {
1085        let input_compressed_accounts = self.get_compressed_sol_accounts(&from.pubkey());
1086        let output_merkle_tree = self.indexer.get_state_merkle_trees()[tree_index.unwrap_or(0)]
1087            .accounts
1088            .merkle_tree;
1089        let recipients = vec![*to];
1090        let transaction_params = if self.keypair_action_config.fee_assert {
1091            Some(TransactionParams {
1092                num_new_addresses: 0,
1093                num_input_compressed_accounts: input_compressed_accounts.len() as u8,
1094                num_output_compressed_accounts: 1u8,
1095                compress: 0,
1096                fee_config: FeeConfig::default(),
1097            })
1098        } else {
1099            None
1100        };
1101        transfer_compressed_sol_test(
1102            &mut self.rpc,
1103            &mut self.indexer,
1104            from,
1105            input_compressed_accounts.as_slice(),
1106            recipients.as_slice(),
1107            &[output_merkle_tree],
1108            transaction_params,
1109        )
1110        .await
1111    }
1112
1113    pub async fn transfer_sol(&mut self, user_index: usize) {
1114        let input_compressed_accounts = self.get_random_compressed_sol_accounts(user_index);
1115
1116        if !input_compressed_accounts.is_empty() {
1117            println!("\n --------------------------------------------------\n\t\t Transfer Sol\n --------------------------------------------------");
1118            let recipients = self
1119                .users
1120                .iter()
1121                .map(|u| u.keypair.pubkey())
1122                .collect::<Vec<Pubkey>>();
1123            let num_output_merkle_trees = Self::safe_gen_range(
1124                &mut self.rng,
1125                1..std::cmp::min(
1126                    self.keypair_action_config
1127                        .max_output_accounts
1128                        .unwrap_or(recipients.len() as u64),
1129                    recipients.len() as u64,
1130                ),
1131                1,
1132            );
1133            let recipients = recipients
1134                .choose_multiple(&mut self.rng, num_output_merkle_trees as usize)
1135                .copied()
1136                .collect::<Vec<_>>();
1137            let output_merkle_trees = self.get_merkle_tree_pubkeys(num_output_merkle_trees);
1138            let transaction_parameters = if self.keypair_action_config.fee_assert {
1139                Some(TransactionParams {
1140                    num_new_addresses: 0,
1141                    num_input_compressed_accounts: input_compressed_accounts.len() as u8,
1142                    num_output_compressed_accounts: num_output_merkle_trees as u8,
1143                    compress: 0,
1144                    fee_config: FeeConfig::default(),
1145                })
1146            } else {
1147                None
1148            };
1149            transfer_compressed_sol_test(
1150                &mut self.rpc,
1151                &mut self.indexer,
1152                &self.users[user_index].keypair,
1153                input_compressed_accounts.as_slice(),
1154                recipients.as_slice(),
1155                output_merkle_trees.as_slice(),
1156                transaction_parameters,
1157            )
1158            .await
1159            .unwrap();
1160            self.stats.sol_transfers += 1;
1161        }
1162    }
1163
1164    pub async fn decompress_sol(&mut self, user_index: usize) {
1165        let input_compressed_accounts = self.get_random_compressed_sol_accounts(user_index);
1166
1167        if !input_compressed_accounts.is_empty() {
1168            println!("\n --------------------------------------------------\n\t\t Decompress Sol\n --------------------------------------------------");
1169            let output_merkle_tree = self.get_merkle_tree_pubkeys(1)[0];
1170            let recipient = self.users
1171                [Self::safe_gen_range(&mut self.rng, 0..std::cmp::min(self.users.len(), 6), 0)]
1172            .keypair
1173            .pubkey();
1174            let balance = input_compressed_accounts
1175                .iter()
1176                .map(|x| x.compressed_account.lamports)
1177                .sum::<u64>();
1178            let decompress_amount = Self::safe_gen_range(&mut self.rng, 1000..balance, balance / 2);
1179            let transaction_paramets = if self.keypair_action_config.fee_assert {
1180                Some(TransactionParams {
1181                    num_new_addresses: 0,
1182                    num_input_compressed_accounts: input_compressed_accounts.len() as u8,
1183                    num_output_compressed_accounts: 1u8,
1184                    compress: 0,
1185                    fee_config: FeeConfig::default(),
1186                })
1187            } else {
1188                None
1189            };
1190            decompress_sol_test(
1191                &mut self.rpc,
1192                &mut self.indexer,
1193                &self.users[user_index].keypair,
1194                &input_compressed_accounts,
1195                &recipient,
1196                decompress_amount,
1197                &output_merkle_tree,
1198                transaction_paramets,
1199            )
1200            .await
1201            .unwrap();
1202            self.stats.sol_decompress += 1;
1203        }
1204    }
1205
1206    pub async fn compress_sol_deterministic(
1207        &mut self,
1208        from: &Keypair,
1209        amount: u64,
1210        tree_index: Option<usize>,
1211    ) {
1212        let input_compressed_accounts = self.get_compressed_sol_accounts(&from.pubkey());
1213        let output_merkle_tree = self.indexer.get_state_merkle_trees()[tree_index.unwrap_or(0)]
1214            .accounts
1215            .merkle_tree;
1216        let transaction_parameters = if self.keypair_action_config.fee_assert {
1217            Some(TransactionParams {
1218                num_new_addresses: 0,
1219                num_input_compressed_accounts: input_compressed_accounts.len() as u8,
1220                num_output_compressed_accounts: 1u8,
1221                compress: amount as i64,
1222                fee_config: FeeConfig::default(),
1223            })
1224        } else {
1225            None
1226        };
1227        compress_sol_test(
1228            &mut self.rpc,
1229            &mut self.indexer,
1230            from,
1231            input_compressed_accounts.as_slice(),
1232            false,
1233            amount,
1234            &output_merkle_tree,
1235            transaction_parameters,
1236        )
1237        .await
1238        .unwrap();
1239    }
1240
1241    pub async fn compress_sol(&mut self, user_index: usize, balance: u64) {
1242        println!("\n --------------------------------------------------\n\t\t Compress Sol\n --------------------------------------------------");
1243        // Limit max compress amount to 1 sol so that context.payer doesn't get depleted by airdrops.
1244        let max_amount = std::cmp::min(balance, 1_000_000_000);
1245        let amount = Self::safe_gen_range(&mut self.rng, 1000..max_amount, max_amount / 2);
1246        let input_compressed_accounts = self.get_random_compressed_sol_accounts(user_index);
1247        let create_output_compressed_accounts_for_input_accounts = false;
1248        // TODO: debug Merkle trees in wrong order
1249        // if input_compressed_accounts.is_empty() {
1250        //     false
1251        // } else {
1252        //     self.rng.gen_bool(0.5)
1253        // };
1254        let output_merkle_tree = self.get_merkle_tree_pubkeys(1)[0];
1255        let transaction_parameters = if self.keypair_action_config.fee_assert {
1256            Some(TransactionParams {
1257                num_new_addresses: 0,
1258                num_input_compressed_accounts: input_compressed_accounts.len() as u8,
1259                num_output_compressed_accounts: 1u8,
1260                compress: amount as i64,
1261                fee_config: FeeConfig::default(),
1262            })
1263        } else {
1264            None
1265        };
1266        compress_sol_test(
1267            &mut self.rpc,
1268            &mut self.indexer,
1269            &self.users[user_index].keypair,
1270            input_compressed_accounts.as_slice(),
1271            create_output_compressed_accounts_for_input_accounts,
1272            amount,
1273            &output_merkle_tree,
1274            transaction_parameters,
1275        )
1276        .await
1277        .unwrap();
1278        airdrop_lamports(
1279            &mut self.rpc,
1280            &self.users[user_index].keypair.pubkey(),
1281            amount,
1282        )
1283        .await
1284        .unwrap();
1285        self.stats.sol_compress += 1;
1286    }
1287
1288    pub async fn create_address(
1289        &mut self,
1290        optional_addresses: Option<Vec<Pubkey>>,
1291        address_tree_index: Option<usize>,
1292    ) -> Vec<Pubkey> {
1293        println!("\n --------------------------------------------------\n\t\t Create Address\n --------------------------------------------------");
1294        // select number of addresses to create
1295        let num_addresses = self.rng.gen_range(1..=2);
1296        let (address_merkle_tree_pubkeys, address_queue_pubkeys) =
1297            if let Some(address_tree_index) = address_tree_index {
1298                (
1299                    vec![
1300                        self.indexer.get_address_merkle_trees()[address_tree_index]
1301                            .accounts
1302                            .merkle_tree;
1303                        num_addresses as usize
1304                    ],
1305                    vec![
1306                        self.indexer.get_address_merkle_trees()[address_tree_index]
1307                            .accounts
1308                            .queue;
1309                        num_addresses as usize
1310                    ],
1311                )
1312            } else {
1313                // select random address Merkle tree(s)
1314                self.get_address_merkle_tree_pubkeys(num_addresses)
1315            };
1316        let mut address_seeds = Vec::new();
1317        let mut created_addresses = Vec::new();
1318
1319        if let Some(addresses) = optional_addresses {
1320            for address in addresses {
1321                let address_seed: [u8; 32] = address.to_bytes();
1322                address_seeds.push(address_seed);
1323                created_addresses.push(address);
1324            }
1325        } else {
1326            for _ in 0..num_addresses {
1327                let address_seed: [u8; 32] =
1328                    bigint_to_be_bytes_array::<32>(&self.rng.gen_biguint(256)).unwrap();
1329                address_seeds.push(address_seed);
1330                created_addresses.push(Pubkey::from(address_seed));
1331            }
1332        }
1333
1334        let output_compressed_accounts = self.get_merkle_tree_pubkeys(num_addresses);
1335        let transaction_parameters = if self.keypair_action_config.fee_assert {
1336            Some(TransactionParams {
1337                num_new_addresses: num_addresses as u8,
1338                num_input_compressed_accounts: 0u8,
1339                num_output_compressed_accounts: num_addresses as u8,
1340                compress: 0,
1341                fee_config: FeeConfig::default(),
1342            })
1343        } else {
1344            None
1345        };
1346        // TODO: add other input compressed accounts
1347        // (to test whether the address generation degrades performance)
1348        create_addresses_test(
1349            &mut self.rpc,
1350            &mut self.indexer,
1351            address_merkle_tree_pubkeys.as_slice(),
1352            address_queue_pubkeys.as_slice(),
1353            output_compressed_accounts,
1354            address_seeds.as_slice(),
1355            &Vec::new(),
1356            false,
1357            transaction_parameters,
1358        )
1359        .await
1360        .unwrap();
1361        self.stats.create_address += num_addresses;
1362        created_addresses
1363    }
1364
1365    pub async fn transfer_spl(&mut self, user_index: usize) {
1366        let user = &self.users[user_index].keypair.pubkey();
1367        println!("\n --------------------------------------------------\n\t\t Tranfer Spl\n --------------------------------------------------");
1368        let (mint, mut token_accounts) = self.select_random_compressed_token_accounts(user).await;
1369        if token_accounts.is_empty() {
1370            let mt_pubkeys = self.get_merkle_tree_pubkeys(1);
1371            mint_tokens_helper(
1372                &mut self.rpc,
1373                &mut self.indexer,
1374                &mt_pubkeys[0],
1375                &self.payer,
1376                &mint,
1377                vec![Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000); 1],
1378                vec![*user; 1],
1379            )
1380            .await;
1381            let (_, _token_accounts) = self.select_random_compressed_token_accounts(user).await;
1382            token_accounts = _token_accounts;
1383        }
1384        let recipients = token_accounts
1385            .iter()
1386            .map(|_| {
1387                self.users
1388                    [Self::safe_gen_range(&mut self.rng, 0..std::cmp::min(self.users.len(), 6), 0)]
1389                .keypair
1390                .pubkey()
1391            })
1392            .collect::<Vec<_>>();
1393        println!("Recipients: {:?}", recipients.len());
1394        let max_amount = token_accounts
1395            .iter()
1396            .map(|token_account| token_account.token_data.amount)
1397            .sum::<u64>();
1398        let amount = Self::safe_gen_range(&mut self.rng, 1000..max_amount, max_amount / 2);
1399        let equal_amount = amount / recipients.len() as u64;
1400        let num_output_compressed_accounts = if max_amount - amount != 0 {
1401            recipients.len() + 1
1402        } else {
1403            recipients.len()
1404        };
1405        // get different amounts for each recipient so that every compressed account is unique
1406        let amounts = recipients
1407            .iter()
1408            .enumerate()
1409            .map(|(i, _)| equal_amount - i as u64)
1410            .collect::<Vec<u64>>();
1411
1412        let output_merkle_tree_pubkeys =
1413            self.get_merkle_tree_pubkeys(num_output_compressed_accounts as u64);
1414        let transaction_paramets = if self.keypair_action_config.fee_assert {
1415            Some(TransactionParams {
1416                num_new_addresses: 0u8,
1417                num_input_compressed_accounts: token_accounts.len() as u8,
1418                num_output_compressed_accounts: output_merkle_tree_pubkeys.len() as u8,
1419                compress: 0,
1420                fee_config: FeeConfig::default(),
1421            })
1422        } else {
1423            None
1424        };
1425        compressed_transfer_test(
1426            &self.rpc.get_payer().insecure_clone(),
1427            &mut self.rpc,
1428            &mut self.indexer,
1429            &mint,
1430            &self.users[user_index].keypair.insecure_clone(),
1431            &recipients,
1432            &amounts,
1433            None,
1434            &token_accounts,
1435            &output_merkle_tree_pubkeys,
1436            None,
1437            false,
1438            transaction_paramets,
1439        )
1440        .await;
1441        self.stats.spl_transfers += 1;
1442    }
1443
1444    pub async fn approve_spl(&mut self, user_index: usize) {
1445        let user = &self.users[user_index].keypair.pubkey();
1446        println!("\n --------------------------------------------------\n\t\t Approve Spl\n --------------------------------------------------");
1447        let (mint, mut token_accounts) = self.select_random_compressed_token_accounts(user).await;
1448        if token_accounts.is_empty() {
1449            let mt_pubkeys = self.get_merkle_tree_pubkeys(1);
1450            mint_tokens_helper(
1451                &mut self.rpc,
1452                &mut self.indexer,
1453                &mt_pubkeys[0],
1454                &self.payer,
1455                &mint,
1456                vec![Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000); 1],
1457                vec![*user; 1],
1458            )
1459            .await;
1460            let (_, _token_accounts) = self.select_random_compressed_token_accounts(user).await;
1461            token_accounts = _token_accounts;
1462        }
1463        println!("token_accounts: {:?}", token_accounts);
1464        let rnd_user_index = self.rng.gen_range(0..self.users.len());
1465        let delegate = self.users[rnd_user_index].keypair.pubkey();
1466        let max_amount = token_accounts
1467            .iter()
1468            .map(|token_account| token_account.token_data.amount)
1469            .sum::<u64>();
1470        let delegate_amount = Self::safe_gen_range(&mut self.rng, 0..max_amount, max_amount / 2);
1471        let num_output_compressed_accounts = if delegate_amount != max_amount { 2 } else { 1 };
1472        let output_merkle_tree_pubkeys = self.get_merkle_tree_pubkeys(2);
1473        let transaction_paramets = if self.keypair_action_config.fee_assert {
1474            Some(TransactionParams {
1475                num_new_addresses: 0u8,
1476                num_input_compressed_accounts: token_accounts.len() as u8,
1477                num_output_compressed_accounts,
1478                compress: 0,
1479                fee_config: FeeConfig::default(),
1480            })
1481        } else {
1482            None
1483        };
1484        approve_test(
1485            &self.users[user_index].keypair,
1486            &mut self.rpc,
1487            &mut self.indexer,
1488            token_accounts,
1489            delegate_amount,
1490            None,
1491            &delegate,
1492            &output_merkle_tree_pubkeys[0],
1493            &output_merkle_tree_pubkeys[1],
1494            transaction_paramets,
1495        )
1496        .await;
1497        self.stats.spl_approved += 1;
1498    }
1499
1500    pub async fn revoke_spl(&mut self, user_index: usize) {
1501        let user = &self.users[user_index].keypair.pubkey();
1502        println!("\n --------------------------------------------------\n\t\t Revoke Spl\n --------------------------------------------------");
1503        let (mint, mut token_accounts) = self
1504            .select_random_compressed_token_accounts_delegated(user, true, None, false)
1505            .await;
1506        if token_accounts.is_empty() {
1507            let mt_pubkeys = self.get_merkle_tree_pubkeys(1);
1508            mint_tokens_helper(
1509                &mut self.rpc,
1510                &mut self.indexer,
1511                &mt_pubkeys[0],
1512                &self.payer,
1513                &mint,
1514                vec![Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000); 1],
1515                vec![*user; 1],
1516            )
1517            .await;
1518            self.approve_spl(user_index).await;
1519            let (_, _token_accounts) = self
1520                .select_random_compressed_token_accounts_delegated(user, true, None, false)
1521                .await;
1522            token_accounts = _token_accounts;
1523        }
1524        let num_output_compressed_accounts = 1;
1525        let output_merkle_tree_pubkeys = self.get_merkle_tree_pubkeys(1);
1526        let transaction_paramets = if self.keypair_action_config.fee_assert {
1527            Some(TransactionParams {
1528                num_new_addresses: 0u8,
1529                num_input_compressed_accounts: token_accounts.len() as u8,
1530                num_output_compressed_accounts,
1531                compress: 0,
1532                fee_config: FeeConfig::default(),
1533            })
1534        } else {
1535            None
1536        };
1537        revoke_test(
1538            &self.users[user_index].keypair,
1539            &mut self.rpc,
1540            &mut self.indexer,
1541            token_accounts,
1542            &output_merkle_tree_pubkeys[0],
1543            transaction_paramets,
1544        )
1545        .await;
1546        self.stats.spl_revoked += 1;
1547    }
1548
1549    pub async fn burn_spl(&mut self, user_index: usize) {
1550        let user = &self.users[user_index].keypair.pubkey();
1551        println!("\n --------------------------------------------------\n\t\t Burn Spl\n --------------------------------------------------");
1552        let (mint, mut token_accounts) = self.select_random_compressed_token_accounts(user).await;
1553        if token_accounts.is_empty() {
1554            let mt_pubkeys = self.get_merkle_tree_pubkeys(1);
1555            mint_tokens_helper(
1556                &mut self.rpc,
1557                &mut self.indexer,
1558                &mt_pubkeys[0],
1559                &self.payer,
1560                &mint,
1561                vec![Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000); 1],
1562                vec![*user; 1],
1563            )
1564            .await;
1565            let (_, _token_accounts) = self.select_random_compressed_token_accounts(user).await;
1566            token_accounts = _token_accounts;
1567        }
1568        let max_amount = token_accounts
1569            .iter()
1570            .map(|token_account| token_account.token_data.amount)
1571            .sum::<u64>();
1572        let burn_amount = Self::safe_gen_range(&mut self.rng, 0..max_amount, max_amount / 2);
1573        let num_output_compressed_accounts = if burn_amount != max_amount { 1 } else { 0 };
1574        let output_merkle_tree_pubkeys = self.get_merkle_tree_pubkeys(1);
1575        let transaction_paramets = if self.keypair_action_config.fee_assert {
1576            Some(TransactionParams {
1577                num_new_addresses: 0u8,
1578                num_input_compressed_accounts: token_accounts.len() as u8,
1579                num_output_compressed_accounts,
1580                compress: 0,
1581                fee_config: FeeConfig::default(),
1582            })
1583        } else {
1584            None
1585        };
1586
1587        burn_test(
1588            &self.users[user_index].keypair,
1589            &mut self.rpc,
1590            &mut self.indexer,
1591            token_accounts,
1592            &output_merkle_tree_pubkeys[0],
1593            burn_amount,
1594            false,
1595            transaction_paramets,
1596        )
1597        .await;
1598        self.stats.spl_burned += 1;
1599    }
1600
1601    pub async fn freeze_spl(&mut self, user_index: usize) {
1602        let user = &self.users[user_index].keypair.pubkey();
1603        println!("\n --------------------------------------------------\n\t\t Freeze Spl\n --------------------------------------------------");
1604        let (mint, mut token_accounts) = self.select_random_compressed_token_accounts(user).await;
1605        if token_accounts.is_empty() {
1606            let mt_pubkeys = self.get_merkle_tree_pubkeys(1);
1607            mint_tokens_helper(
1608                &mut self.rpc,
1609                &mut self.indexer,
1610                &mt_pubkeys[0],
1611                &self.payer,
1612                &mint,
1613                vec![Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000); 1],
1614                vec![*user; 1],
1615            )
1616            .await;
1617            let (_, _token_accounts) = self
1618                .select_random_compressed_token_accounts_delegated(user, false, None, false)
1619                .await;
1620            token_accounts = _token_accounts;
1621        }
1622        let output_merkle_tree_pubkeys = self.get_merkle_tree_pubkeys(1);
1623        let transaction_paramets = if self.keypair_action_config.fee_assert {
1624            Some(TransactionParams {
1625                num_new_addresses: 0u8,
1626                num_input_compressed_accounts: token_accounts.len() as u8,
1627                num_output_compressed_accounts: token_accounts.len() as u8,
1628                compress: 0,
1629                fee_config: FeeConfig::default(),
1630            })
1631        } else {
1632            None
1633        };
1634        freeze_test(
1635            &self.rpc.get_payer().insecure_clone(),
1636            &mut self.rpc,
1637            &mut self.indexer,
1638            token_accounts,
1639            &output_merkle_tree_pubkeys[0],
1640            transaction_paramets,
1641        )
1642        .await;
1643        self.stats.spl_frozen += 1;
1644    }
1645
1646    pub async fn thaw_spl(&mut self, user_index: usize) {
1647        let user = &self.users[user_index].keypair.pubkey();
1648        println!("\n --------------------------------------------------\n\t\t Thaw Spl\n --------------------------------------------------");
1649        let (_, mut token_accounts) = self
1650            .select_random_compressed_token_accounts_frozen(user)
1651            .await;
1652        if token_accounts.is_empty() {
1653            self.freeze_spl(user_index).await;
1654
1655            let (_, _token_accounts) = self
1656                .select_random_compressed_token_accounts_frozen(user)
1657                .await;
1658            token_accounts = _token_accounts;
1659        }
1660        let output_merkle_tree_pubkeys = self.get_merkle_tree_pubkeys(1);
1661        let transaction_paramets = if self.keypair_action_config.fee_assert {
1662            Some(TransactionParams {
1663                num_new_addresses: 0u8,
1664                num_input_compressed_accounts: token_accounts.len() as u8,
1665                num_output_compressed_accounts: token_accounts.len() as u8,
1666                compress: 0,
1667                fee_config: FeeConfig::default(),
1668            })
1669        } else {
1670            None
1671        };
1672
1673        thaw_test(
1674            &self.rpc.get_payer().insecure_clone(),
1675            &mut self.rpc,
1676            &mut self.indexer,
1677            token_accounts,
1678            &output_merkle_tree_pubkeys[0],
1679            transaction_paramets,
1680        )
1681        .await;
1682        self.stats.spl_thawed += 1;
1683    }
1684
1685    pub async fn compress_spl(&mut self, user_index: usize) {
1686        println!("\n --------------------------------------------------\n\t\t Compress Spl\n --------------------------------------------------");
1687        let mut balance = 0;
1688        let mut mint = Pubkey::default();
1689        let mut token_account = Pubkey::default();
1690        for _ in 0..self.users[user_index].token_accounts.len() {
1691            let (_mint, _token_account) = self.users[user_index].token_accounts[self
1692                .rng
1693                .gen_range(0..self.users[user_index].token_accounts.len())];
1694            token_account = _token_account;
1695            mint = _mint;
1696            self.rpc.get_account(_token_account).await.unwrap();
1697            use solana_sdk::program_pack::Pack;
1698            let account = spl_token::state::Account::unpack(
1699                &self
1700                    .rpc
1701                    .get_account(_token_account)
1702                    .await
1703                    .unwrap()
1704                    .unwrap()
1705                    .data,
1706            )
1707            .unwrap();
1708            balance = account.amount;
1709            if balance != 0 {
1710                break;
1711            }
1712        }
1713        if balance != 0 {
1714            self.users[user_index]
1715                .token_accounts
1716                .push((mint, token_account));
1717            let output_merkle_tree_account = self.get_merkle_tree_pubkeys(1);
1718
1719            let amount = Self::safe_gen_range(&mut self.rng, 1000..balance, balance / 2);
1720            let transaction_paramets = if self.keypair_action_config.fee_assert {
1721                Some(TransactionParams {
1722                    num_new_addresses: 0u8,
1723                    num_input_compressed_accounts: 0u8,
1724                    num_output_compressed_accounts: 1u8,
1725                    compress: 0, // sol amount this is a spl compress test
1726                    fee_config: FeeConfig::default(),
1727                })
1728            } else {
1729                None
1730            };
1731            compress_test(
1732                &self.users[user_index].keypair,
1733                &mut self.rpc,
1734                &mut self.indexer,
1735                amount,
1736                &mint,
1737                &output_merkle_tree_account[0],
1738                &token_account,
1739                transaction_paramets,
1740            )
1741            .await;
1742            self.stats.spl_compress += 1;
1743        }
1744    }
1745
1746    pub async fn decompress_spl(&mut self, user_index: usize) {
1747        let user = &self.users[user_index].keypair.pubkey();
1748        println!("\n --------------------------------------------------\n\t\t Decompress Spl\n --------------------------------------------------");
1749        let (mint, mut token_accounts) = self.select_random_compressed_token_accounts(user).await;
1750        if token_accounts.is_empty() {
1751            let mt_pubkeys = self.get_merkle_tree_pubkeys(1);
1752            mint_tokens_helper(
1753                &mut self.rpc,
1754                &mut self.indexer,
1755                &mt_pubkeys[0],
1756                &self.payer,
1757                &mint,
1758                vec![Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000); 1],
1759                vec![*user; 1],
1760            )
1761            .await;
1762            let (_, _token_accounts) = self.select_random_compressed_token_accounts(user).await;
1763            token_accounts = _token_accounts;
1764        }
1765        let token_account = match self.users[user_index]
1766            .token_accounts
1767            .iter()
1768            .find(|t| t.0 == mint)
1769        {
1770            Some(token_account) => token_account.1,
1771            None => {
1772                let token_account_keypair = Keypair::new();
1773                create_token_account(
1774                    &mut self.rpc,
1775                    &mint,
1776                    &token_account_keypair,
1777                    &self.users[user_index].keypair,
1778                )
1779                .await
1780                .unwrap();
1781
1782                token_account_keypair.pubkey()
1783            }
1784        };
1785
1786        self.users[user_index]
1787            .token_accounts
1788            .push((mint, token_account));
1789        let output_merkle_tree_account = self.get_merkle_tree_pubkeys(1);
1790        let max_amount = token_accounts
1791            .iter()
1792            .map(|token_account| token_account.token_data.amount)
1793            .sum::<u64>();
1794        let amount = Self::safe_gen_range(&mut self.rng, 1000..max_amount, max_amount / 2);
1795        let transaction_paramets = if self.keypair_action_config.fee_assert {
1796            Some(TransactionParams {
1797                num_new_addresses: 0u8,
1798                num_input_compressed_accounts: token_accounts.len() as u8,
1799                num_output_compressed_accounts: 1u8,
1800                compress: 0,
1801                fee_config: FeeConfig::default(),
1802            })
1803        } else {
1804            None
1805        };
1806        // decompress
1807        decompress_test(
1808            &self.users[user_index].keypair,
1809            &mut self.rpc,
1810            &mut self.indexer,
1811            token_accounts.clone(),
1812            amount,
1813            &output_merkle_tree_account[0],
1814            &token_account,
1815            transaction_paramets,
1816        )
1817        .await;
1818        self.stats.spl_decompress += 1;
1819    }
1820
1821    pub async fn rollover_state_merkle_tree_and_queue(
1822        &mut self,
1823        index: usize,
1824        payer: &Keypair,
1825        epoch: u64,
1826    ) -> Result<(), RpcError> {
1827        let bundle = self.indexer.get_state_merkle_trees()[index].accounts;
1828        let new_nullifier_queue_keypair = Keypair::new();
1829        let new_merkle_tree_keypair = Keypair::new();
1830        // TODO: move into registry program
1831        let new_cpi_signature_keypair = Keypair::new();
1832        let fee_payer_balance = self
1833            .rpc
1834            .get_balance(&self.indexer.get_payer().pubkey())
1835            .await
1836            .unwrap();
1837        let rollover_signature_and_slot = perform_state_merkle_tree_roll_over_forester(
1838            payer,
1839            &mut self.rpc,
1840            &new_nullifier_queue_keypair,
1841            &new_merkle_tree_keypair,
1842            &new_cpi_signature_keypair,
1843            &bundle.merkle_tree,
1844            &bundle.nullifier_queue,
1845            epoch,
1846            false,
1847        )
1848        .await
1849        .unwrap();
1850        info!("Rollover signature: {:?}", rollover_signature_and_slot.0);
1851        let additional_rent = self
1852            .rpc
1853            .get_minimum_balance_for_rent_exemption(
1854                ProtocolConfig::default().cpi_context_size as usize,
1855            )
1856            .await
1857            .unwrap();
1858        info!("additional_rent: {:?}", additional_rent);
1859        assert_rolled_over_pair(
1860            &self.indexer.get_payer().pubkey(),
1861            &mut self.rpc,
1862            &fee_payer_balance,
1863            &bundle.merkle_tree,
1864            &bundle.nullifier_queue,
1865            &new_merkle_tree_keypair.pubkey(),
1866            &new_nullifier_queue_keypair.pubkey(),
1867            rollover_signature_and_slot.1,
1868            additional_rent,
1869            4,
1870        )
1871        .await;
1872        self.indexer
1873            .get_state_merkle_trees_mut()
1874            .push(StateMerkleTreeBundle {
1875                // TODO: fetch correct fee when this property is used
1876                rollover_fee: 0,
1877                accounts: StateMerkleTreeAccounts {
1878                    merkle_tree: new_merkle_tree_keypair.pubkey(),
1879                    nullifier_queue: new_nullifier_queue_keypair.pubkey(),
1880                    cpi_context: new_cpi_signature_keypair.pubkey(),
1881                },
1882                merkle_tree: Box::new(light_merkle_tree_reference::MerkleTree::<Poseidon>::new(
1883                    STATE_MERKLE_TREE_HEIGHT as usize,
1884                    STATE_MERKLE_TREE_CANOPY_DEPTH as usize,
1885                )),
1886            });
1887        Ok(())
1888    }
1889
1890    pub async fn rollover_address_merkle_tree_and_queue(
1891        &mut self,
1892        index: usize,
1893        payer: &Keypair,
1894        epoch: u64,
1895    ) -> Result<(), RpcError> {
1896        let bundle = self.indexer.get_address_merkle_trees()[index].accounts;
1897        let new_nullifier_queue_keypair = Keypair::new();
1898        let new_merkle_tree_keypair = Keypair::new();
1899        let fee_payer_balance = self
1900            .rpc
1901            .get_balance(&self.indexer.get_payer().pubkey())
1902            .await
1903            .unwrap();
1904        println!("prior balance {}", fee_payer_balance);
1905        perform_address_merkle_tree_roll_over_forester(
1906            payer,
1907            &mut self.rpc,
1908            &new_nullifier_queue_keypair,
1909            &new_merkle_tree_keypair,
1910            &bundle.merkle_tree,
1911            &bundle.queue,
1912            epoch,
1913            false,
1914        )
1915        .await?;
1916        assert_rolled_over_address_merkle_tree_and_queue(
1917            &self.indexer.get_payer().pubkey(),
1918            &mut self.rpc,
1919            &fee_payer_balance,
1920            &bundle.merkle_tree,
1921            &bundle.queue,
1922            &new_merkle_tree_keypair.pubkey(),
1923            &new_nullifier_queue_keypair.pubkey(),
1924        )
1925        .await;
1926        self.indexer.add_address_merkle_tree_accounts(
1927            &new_merkle_tree_keypair,
1928            &new_nullifier_queue_keypair,
1929            None,
1930        );
1931        Ok(())
1932    }
1933
1934    pub fn get_random_compressed_sol_accounts(
1935        &mut self,
1936        user_index: usize,
1937    ) -> Vec<CompressedAccountWithMerkleContext> {
1938        let input_compressed_accounts = self
1939            .indexer
1940            .get_compressed_accounts_by_owner(&self.users[user_index].keypair.pubkey());
1941        let range = std::cmp::min(input_compressed_accounts.len(), 4);
1942        let number_of_compressed_accounts = Self::safe_gen_range(&mut self.rng, 0..=range, 0);
1943        input_compressed_accounts[0..number_of_compressed_accounts].to_vec()
1944    }
1945
1946    pub fn get_compressed_sol_accounts(
1947        &self,
1948        pubkey: &Pubkey,
1949    ) -> Vec<CompressedAccountWithMerkleContext> {
1950        self.indexer.get_compressed_accounts_by_owner(pubkey)
1951    }
1952
1953    pub fn get_merkle_tree_pubkeys(&mut self, num: u64) -> Vec<Pubkey> {
1954        let mut pubkeys = vec![];
1955        for _ in 0..num {
1956            let range_max: usize = std::cmp::min(
1957                self.keypair_action_config
1958                    .max_output_accounts
1959                    .unwrap_or(self.indexer.get_state_merkle_trees().len() as u64),
1960                self.indexer.get_state_merkle_trees().len() as u64,
1961            ) as usize;
1962
1963            let index = Self::safe_gen_range(&mut self.rng, 0..range_max, 0);
1964            pubkeys.push(
1965                self.indexer.get_state_merkle_trees()[index]
1966                    .accounts
1967                    .merkle_tree,
1968            );
1969        }
1970        pubkeys.sort();
1971        pubkeys
1972    }
1973
1974    pub fn get_address_merkle_tree_pubkeys(&mut self, num: u64) -> (Vec<Pubkey>, Vec<Pubkey>) {
1975        let mut pubkeys = vec![];
1976        let mut queue_pubkeys = vec![];
1977        for _ in 0..num {
1978            let index = Self::safe_gen_range(
1979                &mut self.rng,
1980                0..self.indexer.get_address_merkle_trees().len(),
1981                0,
1982            );
1983            pubkeys.push(
1984                self.indexer.get_address_merkle_trees()[index]
1985                    .accounts
1986                    .merkle_tree,
1987            );
1988            queue_pubkeys.push(
1989                self.indexer.get_address_merkle_trees()[index]
1990                    .accounts
1991                    .queue,
1992            );
1993        }
1994        (pubkeys, queue_pubkeys)
1995    }
1996
1997    pub async fn select_random_compressed_token_accounts(
1998        &mut self,
1999        user: &Pubkey,
2000    ) -> (Pubkey, Vec<TokenDataWithContext>) {
2001        self.select_random_compressed_token_accounts_delegated(user, false, None, false)
2002            .await
2003    }
2004
2005    pub async fn select_random_compressed_token_accounts_frozen(
2006        &mut self,
2007        user: &Pubkey,
2008    ) -> (Pubkey, Vec<TokenDataWithContext>) {
2009        self.select_random_compressed_token_accounts_delegated(user, false, None, true)
2010            .await
2011    }
2012
2013    pub async fn select_random_compressed_token_accounts_delegated(
2014        &mut self,
2015        user: &Pubkey,
2016        delegated: bool,
2017        delegate: Option<Pubkey>,
2018        frozen: bool,
2019    ) -> (Pubkey, Vec<TokenDataWithContext>) {
2020        let user_token_accounts = &mut self.indexer.get_compressed_token_accounts_by_owner(user);
2021        // clean up dust so that we don't run into issues that account balances are too low
2022        user_token_accounts.retain(|t| t.token_data.amount > 1000);
2023        let mut token_accounts_with_mint;
2024        let mint;
2025        if user_token_accounts.is_empty() {
2026            mint = self.indexer.get_token_compressed_accounts()[self
2027                .rng
2028                .gen_range(0..self.indexer.get_token_compressed_accounts().len())]
2029            .token_data
2030            .mint;
2031            let number_of_compressed_accounts = Self::safe_gen_range(&mut self.rng, 1..8, 1);
2032            let mt_pubkey = self.indexer.get_state_merkle_trees()[0]
2033                .accounts
2034                .merkle_tree;
2035            mint_tokens_helper(
2036                &mut self.rpc,
2037                &mut self.indexer,
2038                &mt_pubkey,
2039                &self.payer,
2040                &mint,
2041                vec![
2042                    Self::safe_gen_range(&mut self.rng, 100_000..1_000_000, 100_000);
2043                    number_of_compressed_accounts
2044                ],
2045                vec![*user; number_of_compressed_accounts],
2046            )
2047            .await;
2048            token_accounts_with_mint = self
2049                .indexer
2050                .get_compressed_token_accounts_by_owner(user)
2051                .iter()
2052                .filter(|token_account| token_account.token_data.mint == mint)
2053                .cloned()
2054                .collect::<Vec<_>>();
2055        } else {
2056            mint = user_token_accounts
2057                [Self::safe_gen_range(&mut self.rng, 0..user_token_accounts.len(), 0)]
2058            .token_data
2059            .mint;
2060            token_accounts_with_mint = user_token_accounts
2061                .iter()
2062                .filter(|token_account| token_account.token_data.mint == mint)
2063                .map(|token_account| (*token_account).clone())
2064                .collect::<Vec<TokenDataWithContext>>();
2065        }
2066        if delegated {
2067            token_accounts_with_mint = token_accounts_with_mint
2068                .iter()
2069                .filter(|token_account| token_account.token_data.delegate.is_some())
2070                .map(|token_account| (*token_account).clone())
2071                .collect::<Vec<TokenDataWithContext>>();
2072            if token_accounts_with_mint.is_empty() {
2073                return (mint, Vec::new());
2074            }
2075        }
2076        if let Some(delegate) = delegate {
2077            token_accounts_with_mint = token_accounts_with_mint
2078                .iter()
2079                .filter(|token_account| token_account.token_data.delegate.unwrap() == delegate)
2080                .map(|token_account| (*token_account).clone())
2081                .collect::<Vec<TokenDataWithContext>>();
2082        }
2083        if frozen {
2084            token_accounts_with_mint = token_accounts_with_mint
2085                .iter()
2086                .filter(|token_account| token_account.token_data.state == AccountState::Frozen)
2087                .map(|token_account| (*token_account).clone())
2088                .collect::<Vec<TokenDataWithContext>>();
2089            if token_accounts_with_mint.is_empty() {
2090                return (mint, Vec::new());
2091            }
2092        } else {
2093            token_accounts_with_mint = token_accounts_with_mint
2094                .iter()
2095                .filter(|token_account| token_account.token_data.state == AccountState::Initialized)
2096                .map(|token_account| (*token_account).clone())
2097                .collect::<Vec<TokenDataWithContext>>();
2098        }
2099        let range_end = if token_accounts_with_mint.len() == 1 {
2100            1
2101        } else if !token_accounts_with_mint.is_empty() {
2102            self.rng
2103                .gen_range(1..std::cmp::min(token_accounts_with_mint.len(), 4))
2104        } else {
2105            return (mint, Vec::new());
2106        };
2107        let mut get_random_subset_of_token_accounts =
2108            token_accounts_with_mint[0..range_end].to_vec();
2109        // Sorting input and output Merkle tree pubkeys the same way so the pubkey indices do not get out of order
2110        get_random_subset_of_token_accounts.sort_by(|a, b| {
2111            a.compressed_account
2112                .merkle_context
2113                .merkle_tree_pubkey
2114                .cmp(&b.compressed_account.merkle_context.merkle_tree_pubkey)
2115        });
2116        (mint, get_random_subset_of_token_accounts)
2117    }
2118}
2119
2120// Configures probabilities for keypair actions
2121// default sol configuration is all sol actions enabled with 0.5 probability
2122pub struct KeypairActionConfig {
2123    pub compress_sol: Option<f64>,
2124    pub decompress_sol: Option<f64>,
2125    pub transfer_sol: Option<f64>,
2126    pub create_address: Option<f64>,
2127    pub compress_spl: Option<f64>,
2128    pub decompress_spl: Option<f64>,
2129    pub mint_spl: Option<f64>,
2130    pub transfer_spl: Option<f64>,
2131    pub max_output_accounts: Option<u64>,
2132    pub fee_assert: bool,
2133    pub approve_spl: Option<f64>,
2134    pub revoke_spl: Option<f64>,
2135    pub freeze_spl: Option<f64>,
2136    pub thaw_spl: Option<f64>,
2137    pub burn_spl: Option<f64>,
2138}
2139
2140impl KeypairActionConfig {
2141    pub fn inclusion(&self) -> bool {
2142        self.transfer_sol.is_some() || self.transfer_spl.is_some()
2143    }
2144
2145    pub fn non_inclusion(&self) -> bool {
2146        self.create_address.is_some()
2147    }
2148
2149    pub fn sol_default() -> Self {
2150        Self {
2151            compress_sol: Some(0.5),
2152            decompress_sol: Some(0.5),
2153            transfer_sol: Some(0.5),
2154            create_address: None,
2155            compress_spl: None,
2156            decompress_spl: None,
2157            mint_spl: None,
2158            transfer_spl: None,
2159            max_output_accounts: None,
2160            fee_assert: true,
2161            approve_spl: None,
2162            revoke_spl: None,
2163            freeze_spl: None,
2164            thaw_spl: None,
2165            burn_spl: None,
2166        }
2167    }
2168
2169    pub fn spl_default() -> Self {
2170        Self {
2171            compress_sol: None,
2172            decompress_sol: None,
2173            transfer_sol: None,
2174            create_address: None,
2175            compress_spl: Some(0.7),
2176            decompress_spl: Some(0.5),
2177            mint_spl: None,
2178            transfer_spl: Some(0.5),
2179            max_output_accounts: Some(10),
2180            fee_assert: true,
2181            approve_spl: Some(0.5),
2182            revoke_spl: Some(0.5),
2183            freeze_spl: Some(0.5),
2184            thaw_spl: Some(0.5),
2185            burn_spl: Some(0.5),
2186        }
2187    }
2188
2189    pub fn all_default() -> Self {
2190        Self {
2191            compress_sol: Some(0.5),
2192            decompress_sol: Some(1.0),
2193            transfer_sol: Some(1.0),
2194            create_address: Some(0.2),
2195            compress_spl: Some(0.7),
2196            decompress_spl: Some(0.5),
2197            mint_spl: None,
2198            transfer_spl: Some(0.5),
2199            max_output_accounts: Some(10),
2200            fee_assert: true,
2201            approve_spl: Some(0.7),
2202            revoke_spl: Some(0.7),
2203            freeze_spl: Some(0.7),
2204            thaw_spl: Some(0.7),
2205            burn_spl: Some(0.7),
2206        }
2207    }
2208    pub fn all_default_no_fee_assert() -> Self {
2209        Self {
2210            compress_sol: Some(0.5),
2211            decompress_sol: Some(1.0),
2212            transfer_sol: Some(1.0),
2213            create_address: Some(0.2),
2214            compress_spl: Some(0.7),
2215            decompress_spl: Some(0.5),
2216            mint_spl: None,
2217            transfer_spl: Some(0.5),
2218            max_output_accounts: Some(10),
2219            fee_assert: false,
2220            approve_spl: Some(0.7),
2221            revoke_spl: Some(0.7),
2222            freeze_spl: Some(0.7),
2223            thaw_spl: Some(0.7),
2224            burn_spl: Some(0.7),
2225        }
2226    }
2227
2228    pub fn test_default() -> Self {
2229        Self {
2230            compress_sol: Some(1.0),
2231            decompress_sol: Some(1.0),
2232            transfer_sol: Some(1.0),
2233            create_address: Some(1.0),
2234            compress_spl: Some(0.0),
2235            decompress_spl: Some(0.0),
2236            mint_spl: None,
2237            transfer_spl: Some(0.0),
2238            max_output_accounts: Some(10),
2239            fee_assert: true,
2240            approve_spl: None,
2241            revoke_spl: None,
2242            freeze_spl: None,
2243            thaw_spl: None,
2244            burn_spl: None,
2245        }
2246    }
2247
2248    pub fn test_forester_default() -> Self {
2249        Self {
2250            compress_sol: Some(0.0),
2251            decompress_sol: Some(0.0),
2252            transfer_sol: Some(1.0),
2253            create_address: None,
2254            compress_spl: None,
2255            decompress_spl: None,
2256            mint_spl: None,
2257            transfer_spl: None,
2258            max_output_accounts: Some(3),
2259            fee_assert: true,
2260            approve_spl: None,
2261            revoke_spl: None,
2262            freeze_spl: None,
2263            thaw_spl: None,
2264            burn_spl: None,
2265        }
2266    }
2267}
2268
2269// Configures probabilities for general actions
2270pub struct GeneralActionConfig {
2271    pub add_keypair: Option<f64>,
2272    pub create_state_mt: Option<f64>,
2273    pub create_address_mt: Option<f64>,
2274    pub nullify_compressed_accounts: Option<f64>,
2275    pub empty_address_queue: Option<f64>,
2276    pub rollover: Option<f64>,
2277    pub add_forester: Option<f64>,
2278    /// TODO: add this
2279    /// Creates one infinte epoch
2280    pub disable_epochs: bool,
2281}
2282impl Default for GeneralActionConfig {
2283    fn default() -> Self {
2284        Self {
2285            add_keypair: Some(0.3),
2286            create_state_mt: Some(1.0),
2287            create_address_mt: Some(1.0),
2288            nullify_compressed_accounts: Some(0.2),
2289            empty_address_queue: Some(0.2),
2290            rollover: None,
2291            add_forester: None,
2292            disable_epochs: false,
2293        }
2294    }
2295}
2296
2297impl GeneralActionConfig {
2298    pub fn test_forester_default() -> Self {
2299        Self {
2300            add_keypair: None,
2301            create_state_mt: None,
2302            create_address_mt: None,
2303            nullify_compressed_accounts: None,
2304            empty_address_queue: None,
2305            rollover: None,
2306            add_forester: None,
2307            disable_epochs: false,
2308        }
2309    }
2310    pub fn test_with_rollover() -> Self {
2311        Self {
2312            add_keypair: Some(0.3),
2313            create_state_mt: Some(1.0),
2314            create_address_mt: Some(1.0),
2315            nullify_compressed_accounts: Some(0.2),
2316            empty_address_queue: Some(0.2),
2317            rollover: Some(0.5),
2318            add_forester: None,
2319            disable_epochs: false,
2320        }
2321    }
2322}