fix: Appropriate underlying liquidities (#4)

jup-ag · Sep 27, 2023 · 93a97a1 · 93a97a1
1 parent bc50a17
commit 93a97a1
Show file tree

Hide file tree

Showing 7 changed files with 99 additions and 80 deletions.
diff --git a/common/src/pool_pair.rs b/common/src/pool_pair.rs
@@ -125,14 +125,49 @@ pub fn get_account_metas<W: WithdrawStake + ?Sized, D: DepositStake + ?Sized>(
     Ok(metas)
 }
 
+fn prepare_underlying_liquidities(
+    underlying_liquidities: &[Option<&Pubkey>],
+) -> Option<HashSet<Pubkey>> {
+    let uls = HashSet::from_iter(
+        underlying_liquidities
+            .into_iter()
+            .filter_map(|ul| ul.cloned()),
+    );
+    if uls.len() > 0 {
+        Some(uls)
+    } else {
+        None
+    }
+}
+
 #[derive(Clone)]
 pub struct OneWayPoolPair<
     W: WithdrawStake + Clone + Send + Sync + 'static,
     D: DepositStake + Clone + Send + Sync + 'static,
 > {
     pub withdraw: W,
     pub deposit: D,
-    pub clock: Clock,
+    clock: Clock,
+    underlying_liquidities: Option<HashSet<Pubkey>>,
+}
+
+impl<W, D> OneWayPoolPair<W, D>
+where
+    W: WithdrawStake + Clone + Send + Sync,
+    D: DepositStake + Clone + Send + Sync,
+{
+    pub fn new(withdraw: W, deposit: D) -> Self {
+        let underlying_liquidities = prepare_underlying_liquidities(&[
+            withdraw.underlying_liquidity(),
+            deposit.underlying_liquidity(),
+        ]);
+        Self {
+            withdraw,
+            deposit,
+            clock: Clock::default(),
+            underlying_liquidities,
+        }
+    }
 }
 
 impl<W, D> Amm for OneWayPoolPair<W, D>
@@ -231,9 +266,7 @@ where
     }
 
     fn underlying_liquidities(&self) -> Option<HashSet<Pubkey>> {
-        self.deposit
-            .underlying_liquidity()
-            .map(|ul| HashSet::from([*ul]))
+        self.underlying_liquidities.clone()
     }
 
     fn program_dependencies(&self) -> Vec<(Pubkey, String)> {
@@ -257,7 +290,29 @@ pub struct TwoWayPoolPair<
 > {
     pub p1: P1,
     pub p2: P2,
-    pub clock: Clock,
+    clock: Clock,
+    underlying_liquidities: Option<HashSet<Pubkey>>,
+}
+
+impl<P1, P2> TwoWayPoolPair<P1, P2>
+where
+    P1: DepositStake + WithdrawStake + Clone + Send + Sync,
+    P2: DepositStake + WithdrawStake + Clone + Send + Sync,
+{
+    pub fn new(p1: P1, p2: P2) -> Self {
+        let underlying_liquidities = prepare_underlying_liquidities(&[
+            DepositStake::underlying_liquidity(&p1),
+            WithdrawStake::underlying_liquidity(&p1),
+            DepositStake::underlying_liquidity(&p2),
+            WithdrawStake::underlying_liquidity(&p2),
+        ]);
+        Self {
+            p1,
+            p2,
+            clock: Clock::default(),
+            underlying_liquidities,
+        }
+    }
 }
 
 impl<P1, P2> Amm for TwoWayPoolPair<P1, P2>
@@ -359,6 +414,10 @@ where
         1 + WithdrawStake::accounts_len(&self.p1) + DepositStake::accounts_len(&self.p2) + 1
     }
 
