Validator list balance solution

programs/steward/src/constants.rs

@@ -8,6 +8,7 @@ pub const MAX_VALIDATORS: usize = 5_000;
 pub const BASIS_POINTS_MAX: u16 = 10_000;
 pub const COMMISSION_MAX: u8 = 100;
 pub const SORTED_INDEX_DEFAULT: u16 = u16::MAX;
+pub const LAMPORT_BALANCE_DEFAULT: u64 = u64::MAX;
 // Need at least 1% of slots remaining (4320 slots) to execute steps in state machine
 pub const EPOCH_PROGRESS_MAX: f64 = 0.99;
 // Cannot go more than 100 epochs without scoring

programs/steward/src/delegation.rs

@@ -1,6 +1,7 @@
 use anchor_lang::prelude::*;
 use spl_stake_pool::big_vec::BigVec;
 
+use crate::constants::LAMPORT_BALANCE_DEFAULT;
 use crate::events::DecreaseComponents;
 use crate::{
     errors::StewardError,
@@ -250,6 +251,7 @@ impl UnstakeState {
         // either to the target or to the previous balance before the deposit, whichever is lower in terms of total lamports unstaked
         if current_lamports > state.validator_lamport_balances[index]
             && self.stake_deposit_unstake_total < self.stake_deposit_unstake_cap
+            && state.validator_lamport_balances[index] != LAMPORT_BALANCE_DEFAULT
         {
             let lamports_above_target = current_lamports
                 .checked_sub(target_lamports)

programs/steward/src/instructions/mod.rs

@@ -16,6 +16,7 @@ pub mod realloc_state;
 pub mod rebalance;
 pub mod remove_validator_from_blacklist;
 pub mod reset_steward_state;
+pub mod reset_validator_lamport_balances;
 pub mod resume_steward;
 pub mod set_new_authority;
 pub mod spl_passthrough;
@@ -38,6 +39,7 @@ pub use realloc_state::*;
 pub use rebalance::*;
 pub use remove_validator_from_blacklist::*;
 pub use reset_steward_state::*;
+pub use reset_validator_lamport_balances::*;
 pub use resume_steward::*;
 pub use set_new_authority::*;
 pub use spl_passthrough::*;

programs/steward/src/instructions/realloc_state.rs

@@ -1,6 +1,6 @@
 use crate::{
     bitmask::BitMask,
-    constants::{MAX_ALLOC_BYTES, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_ALLOC_BYTES, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
     errors::StewardError,
     state::{Config, StewardStateAccount},
     utils::get_validator_list,
@@ -74,6 +74,7 @@ pub fn handler(ctx: Context<ReallocState>) -> Result<()> {
 
         state_account.state.state_tag = StewardStateEnum::ComputeScores;
         state_account.state.num_pool_validators = validator_list.len() as u64;
+        state_account.state.validator_lamport_balances = [LAMPORT_BALANCE_DEFAULT; MAX_VALIDATORS];
         state_account.state.scores = [0; MAX_VALIDATORS];
         state_account.state.sorted_score_indices = [SORTED_INDEX_DEFAULT; MAX_VALIDATORS];
         state_account.state.yield_scores = [0; MAX_VALIDATORS];

programs/steward/src/instructions/reset_steward_state.rs

@@ -1,5 +1,5 @@
 use crate::{
-    constants::{MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
     errors::StewardError,
     state::{Config, StewardStateAccount},
     utils::{deserialize_stake_pool, get_config_admin, get_stake_pool_address},
@@ -47,6 +47,7 @@ pub fn handler(ctx: Context<ResetStewardState>) -> Result<()> {
 
     state_account.state.state_tag = StewardStateEnum::ComputeScores;
     state_account.state.num_pool_validators = validator_list.len() as u64;
+    state_account.state.validator_lamport_balances = [LAMPORT_BALANCE_DEFAULT; MAX_VALIDATORS];
     state_account.state.scores = [0; MAX_VALIDATORS];
     state_account.state.sorted_score_indices = [SORTED_INDEX_DEFAULT; MAX_VALIDATORS];
     state_account.state.yield_scores = [0; MAX_VALIDATORS];

programs/steward/src/instructions/reset_validator_lamport_balances.rs

@@ -0,0 +1,29 @@
+use crate::{
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS},
+    state::*,
+    utils::get_config_admin,
+};
+use anchor_lang::prelude::*;
+
+#[derive(Accounts)]
+pub struct ResetValidatorLamportBalances<'info> {
+    #[account(
+        mut,
+        seeds = [StewardStateAccount::SEED, config.key().as_ref()],
+        bump
+    )]
+    pub steward_state: AccountLoader<'info, StewardStateAccount>,
+
+    pub config: AccountLoader<'info, Config>,
+
+    #[account(address = get_config_admin(&config)?)]
+    pub authority: Signer<'info>,
+}
+
+pub fn handler(ctx: Context<ResetValidatorLamportBalances>) -> Result<()> {
+    let state_account = &mut ctx.accounts.steward_state.load_mut()?;
+
+    state_account.state.validator_lamport_balances = [LAMPORT_BALANCE_DEFAULT; MAX_VALIDATORS];
+
+    Ok(())
+}

programs/steward/src/instructions/spl_passthrough.rs

@@ -4,7 +4,7 @@
 // is that the config, stake pool address, staker, signer, and sometimes state account match up.
 // Otherwise these instructions are intented to be minimally restrictive.
 
-use crate::constants::MAX_VALIDATORS;
+use crate::constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS};
 use crate::errors::StewardError;
 use crate::state::Config;
 use crate::utils::{
@@ -400,9 +400,11 @@ pub fn increase_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_add(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_add(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(
@@ -521,9 +523,11 @@ pub fn decrease_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_sub(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_sub(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(
@@ -641,9 +645,11 @@ pub fn increase_additional_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_add(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_add(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(
@@ -765,9 +771,11 @@ pub fn decrease_additional_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_sub(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_sub(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(

programs/steward/src/lib.rs

@@ -207,6 +207,13 @@ pub mod steward {
         )
     }
 
+    /// Reset validator_lamport_balances to default
+    pub fn reset_validator_lamport_balances(
+        ctx: Context<ResetValidatorLamportBalances>,
+    ) -> Result<()> {
+        instructions::reset_validator_lamport_balances::handler(ctx)
+    }
+
     /// Closes Steward PDA accounts associated with a given Config (StewardStateAccount, and Staker).
     /// Config is not closed as it is a Keypair, so lamports can simply be withdrawn.
     /// Reclaims lamports to authority

programs/steward/src/state/steward_state.rs

@@ -3,7 +3,7 @@ use std::fmt::Display;
 
 use crate::{
     bitmask::BitMask,
-    constants::{MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
     delegation::{
         decrease_stake_calculation, increase_stake_calculation, RebalanceType, UnstakeState,
     },
@@ -894,6 +894,11 @@ impl StewardState {
                 In all cases where the current_lamports is now below the target or internal balance, we update the internal balance.
                 Otherwise, keep the internal balance the same to ensure we still see the stake deposit delta, until it can be unstaked.
                 */
+
+                if self.validator_lamport_balances[index] == LAMPORT_BALANCE_DEFAULT {
+                    self.validator_lamport_balances[index] = current_lamports;
+                }
+
                 self.validator_lamport_balances[index] = match (
                     current_lamports < self.validator_lamport_balances[index],
                     current_lamports < target_lamports,
@@ -1005,9 +1010,12 @@ impl StewardState {
                     }
                 }
                 RebalanceType::Increase(amount) => {
-                    self.validator_lamport_balances[index] = self.validator_lamport_balances[index]
-                        .checked_add(amount)
-                        .ok_or(StewardError::ArithmeticError)?;
+                    if self.validator_lamport_balances[index] != LAMPORT_BALANCE_DEFAULT {
+                        self.validator_lamport_balances[index] = self.validator_lamport_balances
+                            [index]
+                            .checked_add(amount)
+                            .ok_or(StewardError::ArithmeticError)?;
+                    }
                 }
                 RebalanceType::None => {}
             }

tests/tests/steward/test_spl_passthrough.rs

@@ -5,7 +5,7 @@ use anchor_lang::{
     AccountDeserialize, AnchorDeserialize, InstructionData, ToAccountMetas,
 };
 use jito_steward::{
-    constants::MAX_VALIDATORS,
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS},
     derive_steward_state_address,
     utils::{StakePool, ValidatorList},
     Config, Delegation, StewardStateAccount, StewardStateEnum,
@@ -300,7 +300,7 @@ async fn _increase_and_check_stake(
 
     let state_account_before: StewardStateAccount =
         fixture.load_and_deserialize(&fixture.steward_state).await;
-    let lamports_before_increase = *state_account_before
+    let _lamports_before_increase = *state_account_before
         .state
         .validator_lamport_balances
         .get(validator_list_index)
@@ -354,10 +354,7 @@ async fn _increase_and_check_stake(
         .get(validator_list_index)
         .expect("Lamport balance out of bounds");
 
-    assert_eq!(
-        lamports_after_increase,
-        lamports_before_increase + lamports_to_stake
-    );
+    assert_eq!(lamports_after_increase, LAMPORT_BALANCE_DEFAULT);
 }
 
 async fn _increase_and_check_additional_stake(
@@ -386,7 +383,7 @@ async fn _increase_and_check_additional_stake(
 
     let state_account_before: StewardStateAccount =
         fixture.load_and_deserialize(&fixture.steward_state).await;
-    let lamports_before_increase = *state_account_before
+    let _lamports_before_increase = *state_account_before
         .state
         .validator_lamport_balances
         .get(validator_list_index)
@@ -447,10 +444,7 @@ async fn _increase_and_check_additional_stake(
         .get(validator_list_index)
         .expect("Lamport balance out of bounds");
 
-    assert_eq!(
-        lamports_after_increase,
-        lamports_before_increase + lamports_to_stake
-    );
+    assert_eq!(lamports_after_increase, LAMPORT_BALANCE_DEFAULT);
 }
 
 pub async fn _set_staker(fixture: &TestFixture, new_staker: Pubkey) {