[LST-415] fix rebalance

controller/core/src/sync_sol_val.rs

@@ -16,7 +16,9 @@ pub struct SyncSolVal {
 }
 
 impl SyncSolVal {
-    /// Returns new pool total SOL value
+    /// # Returns
+    /// New pool total SOL value.
+    /// `None` on overflow
     ///
     /// This is rly just a wrapper for return
     /// `old_pool_total_sol_value - self.lst_sol_val.old() + self.lst_sol_val.new()`
@@ -34,8 +36,6 @@ impl SyncSolVal {
 impl PoolSvLamports {
     /// Applies a [`SyncSolVal`] followed by an [`UpdateYield`] based on the changes
     /// the sync made.
-    ///
-    /// Assumes INF mint supply did not change
     #[inline]
     pub const fn aft_ssv_uy(self, sync: &SyncSolVal) -> Option<Self> {
         let new_total_sol_value = match sync.exec(*self.total()) {
@@ -50,27 +50,7 @@ impl PoolSvLamports {
     }
 }
 
-impl PoolSvMutRefs<'_> {
-    /// Returns None on overflow
-    #[inline]
-    pub const fn apply_sync_sol_val(&mut self, sync: &SyncSolVal) -> Option<&mut Self> {
-        let new_total = match sync.exec(**self.total()) {
-            None => return None,
-            Some(nt) => nt,
-        };
-        **self.total_mut() = new_total;
-        Some(self)
-    }
-}
-
 impl PoolStateV2 {
-    /// Returns None on overflow
-    #[inline]
-    pub fn apply_sync_sol_val(&mut self, sync: &SyncSolVal) -> Option<&mut Self> {
-        PoolSvMutRefs::from_pool_state_v2(self).apply_sync_sol_val(sync)?;
-        Some(self)
-    }
-
     /// Applies a [`SyncSolVal`] followed by an [`UpdateYield`] based on the changes
     /// the sync made.
     ///

controller/program/Cargo.toml

@@ -40,6 +40,7 @@ inf1-pp-core = { workspace = true }
 inf1-pp-flatslab-std = { workspace = true }
 inf1-test-utils = { workspace = true }
 mollusk-svm = { workspace = true }
+mollusk-svm-programs-token = { workspace = true }
 proptest = { workspace = true, features = ["std"] }
 solana-account = { workspace = true }
 solana-instruction = { workspace = true }

controller/program/src/instructions/admin/set_sol_value_calculator.rs

@@ -1,5 +1,8 @@
 use inf1_ctl_jiminy::{
-    account_utils::{lst_state_list_checked, lst_state_list_checked_mut, pool_state_v2_checked},
+    account_utils::{
+        lst_state_list_checked, lst_state_list_checked_mut, pool_state_v2_checked,
+        pool_state_v2_checked_mut,
+    },
     cpi::SetSolValueCalculatorIxPreAccountHandles,
     err::Inf1CtlErr,
     instructions::{
@@ -20,9 +23,10 @@ use jiminy_cpi::{
 
 use inf1_core::instructions::admin::set_sol_value_calculator::SetSolValueCalculatorIxAccs;
 use inf1_core::instructions::sync_sol_value::SyncSolValueIxAccs;
+use jiminy_sysvar_clock::Clock;
 
 use crate::{
-    svc::lst_sync_sol_val,
+    svc::lst_ssv_uy,
     utils::split_suf_accs,
     verify::{
         verify_not_rebalancing_and_not_disabled, verify_pks, verify_signers,
@@ -94,18 +98,23 @@ pub fn process_set_sol_value_calculator(
         calc,
     }: &SetSolValueCalculatorIxAccounts,
     lst_idx: usize,
+    clock: &Clock,
 ) -> Result<(), ProgramError> {
-    let calc_key = *abr.get(*calc_prog).key();
+    pool_state_v2_checked_mut(abr.get_mut(*ix_prefix.pool_state()))?
+        .release_yield(clock.slot)
+        .map_err(Inf1CtlCustomProgErr)?;
+
+    let new_calc_prog = *abr.get(*calc_prog).key();
 
     let list = lst_state_list_checked_mut(abr.get_mut(*ix_prefix.lst_state_list()))?;
     let lst_state = list
         .0
         .get_mut(lst_idx)
         .ok_or(Inf1CtlCustomProgErr(Inf1CtlErr::InvalidLstIndex))?;
 
-    lst_state.sol_value_calculator = calc_key;
+    lst_state.sol_value_calculator = new_calc_prog;
 
-    lst_sync_sol_val(
+    lst_ssv_uy(
         abr,
         cpi,
         &SyncSolValueIxAccs {
@@ -115,7 +124,7 @@ pub fn process_set_sol_value_calculator(
                 .with_lst_state_list(*ix_prefix.lst_state_list())
                 .with_pool_reserves(*ix_prefix.pool_reserves())
                 .build(),
-            calc_prog: *abr.get(*calc_prog).key(),
+            calc_prog: new_calc_prog,
             calc,
         },
         lst_idx,

controller/program/src/instructions/rebalance/end.rs

@@ -14,20 +14,25 @@ use inf1_ctl_jiminy::{
     keys::{LST_STATE_LIST_ID, POOL_STATE_ID, REBALANCE_RECORD_ID},
     pda_onchain::create_raw_pool_reserves_addr,
     program_err::Inf1CtlCustomProgErr,
-    typedefs::u8bool::U8BoolMut,
+    sync_sol_val::SyncSolVal,
+    typedefs::{
+        pool_sv::{PoolSvLamports, PoolSvMutRefs},
+        u8bool::U8BoolMut,
+    },
+    yields::update::UpdateYield,
 };
 use jiminy_cpi::{
     account::{Abr, AccountHandle},
-    program_error::{ProgramError, NOT_ENOUGH_ACCOUNT_KEYS},
+    program_error::ProgramError,
 };
 
 use inf1_core::instructions::{
     rebalance::end::EndRebalanceIxAccs, sync_sol_value::SyncSolValueIxAccs,
 };
 
 use crate::{
-    svc::lst_sync_sol_val,
-    utils::split_suf_accs,
+    svc::{cpi_lst_reserves_sol_val, update_lst_state_sol_val},
+    utils::{accs_split_first_chunk, split_suf_accs},
     verify::{verify_is_rebalancing, verify_pks, verify_signers},
     Cpi,
 };
@@ -42,9 +47,7 @@ fn end_rebalance_accs_checked<'a, 'acc>(
     abr: &Abr,
     accounts: &'a [AccountHandle<'acc>],
 ) -> Result<EndRebalanceIxAccounts<'a, 'acc>, ProgramError> {
-    let (ix_prefix, suf) = accounts
-        .split_first_chunk()
-        .ok_or(NOT_ENOUGH_ACCOUNT_KEYS)?;
+    let (ix_prefix, suf) = accs_split_first_chunk(accounts)?;
     let ix_prefix = EndRebalanceIxPreAccs(*ix_prefix);
 
     let pool = pool_state_v2_checked(abr.get(*ix_prefix.pool_state()))?;
@@ -105,17 +108,14 @@ pub fn process_end_rebalance(
         inp_calc,
     } = end_rebalance_accs_checked(abr, accounts)?;
 
-    let pool_acc = abr.get_mut(*ix_prefix.pool_state());
-    let pool = pool_state_v2_checked_mut(pool_acc)?;
-    U8BoolMut(&mut pool.is_rebalancing).set_false();
-
     let (old_total_sol_value, inp_lst_idx) = {
         let rr = rebalance_record_checked(abr.get(*ix_prefix.rebalance_record()))?;
-
         (rr.old_total_sol_value, rr.inp_lst_index as usize)
     };
 
-    lst_sync_sol_val(
+    abr.close(*ix_prefix.rebalance_record(), *ix_prefix.pool_state())?;
+
+    let inp_lst_new = cpi_lst_reserves_sol_val(
         abr,
         cpi,
         &SyncSolValueIxAccs {
@@ -128,19 +128,35 @@ pub fn process_end_rebalance(
             calc_prog: *abr.get(inp_calc_prog).key(),
             calc: inp_calc,
         },
-        inp_lst_idx,
     )?;
+    let inp_sol_val =
+        update_lst_state_sol_val(abr, *ix_prefix.lst_state_list(), inp_lst_idx, inp_lst_new)?;
 
-    let new_total_sol_value = {
-        let pool = pool_state_v2_checked(abr.get(*ix_prefix.pool_state()))?;
-        pool.total_sol_value
-    };
+    let pool_acc = abr.get_mut(*ix_prefix.pool_state());
+    let pool = pool_state_v2_checked_mut(pool_acc)?;
+
+    U8BoolMut(&mut pool.is_rebalancing).set_false();
+
+    let new_total_sol_value = SyncSolVal {
+        lst_sol_val: inp_sol_val,
+    }
+    .exec(pool.total_sol_value)
+    .ok_or(Inf1CtlCustomProgErr(Inf1CtlErr::MathError))?;
 
     if new_total_sol_value < old_total_sol_value {
         return Err(Inf1CtlCustomProgErr(Inf1CtlErr::PoolWouldLoseSolValue).into());
     }
 
-    abr.close(*ix_prefix.rebalance_record(), *ix_prefix.pool_state())?;
+    let new = UpdateYield {
+        new_total_sol_value,
+        old: PoolSvLamports::from_pool_state_v2(pool)
+            // compare against val stored in RebalanceRecord
+            .with_total(old_total_sol_value),
+    }
+    .exec()
+    .ok_or(Inf1CtlCustomProgErr(Inf1CtlErr::MathError))?;
+
+    PoolSvMutRefs::from_pool_state_v2(pool).update(new);
 
     Ok(())
 }