stake-pool: make update_validator_list_balance() more ergonomic (#6058)

* remove unnecessary vote acc slice

* nightly fmt

* fix clippy

* remove redundant check

* restore deprecated update_validator_list_balance, rename update_validator_list_balance_chunk

* fix tests

* option -> result

* fmt

Author: billythedummy

stake-pool/program/src/instruction.rs

@@ -8,12 +8,13 @@ use {
         find_stake_program_address, find_transient_stake_program_address,
         find_withdraw_authority_program_address,
         inline_mpl_token_metadata::{self, pda::find_metadata_account},
-        state::{Fee, FeeType, StakePool, ValidatorList},
+        state::{Fee, FeeType, StakePool, ValidatorList, ValidatorStakeInfo},
         MAX_VALIDATORS_TO_UPDATE,
     },
     borsh::{BorshDeserialize, BorshSchema, BorshSerialize},
     solana_program::{
         instruction::{AccountMeta, Instruction},
+        program_error::ProgramError,
         pubkey::Pubkey,
         stake, system_program, sysvar,
     },
@@ -1363,6 +1364,10 @@ pub fn decrease_additional_validator_stake_with_vote(
 
 /// Creates `UpdateValidatorListBalance` instruction (update validator stake
 /// account balances)
+#[deprecated(
+    since = "1.1.0",
+    note = "please use `update_validator_list_balance_chunk`"
+)]
 pub fn update_validator_list_balance(
     program_id: &Pubkey,
     stake_pool: &Pubkey,
@@ -1423,6 +1428,74 @@ pub fn update_validator_list_balance(
     }
 }
 
