Pool: Split deposit should target the new pool ratio

contracts/pool/src/contract.rs

@@ -17,10 +17,7 @@ use crate::{
     token_contract,
 };
 use decimal::Decimal;
-use phoenix::{
-    utils::{is_approx_ratio, LiquidityPoolInitInfo},
-    validate_bps, validate_int_parameters,
-};
+use phoenix::{utils::LiquidityPoolInitInfo, validate_bps, validate_int_parameters};
 
 // Metadata that is added on to the WASM custom section
 contractmeta!(
@@ -295,7 +292,7 @@ impl LiquidityPoolTrait for LiquidityPool {
                     a,
                     &config.token_a,
                 );
-                do_swap(
+                let actual_b_from_swap = do_swap(
                     env.clone(),
                     sender.clone(),
                     // FIXM: Disable Referral struct
@@ -307,7 +304,16 @@ impl LiquidityPoolTrait for LiquidityPool {
                     None,
                 );
                 // return: rest of Token A amount, simulated result of swap of portion A
-                (a - a_for_swap, b_from_swap)
+                if (actual_b_from_swap - b_from_swap).abs() > 1 {
+                    log!(
+                        &env,
+                        "Pool: ProvideLiquidity: B token off by more than rounding error! actual: {}, predicted: {}",
+                        actual_b_from_swap,
+                        b_from_swap
+                    );
+                    panic_with_error!(env, ContractError::OffByMoreThanRoundingError)
+                }
+                (a - a_for_swap, actual_b_from_swap)
             }
             // Only token B is provided
             (None, Some(b)) if b > 0 => {
@@ -319,7 +325,7 @@ impl LiquidityPoolTrait for LiquidityPool {
                     b,
                     &config.token_b,
                 );
-                do_swap(
+                let actual_a_from_swap = do_swap(
                     env.clone(),
                     sender.clone(),
                     // FIXM: Disable Referral struct
@@ -331,7 +337,16 @@ impl LiquidityPoolTrait for LiquidityPool {
                     None,
                 );
                 // return: simulated result of swap of portion B, rest of Token B amount
-                (a_from_swap, b - b_for_swap)
+                if (actual_a_from_swap - a_from_swap).abs() > 1 {
+                    log!(
+                        &env,
+                        "Pool: ProvideLiquidity: A token off by more than rounding error! actual: {}, predicted: {}",
+                        actual_a_from_swap,
+                        a_from_swap
+                    );
+                    panic_with_error!(env, ContractError::OffByMoreThanRoundingError)
+                }
+                (actual_a_from_swap, b - b_for_swap)
             }
             // None or invalid amounts are provided
             _ => {
@@ -845,10 +860,10 @@ fn do_swap(
 
 /// This function divides the deposit in such a way that when swapping it for the other token,
 /// the resulting amounts of tokens maintain the current pool's ratio.
+/// * `config` - The configuration of the liquidity pool.
 /// * `a_pool` - The current amount of Token A in the liquidity pool.
 /// * `b_pool` - The current amount of Token B in the liquidity pool.
 /// * `deposit` - The total amount of tokens that the user wants to deposit into the liquidity pool.
-/// * `sell_a` - A boolean that indicates whether the deposit is in Token A (if true) or in Token B (if false).
 /// # Returns
 /// * A tuple `(final_offer_amount, final_ask_amount)`, where `final_offer_amount` is the amount of deposit tokens
 ///   to be swapped, and `final_ask_amount` is the amount of the other tokens that will be received in return.
@@ -878,57 +893,45 @@ fn split_deposit_based_on_pool_ratio(
         );
     }
 
-    // Calculate the current ratio in the pool
-    let target_ratio = Decimal::from_ratio(b_pool, a_pool);
-    // Define boundaries for binary search algorithm
-    let mut low = 0;
-    let mut high = deposit;
-
-    // Tolerance is the smallest difference in deposit that we care about
-    let tolerance = 500;
+    // get pool fee rate
+    let fee = config.protocol_fee_rate();
+    // determine which pool is offer and which is ask
+    let (offer_pool, ask_pool) = if offer_asset == &config.token_a {
+        (a_pool, b_pool)
+    } else {
+        (b_pool, a_pool)
+    };
 
-    let mut final_offer_amount = deposit; // amount of deposit tokens to be swapped
-    let mut final_ask_amount = 0; // amount of other tokens to be received
+    // formula to calculate final_offer_amount
+    let final_offer_amount = {
+        let numerator = deposit * fee - 2 * offer_pool
+            + (deposit * deposit * fee * fee
+                + 4 * deposit * offer_pool
+                + 4 * offer_pool * offer_pool)
+                .sqrt();
+        let denominator = 2 * (fee + Decimal::one());
+        numerator / denominator
+    };
 
-    while high - low > tolerance {
-        let mid = (low + high) / 2; // Calculate middle point
+    // formula to calculate final_ask_amount
+    // we need to handle the fee here as well
+    // we don't change ask_pool offer_pool values based on the fee prior to this method
+    let final_ask_amount = {
+        let numerator = ask_pool * final_offer_amount;
+        let denominator = offer_pool + final_offer_amount;
+        numerator / denominator
+    };
 
-        // Simulate swap to get amount of other tokens to be received for `mid` amount of deposit tokens
-        let SimulateSwapResponse {
-            ask_amount,
-            spread_amount: _,
-            commission_amount: _,
-            total_return: _,
-        } = LiquidityPool::simulate_swap(env.clone(), offer_asset.clone(), mid);
-
-        // Update final amounts
-        final_offer_amount = mid;
-        final_ask_amount = ask_amount;
-
-        // Calculate the ratio that would result from swapping `mid` deposit tokens
-        let ratio = if offer_asset == &config.token_a {
-            Decimal::from_ratio(ask_amount, deposit - mid)
-        } else {
-            Decimal::from_ratio(deposit - mid, ask_amount)
-        };
+    log!(
+        &env,
+        "log",
+        a_pool,
+        b_pool,
+        deposit,
+        final_offer_amount,
+        final_ask_amount
+    );
 
-        // If the resulting ratio is approximately equal (1%) to the target ratio, break the loop
-        if is_approx_ratio(ratio, target_ratio, Decimal::percent(1)) {
-            break;
-        }
-        // Update boundaries for the next iteration of the binary search
-        if ratio > target_ratio {
-            if offer_asset == &config.token_a {
-                high = mid;
-            } else {
-                low = mid;
-            }
-        } else if offer_asset == &config.token_a {
-            low = mid;
-        } else {
-            high = mid;
-        };
-    }
     (final_offer_amount, final_ask_amount)
 }
 

contracts/pool/src/error.rs

@@ -27,6 +27,6 @@ pub enum ContractError {
     TokenABiggerThanTokenB = 18,
     InvalidBps = 19,
     SlippageInvalid = 20,
-
     SwapMinReceivedBiggerThanReturn = 21,
+    OffByMoreThanRoundingError = 22,
 }

contracts/pool/src/tests/liquidity.rs

@@ -380,7 +380,7 @@ fn provide_liqudity_single_asset_equal_with_fees() {
     let stake_manager = Address::generate(&env);
     let stake_owner = Address::generate(&env);
 
-    let swap_fees = 1_000i64; // 10% bps
+    let swap_fees = 1; // 10% bps
     let pool = deploy_liquidity_pool_contract(
         &env,
         None,
@@ -562,6 +562,7 @@ fn provide_liqudity_single_asset_one_third_with_fees() {
     token2.mint(&user1, &100_000);
     // providing liquidity with a single asset - token2
     pool.provide_liquidity(&user1, &None, &None, &Some(100_000), &None, &None);
+    soroban_sdk::testutils::arbitrary::std::dbg!("after2");
     // before swap : A(10_000_000), B(30_000_000)
     // since pool is 1/3 algorithm will split it around 15794/52734
     // swap 47_226k B for A = 17_548 (-10% fee = 15_793)

contracts/pool/src/tests/swap.rs

@@ -1,4 +1,5 @@
 extern crate std;
+use phoenix::assert_approx_eq;
 use pretty_assertions::assert_eq;
 use soroban_sdk::{
     symbol_short,
@@ -8,7 +9,10 @@ use soroban_sdk::{
 use test_case::test_case;
 
 use super::setup::{deploy_liquidity_pool_contract, deploy_token_contract};
-use crate::storage::{Asset, PoolResponse, SimulateReverseSwapResponse, SimulateSwapResponse};
+use crate::{
+    storage::{Asset, PoolResponse, SimulateReverseSwapResponse, SimulateSwapResponse},
+    token_contract,
+};
 use decimal::Decimal;
 
 #[test]
@@ -1032,3 +1036,68 @@ fn test_should_fail_when_invalid_ask_asset_min_amount() {
     let spread = 100i64; // 1% maximum spread allowed
     pool.swap(&user, &token1.address, &1, &Some(10), &Some(spread));
 }
+
+#[test_case(0; "when fee is 0 percent")]
+#[test_case(1_000; "when fee is 10 percent")]
+fn provide_liqudity_single_asset_poc_split_good_target(swap_fees: i64) {
+    let env = Env::default();
+    env.mock_all_auths();
+    env.budget().reset_unlimited();
+
+    let mut admin1 = Address::generate(&env);
+    let mut admin2 = Address::generate(&env);
+
+    let mut token1 = deploy_token_contract(&env, &admin1);
+    let mut token2 = deploy_token_contract(&env, &admin2);
+    if token2.address < token1.address {
+        std::mem::swap(&mut token1, &mut token2);
+        std::mem::swap(&mut admin1, &mut admin2);
+    }
+    let user1 = Address::generate(&env);
+    let user2 = Address::generate(&env);
+
+    let pool = deploy_liquidity_pool_contract(
+        &env,
+        None,
+        (&token1.address, &token2.address),
+        swap_fees,
+        None,
+        None,
+        None,
+        Address::generate(&env),
+        Address::generate(&env),
+    );
+
+    token1.mint(&user1, &10_000_000);
+    token2.mint(&user1, &10_000_000);
+
+    pool.provide_liquidity(
+        &user1,
+        &Some(10_000_000),
+        &Some(10_000_000),
+        &Some(10_000_000),
+        &Some(10_000_000),
+        &None,
+    );
+    assert_eq!(token1.balance(&pool.address), 10_000_000);
+    assert_eq!(token2.balance(&pool.address), 10_000_000);
+
+    token1.mint(&user2, &500_000);
+
+    let user2_token_a_balance_before = token1.balance(&user2);
+
+    pool.provide_liquidity(&user2, &Some(500_000), &None, &None, &None, &None);
+
+    let share_token = pool.query_share_token_address();
+
+    let user2_lp_balance = token_contract::Client::new(&env, &share_token).balance(&user2);
+
+    // @audit User 2 withdraw the liquidity by burning all its lp token balance.
+    let (_, amount_b) = pool.withdraw_liquidity(&user2, &user2_lp_balance, &1, &1);
+
+    pool.swap(&user2, &token2.address, &amount_b, &None, &None);
+
+    let user2_token_a_balance_after = token1.balance(&user2);
+
+    assert_approx_eq!(user2_token_a_balance_before, user2_token_a_balance_after, 5);
+}

packages/phoenix/src/utils.rs

@@ -38,6 +38,36 @@ pub fn is_approx_ratio(a: Decimal, b: Decimal, tolerance: Decimal) -> bool {
     diff <= tolerance
 }
 
+#[macro_export]
+macro_rules! assert_approx_eq {
+    ($a:expr, $b:expr) => {{
+        let eps = 1.0e-6;
+        let (a, b) = (&$a, &$b);
+        assert!(
+            (*a - *b).abs() < eps,
+            "assertion failed: `(left !== right)` \
+             (left: `{:?}`, right: `{:?}`, expect diff: `{:?}`, real diff: `{:?}`)",
+            *a,
+            *b,
+            eps,
+            (*a - *b).abs()
+        );
+    }};
+    ($a:expr, $b:expr, $eps:expr) => {{
+        let (a, b) = (&$a, &$b);
+        let eps = $eps;
+        assert!(
+            (*a - *b).abs() < eps,
+            "assertion failed: `(left !== right)` \
+             (left: `{:?}`, right: `{:?}`, expect diff: `{:?}`, real diff: `{:?}`)",
+            *a,
+            *b,
+            eps,
+            (*a - *b).abs()
+        );
+    }};
+}
+
 #[contracttype]
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct TokenInitInfo {