+    fn underlying_liquidities(&self) -> Option<HashSet<Pubkey>> {
+        self.underlying_liquidities.clone()
+    }
+
     fn program_dependencies(&self) -> Vec<(Pubkey, String)> {
         vec![
             (

diff --git a/common/src/withdraw_stake.rs b/common/src/withdraw_stake.rs
@@ -147,5 +147,9 @@ pub trait WithdrawStake: BaseStakePoolAmm {
 
     fn virtual_ix(&self, quote: &WithdrawStakeQuote) -> Result<Instruction>;
 
+    fn underlying_liquidity(&self) -> Option<&Pubkey> {
+        None
+    }
+
     fn accounts_len(&self) -> usize;
 }
diff --git a/libs/eversol_stake_pool/src/stakedex_traits.rs b/libs/eversol_stake_pool/src/stakedex_traits.rs
@@ -397,4 +397,8 @@ impl WithdrawStake for EversolStakePoolStakedex {
     fn accounts_len(&self) -> usize {
         EVERSOL_STAKE_POOL_WITHDRAW_STAKE_IX_ACCOUNTS_LEN
     }
+
+    fn underlying_liquidity(&self) -> Option<&Pubkey> {
+        Some(&eversol_stake_pool::ID)
+    }
 }
diff --git a/libs/lido/src/stakedex_traits.rs b/libs/lido/src/stakedex_traits.rs
@@ -199,6 +199,10 @@ impl WithdrawStake for LidoStakedex {
     fn accounts_len(&self) -> usize {
         LIDO_WITHDRAW_STAKE_IX_ACCOUNTS_LEN
     }
+
+    fn underlying_liquidity(&self) -> Option<&Pubkey> {
+        Some(&lido_state::ID)
+    }
 }
 
 impl DepositSol for LidoStakedex {

diff --git a/libs/socean_stake_pool/src/stakedex_traits.rs b/libs/socean_stake_pool/src/stakedex_traits.rs
@@ -385,4 +385,8 @@ impl WithdrawStake for SoceanStakePoolStakedex {
     fn accounts_len(&self) -> usize {
         SOCEAN_STAKE_POOL_WITHDRAW_STAKE_IX_ACCOUNTS_LEN
     }
+
+    fn underlying_liquidity(&self) -> Option<&Pubkey> {
+        Some(&socean_stake_pool::ID)
+    }
 }
diff --git a/libs/spl_stake_pool/src/stakedex_traits.rs b/libs/spl_stake_pool/src/stakedex_traits.rs
@@ -396,4 +396,8 @@ impl WithdrawStake for SplStakePoolStakedex {
     fn accounts_len(&self) -> usize {
         SPL_STAKE_POOL_WITHDRAW_STAKE_IX_ACCOUNTS_LEN
     }
+
+    fn underlying_liquidity(&self) -> Option<&Pubkey> {
+        Some(&self.stake_pool_addr)
+    }
 }
diff --git a/stakedex_sdk/src/lib.rs b/stakedex_sdk/src/lib.rs
@@ -572,109 +572,49 @@ impl Stakedex {
         for (first_stakedex, second_stakedex) in stakedexes.into_iter().tuple_combinations() {
             match (first_stakedex, second_stakedex) {
                 (Stakedex::SplStakePool(p1), Stakedex::SplStakePool(p2)) => {
-                    amms.push(Box::new(TwoWayPoolPair {
-                        p1,
-                        p2,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(TwoWayPoolPair::new(p1, p2)));
                 }
                 match_stakedexes!(SplStakePool, Socean, p1, p2) => {
-                    amms.push(Box::new(TwoWayPoolPair {
-                        p1,
-                        p2,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(TwoWayPoolPair::new(p1, p2)));
                 }
                 match_stakedexes!(SplStakePool, Eversol, p1, p2) => {
-                    amms.push(Box::new(TwoWayPoolPair {
-                        p1,
-                        p2,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(TwoWayPoolPair::new(p1, p2)));
                 }
                 match_stakedexes!(SplStakePool, Marinade, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(SplStakePool, UnstakeIt, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Socean, Eversol, p1, p2) => {
-                    amms.push(Box::new(TwoWayPoolPair {
-                        p1,
-                        p2,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(TwoWayPoolPair::new(p1, p2)));
                 }
                 match_stakedexes!(Socean, UnstakeIt, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Socean, Marinade, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Eversol, UnstakeIt, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Eversol, Marinade, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Lido, SplStakePool, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Lido, Socean, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Lido, Eversol, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Lido, UnstakeIt, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Lido, Marinade, withdraw, deposit) => {
-                    amms.push(Box::new(OneWayPoolPair {
-                        withdraw,
-                        deposit,
-                        clock: Clock::default(),
-                    }));
+                    amms.push(Box::new(OneWayPoolPair::new(withdraw, deposit)));
                 }
                 match_stakedexes!(Marinade, UnstakeIt, _, _) => (), // Cannot do anything with those two
                 match_same_stakedex!(Socean)