+/// Creates an `UpdateValidatorListBalance` instruction (update validator stake
+/// account balances) to update `validator_list[start_index..start_index +
+/// len]`.
+///
+/// Returns `Err(ProgramError::InvalidInstructionData)` if:
+/// - `start_index..start_index + len` is out of bounds for
+///   `validator_list.validators`
+pub fn update_validator_list_balance_chunk(
+    program_id: &Pubkey,
+    stake_pool: &Pubkey,
+    stake_pool_withdraw_authority: &Pubkey,
+    validator_list_address: &Pubkey,
+    reserve_stake: &Pubkey,
+    validator_list: &ValidatorList,
+    len: usize,
+    start_index: usize,
+    no_merge: bool,
+) -> Result<Instruction, ProgramError> {
+    let mut accounts = vec![
+        AccountMeta::new_readonly(*stake_pool, false),
+        AccountMeta::new_readonly(*stake_pool_withdraw_authority, false),
+        AccountMeta::new(*validator_list_address, false),
+        AccountMeta::new(*reserve_stake, false),
+        AccountMeta::new_readonly(sysvar::clock::id(), false),
+        AccountMeta::new_readonly(sysvar::stake_history::id(), false),
+        AccountMeta::new_readonly(stake::program::id(), false),
+    ];
+    let validator_list_subslice = validator_list
+        .validators
+        .get(start_index..start_index.saturating_add(len))
+        .ok_or(ProgramError::InvalidInstructionData)?;
+    accounts.extend(validator_list_subslice.iter().flat_map(
+        |ValidatorStakeInfo {
+             vote_account_address,
+             validator_seed_suffix,
+             transient_seed_suffix,
+             ..
+         }| {
+            let (validator_stake_account, _) = find_stake_program_address(
+                program_id,
+                vote_account_address,
+                stake_pool,
+                NonZeroU32::new((*validator_seed_suffix).into()),
+            );
+            let (transient_stake_account, _) = find_transient_stake_program_address(
+                program_id,
+                vote_account_address,
+                stake_pool,
+                (*transient_seed_suffix).into(),
+            );
+            [
+                AccountMeta::new(validator_stake_account, false),
+                AccountMeta::new(transient_stake_account, false),
+            ]
+        },
+    ));
+    Ok(Instruction {
+        program_id: *program_id,
+        accounts,
+        data: StakePoolInstruction::UpdateValidatorListBalance {
+            start_index: start_index.try_into().unwrap(),
+            no_merge,
+        }
+        .try_to_vec()
+        .unwrap(),
+    })
+}
+
 /// Creates `UpdateStakePoolBalance` instruction (pool balance from the stake
 /// account list balances)
 pub fn update_stake_pool_balance(
@@ -1482,31 +1555,29 @@ pub fn update_stake_pool(
     stake_pool_address: &Pubkey,
     no_merge: bool,
 ) -> (Vec<Instruction>, Vec<Instruction>) {
-    let vote_accounts: Vec<Pubkey> = validator_list
-        .validators
-        .iter()
-        .map(|item| item.vote_account_address)
-        .collect();
-
     let (withdraw_authority, _) =
         find_withdraw_authority_program_address(program_id, stake_pool_address);
 
-    let mut update_list_instructions: Vec<Instruction> = vec![];
-    let mut start_index = 0;
-    for accounts_chunk in vote_accounts.chunks(MAX_VALIDATORS_TO_UPDATE) {
-        update_list_instructions.push(update_validator_list_balance(
-            program_id,
-            stake_pool_address,
-            &withdraw_authority,
-            &stake_pool.validator_list,
-            &stake_pool.reserve_stake,
-            validator_list,
-            accounts_chunk,
-            start_index,
-            no_merge,
-        ));
-        start_index = start_index.saturating_add(MAX_VALIDATORS_TO_UPDATE as u32);
-    }
+    let update_list_instructions = validator_list
+        .validators
+        .chunks(MAX_VALIDATORS_TO_UPDATE)
+        .enumerate()
+        .map(|(i, chunk)| {
+            // unwrap-safety: chunk len and offset are derived
+            update_validator_list_balance_chunk(
+                program_id,
+                stake_pool_address,
+                &withdraw_authority,
+                &stake_pool.validator_list,
+                &stake_pool.reserve_stake,
+                validator_list,
+                chunk.len(),
+                i.saturating_mul(MAX_VALIDATORS_TO_UPDATE),
+                no_merge,
+            )
+            .unwrap()
+        })
+        .collect();
 
     let final_instructions = vec![
         update_stake_pool_balance(

stake-pool/program/tests/decrease.rs

@@ -624,7 +624,6 @@ async fn fail_additional_with_increasing() {
             &mut context.banks_client,
             &context.payer,
             &last_blockhash,
-            &[validator_stake.vote.pubkey()],
             false,
         )
         .await;

stake-pool/program/tests/deposit.rs

@@ -94,7 +94,6 @@ async fn setup(
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &[validator_stake_account.vote.pubkey()],
             false,
         )
         .await;

stake-pool/program/tests/deposit_edge_cases.rs

@@ -87,7 +87,6 @@ async fn setup(
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &[validator_stake_account.vote.pubkey()],
             false,
         )
         .await;

stake-pool/program/tests/force_destake.rs

@@ -155,7 +155,6 @@ async fn success_update() {
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &[voter_pubkey],
             false,
         )
         .await;
@@ -272,7 +271,6 @@ async fn success_remove_validator() {
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &[voter_pubkey],
             false,
         )
         .await;
@@ -311,7 +309,6 @@ async fn success_remove_validator() {
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &[voter_pubkey],
             false,
         )
         .await;

stake-pool/program/tests/helpers/mod.rs

@@ -32,7 +32,7 @@ use {
         instruction, minimum_delegation,
         processor::Processor,
         state::{self, FeeType, FutureEpoch, StakePool, ValidatorList},
-        MINIMUM_RESERVE_LAMPORTS,
+        MAX_VALIDATORS_TO_UPDATE, MINIMUM_RESERVE_LAMPORTS,
     },
     spl_token_2022::{
         extension::{ExtensionType, StateWithExtensionsOwned},
@@ -1412,21 +1412,22 @@ impl StakePoolAccounts {
         banks_client: &mut BanksClient,
         payer: &Keypair,
         recent_blockhash: &Hash,
-        validator_vote_accounts: &[Pubkey],
+        len: usize,
         no_merge: bool,
     ) -> Option<TransportError> {
         let validator_list = self.get_validator_list(banks_client).await;
-        let mut instructions = vec![instruction::update_validator_list_balance(
+        let mut instructions = vec![instruction::update_validator_list_balance_chunk(
             &id(),
             &self.stake_pool.pubkey(),
             &self.withdraw_authority,
             &self.validator_list.pubkey(),
             &self.reserve_stake.pubkey(),
             &validator_list,
-            validator_vote_accounts,
+            len,
             0,
             no_merge,
-        )];
+        )
+        .unwrap()];
         self.maybe_add_compute_budget_instruction(&mut instructions);
         let transaction = Transaction::new_signed_with_payer(
             &instructions,
@@ -1501,22 +1502,31 @@ impl StakePoolAccounts {
         banks_client: &mut BanksClient,
         payer: &Keypair,
         recent_blockhash: &Hash,
-        validator_vote_accounts: &[Pubkey],
         no_merge: bool,
     ) -> Option<TransportError> {
         let validator_list = self.get_validator_list(banks_client).await;
-        let mut instructions = vec![
-            instruction::update_validator_list_balance(
-                &id(),
-                &self.stake_pool.pubkey(),
-                &self.withdraw_authority,
-                &self.validator_list.pubkey(),
-                &self.reserve_stake.pubkey(),
-                &validator_list,
-                validator_vote_accounts,
-                0,
-                no_merge,
-            ),
+        let mut instructions = vec![];
+        for (i, chunk) in validator_list
+            .validators
+            .chunks(MAX_VALIDATORS_TO_UPDATE)
+            .enumerate()
+        {
+            instructions.push(
+                instruction::update_validator_list_balance_chunk(
+                    &id(),
+                    &self.stake_pool.pubkey(),
+                    &self.withdraw_authority,
+                    &self.validator_list.pubkey(),
+                    &self.reserve_stake.pubkey(),
+                    &validator_list,
+                    chunk.len(),
+                    i * MAX_VALIDATORS_TO_UPDATE,
+                    no_merge,
+                )
+                .unwrap(),
+            );
+        }
+        instructions.extend([
             instruction::update_stake_pool_balance(
                 &id(),
                 &self.stake_pool.pubkey(),
@@ -1532,7 +1542,7 @@ impl StakePoolAccounts {
                 &self.stake_pool.pubkey(),
                 &self.validator_list.pubkey(),
             ),
-        ];
+        ]);
         self.maybe_add_compute_budget_instruction(&mut instructions);
         let transaction = Transaction::new_signed_with_payer(
             &instructions,

stake-pool/program/tests/huge_pool.rs

@@ -172,15 +172,15 @@ async fn setup(
 #[test_case(MAX_POOL_SIZE; "no-compute-budget")]
 #[tokio::test]
 async fn update(max_validators: u32) {
-    let (mut context, stake_pool_accounts, vote_account_pubkeys, _, _, _, _) =
+    let (mut context, stake_pool_accounts, _, _, _, _, _) =
         setup(max_validators, max_validators, STAKE_AMOUNT).await;
 
     let error = stake_pool_accounts
         .update_validator_list_balance(
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &vote_account_pubkeys[0..MAX_VALIDATORS_TO_UPDATE],
+            MAX_VALIDATORS_TO_UPDATE,
             false, /* no_merge */
         )
         .await;
@@ -335,23 +335,24 @@ async fn remove_validator_from_pool(max_validators: u32) {
             &mut context.banks_client,
             &context.payer,
             &context.last_blockhash,
-            &[first_vote],
+            1,
             false, /* no_merge */
         )
         .await;
     assert!(error.is_none(), "{:?}", error);
 
-    let mut instructions = vec![instruction::update_validator_list_balance(
+    let mut instructions = vec![instruction::update_validator_list_balance_chunk(
         &id(),
         &stake_pool_accounts.stake_pool.pubkey(),
         &stake_pool_accounts.withdraw_authority,
         &stake_pool_accounts.validator_list.pubkey(),
         &stake_pool_accounts.reserve_stake.pubkey(),
         &validator_list,
-        &[middle_vote],
-        middle_index as u32,
+        1,
+        middle_index,
         /* no_merge = */ false,
-    )];
+    )
+    .unwrap()];
     stake_pool_accounts.maybe_add_compute_budget_instruction(&mut instructions);
     let transaction = Transaction::new_signed_with_payer(
         &instructions,
@@ -366,17 +367,18 @@ async fn remove_validator_from_pool(max_validators: u32) {
         .err();
     assert!(error.is_none(), "{:?}", error);
 
-    let mut instructions = vec![instruction::update_validator_list_balance(
+    let mut instructions = vec![instruction::update_validator_list_balance_chunk(
         &id(),
         &stake_pool_accounts.stake_pool.pubkey(),
         &stake_pool_accounts.withdraw_authority,
         &stake_pool_accounts.validator_list.pubkey(),
         &stake_pool_accounts.reserve_stake.pubkey(),
         &validator_list,
-        &[last_vote],
-        last_index as u32,
+        1,
+        last_index,
         /* no_merge = */ false,
-    )];
+    )
+    .unwrap()];
     stake_pool_accounts.maybe_add_compute_budget_instruction(&mut instructions);
     let transaction = Transaction::new_signed_with_payer(
         &instructions,

stake-pool/program/tests/increase.rs

@@ -548,7 +548,6 @@ async fn fail_additional_with_decreasing() {
             &mut context.banks_client,
             &context.payer,
             &last_blockhash,
-            &[validator_stake.vote.pubkey()],
             false,
         )
         .await;