@@ -104,7 +104,7 @@ pub struct StakePool {
104104pub epoch_fee : Fee ,
105105
106106 /// Fee for next epoch
107- pub next_epoch_fee : Option < Fee > ,
107+ pub next_epoch_fee : FutureEpoch < Fee > ,
108108
109109 /// Preferred deposit validator vote account pubkey
110110pub preferred_deposit_validator_vote_address : Option < Pubkey > ,
@@ -119,7 +119,7 @@ pub struct StakePool {
119119pub stake_withdrawal_fee : Fee ,
120120
121121 /// Future stake withdrawal fee, to be set for the following epoch
122- pub next_stake_withdrawal_fee : Option < Fee > ,
122+ pub next_stake_withdrawal_fee : FutureEpoch < Fee > ,
123123
124124 /// Fees paid out to referrers on referred stake deposits.
125125/// Expressed as a percentage (0 - 100) of deposit fees.
@@ -148,7 +148,7 @@ pub struct StakePool {
148148pub sol_withdrawal_fee : Fee ,
149149
150150 /// Future SOL withdrawal fee, to be set for the following epoch
151- pub next_sol_withdrawal_fee : Option < Fee > ,
151+ pub next_sol_withdrawal_fee : FutureEpoch < Fee > ,
152152
153153 /// Last epoch's total pool tokens, used only for APR estimation
154154pub last_epoch_pool_token_supply : u64 ,
@@ -483,14 +483,14 @@ impl StakePool {
483483 match fee {
484484 FeeType :: SolReferral ( new_fee) => self . sol_referral_fee = * new_fee,
485485 FeeType :: StakeReferral ( new_fee) => self . stake_referral_fee = * new_fee,
486- FeeType :: Epoch ( new_fee) => self . next_epoch_fee = Some ( * new_fee) ,
486+ FeeType :: Epoch ( new_fee) => self . next_epoch_fee = FutureEpoch :: new ( * new_fee) ,
487487 FeeType :: StakeWithdrawal ( new_fee) => {
488488 new_fee. check_withdrawal ( & self . stake_withdrawal_fee ) ?;
489- self . next_stake_withdrawal_fee = Some ( * new_fee)
489+ self . next_stake_withdrawal_fee = FutureEpoch :: new ( * new_fee)
490490 }
491491 FeeType :: SolWithdrawal ( new_fee) => {
492492 new_fee. check_withdrawal ( & self . sol_withdrawal_fee ) ?;
493- self . next_sol_withdrawal_fee = Some ( * new_fee)
493+ self . next_sol_withdrawal_fee = FutureEpoch :: new ( * new_fee)
494494 }
495495 FeeType :: SolDeposit ( new_fee) => self . sol_deposit_fee = * new_fee,
496496 FeeType :: StakeDeposit ( new_fee) => self . stake_deposit_fee = * new_fee,
@@ -793,6 +793,62 @@ impl ValidatorListHeader {
793793 }
794794}
795795
796+ /// Wrapper type that "counts down" epochs, which is Borsh-compatible with the
797+ /// native `Option`
798+ #[ repr( C ) ]
799+ #[ derive( Clone , Copy , Debug , PartialEq , BorshSerialize , BorshDeserialize , BorshSchema ) ]
800+ pub enum FutureEpoch < T > {
801+ /// Nothing is set
802+ None ,
803+ /// Value is ready after the next epoch boundary
804+ One ( T ) ,
805+ /// Value is ready after two epoch boundaries
806+ Two ( T ) ,
807+ }
808+ impl < T > Default for FutureEpoch < T > {
809+ fn default ( ) -> Self {
810+ Self :: None
811+ }
812+ }
813+ impl < T > FutureEpoch < T > {
814+ /// Create a new value to be unlocked in a two epochs
815+ pub fn new ( value : T ) -> Self {
816+ Self :: Two ( value)
817+ }
818+ }
819+ impl < T : Clone > FutureEpoch < T > {
820+ /// Update the epoch, to be done after `get`ting the underlying value
821+ pub fn update_epoch ( & mut self ) {
822+ match self {
823+ Self :: None => { }
824+ Self :: One ( _) => {
825+ // The value has waited its last epoch
826+ * self = Self :: None ;
827+ }
828+ // The value still has to wait one more epoch after this
829+ Self :: Two ( v) => {
830+ * self = Self :: One ( v. clone ( ) ) ;
831+ }
832+ }
833+ }
834+
835+ /// Get the value if it's ready, which is only at `One` epoch remaining
836+ pub fn get ( & self ) -> Option < & T > {
837+ match self {
838+ Self :: None | Self :: Two ( _) => None ,
839+ Self :: One ( v) => Some ( v) ,
840+ }
841+ }
842+ }
843+ impl < T > From < FutureEpoch < T > > for Option < T > {
844+ fn from ( v : FutureEpoch < T > ) -> Option < T > {
845+ match v {
846+ FutureEpoch :: None => None ,
847+ FutureEpoch :: One ( inner) | FutureEpoch :: Two ( inner) => Some ( inner)
848+ }
849+ }
850+ }
851+
796852/// Fee rate as a ratio, minted on `UpdateStakePoolBalance` as a proportion of
797853/// the rewards
798854/// If either the numerator or the denominator is 0, the fee is considered to be 0
0 commit comments