feat(smart-contracts): make provide liquidity for stableswap generic …

…so we can provide liquidity with 2, 3, or N assets

contracts/liquidity_hub/pool-manager/src/helpers.rs

@@ -668,54 +668,47 @@ pub fn get_asset_indexes_in_pool(
 
 // TODO: handle unwraps properly
 #[allow(clippy::unwrap_used)]
-pub fn compute_d(
-    amp_factor: &u64,
-    amount_a: Uint128,
-    amount_b: Uint128,
-    amount_c: Uint128,
-) -> Option<Uint256> {
-    let sum_x = amount_a
-        .checked_add(amount_b.checked_add(amount_c).unwrap())
-        .unwrap(); // sum(x_i), a.k.a S
+pub fn compute_d(amp_factor: &u64, deposits: &Vec<Coin>) -> Option<Uint256> {
+    let n_coins = Uint128::from(deposits.len() as u128);
+
+    // sum(x_i), a.k.a S
+    let sum_x = deposits
+        .iter()
+        .fold(Uint128::zero(), |acc, x| acc.checked_add(x.amount).unwrap())
+        .clone();
+
     if sum_x == Uint128::zero() {
         Some(Uint256::zero())
     } else {
-        let amount_a_times_coins = amount_a.checked_mul(N_COINS.into()).unwrap();
-        let amount_b_times_coins = amount_b.checked_mul(N_COINS.into()).unwrap();
-        let amount_c_times_coins = amount_c.checked_mul(N_COINS.into()).unwrap();
+        // do as below but for a generic number of assets
+        let amount_times_coins: Vec<Uint128> = deposits
+            .into_iter()
+            .map(|coin| coin.amount.checked_mul(n_coins).unwrap())
+            .collect();
 
         // Newton's method to approximate D
         let mut d_prev: Uint256;
         let mut d: Uint256 = sum_x.into();
-        for _ in 0..256 {
-            let mut d_prod = d;
-            d_prod = d_prod
-                .checked_mul(d)
-                .unwrap()
-                .checked_div(amount_a_times_coins.into())
-                .unwrap();
-            d_prod = d_prod
-                .checked_mul(d)
-                .unwrap()
-                .checked_div(amount_b_times_coins.into())
-                .unwrap();
-            d_prod = d_prod
-                .checked_mul(d)
-                .unwrap()
-                .checked_div(amount_c_times_coins.into())
-                .unwrap();
-            d_prev = d;
-            d = compute_next_d(amp_factor, d, d_prod, sum_x).unwrap();
-            // Equality with the precision of 1
-            if d > d_prev {
-                if d.checked_sub(d_prev).unwrap() <= Uint256::one() {
+        for amount in amount_times_coins.into_iter() {
+            for _ in 0..256 {
+                let mut d_prod = d;
+                d_prod = d_prod
+                    .checked_mul(d)
+                    .unwrap()
+                    .checked_div(amount.into())
+                    .unwrap();
+                d_prev = d;
+                d = compute_next_d(amp_factor, d, d_prod, sum_x).unwrap();
+                // Equality with the precision of 1
+                if d > d_prev {
+                    if d.checked_sub(d_prev).unwrap() <= Uint256::one() {
+                        break;
+                    }
+                } else if d_prev.checked_sub(d).unwrap() <= Uint256::one() {
                     break;
                 }
-            } else if d_prev.checked_sub(d).unwrap() <= Uint256::one() {
-                break;
             }
         }
-
         Some(d)
     }
 }
@@ -756,28 +749,28 @@ fn compute_next_d(
 #[allow(clippy::unwrap_used, clippy::too_many_arguments)]
 pub fn compute_mint_amount_for_deposit(
     amp_factor: &u64,
-    deposit_amount_a: Uint128,
-    deposit_amount_b: Uint128,
-    deposit_amount_c: Uint128,
-    swap_amount_a: Uint128,
-    swap_amount_b: Uint128,
-    swap_amount_c: Uint128,
+    deposits: &Vec<Coin>,
+    pool_assets: &Vec<Coin>,
     pool_token_supply: Uint128,
 ) -> Option<Uint128> {
     // Initial invariant
-    let d_0 = compute_d(amp_factor, swap_amount_a, swap_amount_b, swap_amount_c)?;
-    let new_balances = [
-        swap_amount_a.checked_add(deposit_amount_a).unwrap(),
-        swap_amount_b.checked_add(deposit_amount_b).unwrap(),
-        swap_amount_c.checked_add(deposit_amount_c).unwrap(),
-    ];
+    let d_0 = compute_d(amp_factor, deposits)?;
+
+    let new_balances: Vec<Coin> = pool_assets
+        .iter()
+        .enumerate()
+        .map(|(i, pool_asset)| {
+            let deposit_amount = deposits[i].amount;
+            let new_amount = pool_asset.amount.checked_add(deposit_amount).unwrap();
+            Coin {
+                denom: pool_asset.denom.clone(),
+                amount: new_amount,
+            }
+        })
+        .collect();
+
     // Invariant after change
-    let d_1 = compute_d(
-        amp_factor,
-        new_balances[0],
-        new_balances[1],
-        new_balances[2],
-    )?;
+    let d_1 = compute_d(amp_factor, &new_balances)?;
     if d_1 <= d_0 {
         None
     } else {

contracts/liquidity_hub/pool-manager/src/liquidity/commands.rs

@@ -271,27 +271,11 @@ pub fn provide_liquidity(
                     // Make sure at least MINIMUM_LIQUIDITY_AMOUNT is deposited to mitigate the risk of the first
                     // depositor preventing small liquidity providers from joining the pool
                     let min_lp_token_amount = MINIMUM_LIQUIDITY_AMOUNT * Uint128::from(3u8);
-                    let share = Uint128::try_from(
-                        compute_d(
-                            amp_factor,
-                            deposits[0].amount,
-                            deposits[1].amount,
-                            deposits[2].amount,
-                        )
-                        .unwrap(),
-                    )?
-                    .checked_sub(min_lp_token_amount)
-                    .map_err(|_| {
-                        ContractError::InvalidInitialLiquidityAmount(min_lp_token_amount)
-                    })?;
-
-                    // TODO: is this needed? I see it below after locking logic
-                    // messages.append(&mut mint_lp_token_msg(
-                    //     liquidity_token.clone(),
-                    //     env.contract.address.to_string(),
-                    //     env.contract.address.to_string(),
-                    //     min_lp_token_amount,
-                    // )?);
+                    let share = Uint128::try_from(compute_d(amp_factor, &deposits).unwrap())?
+                        .checked_sub(min_lp_token_amount)
+                        .map_err(|_| {
+                            ContractError::InvalidInitialLiquidityAmount(min_lp_token_amount)
+                        })?;
 
                     // share should be above zero after subtracting the min_lp_token_amount
                     if share.is_zero() {
@@ -304,12 +288,8 @@ pub fn provide_liquidity(
                 } else {
                     let amount = compute_mint_amount_for_deposit(
                         amp_factor,
-                        deposits[0].amount,
-                        deposits[1].amount,
-                        deposits[2].amount,
-                        pool_assets[0].amount,
-                        pool_assets[1].amount,
-                        pool_assets[2].amount,
+                        &deposits,
+                        &pool_assets,
                         total_share,
                     )
                     .unwrap();

contracts/liquidity_hub/pool-manager/src/tests/integration_tests.rs

@@ -2023,6 +2023,8 @@ mod swapping {
         );
     }
 
+    // TODO: make sure stableswap test works when pool has ONLY 2 assets
+
     #[test]
     fn basic_swapping_test_stable_swap() {
         let mut suite = TestingSuite::default_with_balances(vec![
@@ -3767,8 +3769,6 @@ mod provide_liquidity {
                 "uusd".to_string(),
             )],
             |result| {
-                // Find the key with 'offer_amount' and the key with 'return_amount'
-                // Ensure that the offer amount is 1000 and the return amount is greater than 0
                 let mut return_amount = String::new();
                 let mut offer_amount = String::new();
 
@@ -3793,6 +3793,198 @@ mod provide_liquidity {
             },
         );
     }
+
+    // This test is to ensure that the edge case of providing liquidity with 3 assets
+    #[test]
+    fn provide_liquidity_stable_swap_edge_case() {
+        let mut suite = TestingSuite::default_with_balances(vec![
+            coin(1_000_000_001u128, "uwhale".to_string()),
+            coin(1_000_000_000u128, "uluna".to_string()),
+            coin(1_000_000_001u128, "uusd".to_string()),
+        ]);
+        let creator = suite.creator();
+        let _other = suite.senders[1].clone();
+        let _unauthorized = suite.senders[2].clone();
+        // Asset infos with uwhale and uluna
+
+        let asset_infos = vec![
+            "uwhale".to_string(),
+            "uluna".to_string(),
+            "uusd".to_string(),
+        ];
+
+        // Protocol fee is 0.01% and swap fee is 0.02% and burn fee is 0%
+        #[cfg(not(feature = "osmosis"))]
+        let pool_fees = PoolFee {
+            protocol_fee: Fee {
+                share: Decimal::from_ratio(1u128, 1000u128),
+            },
+            swap_fee: Fee {
+                share: Decimal::from_ratio(1u128, 10_000_u128),
+            },
+            burn_fee: Fee {
+                share: Decimal::zero(),
+            },
+            extra_fees: vec![],
+        };
+
+        #[cfg(feature = "osmosis")]
+        let pool_fees = PoolFee {
+            protocol_fee: Fee {
+                share: Decimal::from_ratio(1u128, 1000u128),
+            },
+            swap_fee: Fee {
+                share: Decimal::from_ratio(1u128, 100_00u128),
+            },
+            burn_fee: Fee {
+                share: Decimal::zero(),
+            },
+            osmosis_fee: Fee {
+                share: Decimal::from_ratio(1u128, 1000u128),
+            },
+            extra_fees: vec![],
+        };
+
+        // Create a pool with 3 assets
+        suite.instantiate_default().create_pool(
+            creator.clone(),
+            asset_infos,
+            vec![6u8, 6u8, 6u8],
+            pool_fees,
+            PoolType::StableSwap { amp: 100 },
+            Some("whale-uluna-uusd".to_string()),
+            vec![coin(1000, "uusd")],
+            |result| {
+                result.unwrap();
+            },
+        );
+
+        // Adding liquidity with less than the minimum liquidity amount should fail
+        suite.provide_liquidity(
+            creator.clone(),
+            "whale-uluna-uusd".to_string(),
+            None,
+            None,
+            None,
+            None,
+            vec![
+                Coin {
+                    denom: "uwhale".to_string(),
+                    amount: Uint128::from(MINIMUM_LIQUIDITY_AMOUNT),
+                },
+                Coin {
+                    denom: "uluna".to_string(),
+                    amount: Uint128::from(MINIMUM_LIQUIDITY_AMOUNT),
+                },
+                Coin {
+                    denom: "uusd".to_string(),
+                    amount: Uint128::from(MINIMUM_LIQUIDITY_AMOUNT),
+                },
+            ],
+            |result| {
+                assert_eq!(
+                    result.unwrap_err().downcast_ref::<ContractError>(),
+                    Some(&ContractError::InvalidInitialLiquidityAmount(
+                        MINIMUM_LIQUIDITY_AMOUNT * Uint128::from(3u128)
+                    ))
+                );
+            },
+        );
+
+        // Lets try to add liquidity with the correct amount (1_000_000 of each asset)
+        suite.provide_liquidity(
+            creator.clone(),
+            "whale-uluna-uusd".to_string(),
+            None,
+            None,
+            None,
+            None,
+            vec![
+                Coin {
+                    denom: "uwhale".to_string(),
+                    amount: Uint128::from(1_000_000u128),
+                },
+                Coin {
+                    denom: "uluna".to_string(),
+                    amount: Uint128::from(1_000_000u128),
+                },
+                Coin {
+                    denom: "uusd".to_string(),
+                    amount: Uint128::from(1_000_000u128),
+                },
+            ],
+            |result| {
+                // Ensure we got 999000 in the response which is 1mil less the initial liquidity amount
+                for event in result.unwrap().events {
+                    for attribute in event.attributes {
+                        if attribute.key == "share" {
+                            assert_approx_eq!(
+                                attribute.value.parse::<u128>().unwrap(),
+                                1_000_000u128 * 3,
+                                "0.002"
+                            );
+                        }
+                    }
+                }
+            },
+        );
+
+        let simulated_return_amount = RefCell::new(Uint128::zero());
+        suite.query_simulation(
+            "whale-uluna-uusd".to_string(),
+            Coin {
+                denom: "uwhale".to_string(),
+                amount: Uint128::from(1_000u128),
+            },
+            "uluna".to_string(),
+            |result| {
+                *simulated_return_amount.borrow_mut() = result.unwrap().return_amount;
+            },
+        );
+
+        // Now lets try a swap
+        suite.swap(
+            creator.clone(),
+            "uluna".to_string(),
+            None,
+            None,
+            None,
+            "whale-uluna-uusd".to_string(),
+            vec![coin(1_000u128, "uwhale".to_string())],
+            |result| {
+                // Find the key with 'offer_amount' and the key with 'return_amount'
+                // Ensure that the offer amount is 1000 and the return amount is greater than 0
+                let mut return_amount = String::new();
+                let mut offer_amount = String::new();
+
+                for event in result.unwrap().events {
+                    if event.ty == "wasm" {
+                        for attribute in event.attributes {
+                            match attribute.key.as_str() {
+                                "return_amount" => return_amount = attribute.value,
+                                "offer_amount" => offer_amount = attribute.value,
+                                _ => {}
+                            }
+                        }
+                    }
+                }
+                // Because the Pool was created and 1_000_000 of each token has been provided as liquidity
+                // Assuming no fees we should expect a small swap of 1000 to result in not too much slippage
+                // Expect 1000 give or take 0.002 difference
+                // Once fees are added and being deducted properly only the "0.002" should be changed.
+                assert_approx_eq!(
+                    offer_amount.parse::<u128>().unwrap(),
+                    return_amount.parse::<u128>().unwrap(),
+                    "0.002"
+                );
+                assert_approx_eq!(
+                    simulated_return_amount.borrow().u128(),
+                    return_amount.parse::<u128>().unwrap(),
+                    "0.002"
+                );
+            },
+        );
+    }
 }
 
 mod multiple_pools {