feat(smart-contracts): add stableswap support for provide liquidity

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

@@ -18,15 +18,15 @@
 
 // todo isn't this for the 3pool? shouldn't it be 3
 // the number of assets in the pool
-const N_COINS: Uint256 = Uint256::from_u128(2);
+pub const N_COINS: u8 = 3;
 
 fn calculate_stableswap_d(
     offer_pool: Decimal256,
     ask_pool: Decimal256,
     amp: &u64,
     precision: u8,
 ) -> Result<Decimal256, ContractError> {
-    let n_coins = Decimal256::from_ratio(N_COINS, Uint256::from_u128(1));
+    let n_coins = Decimal256::from_ratio(Uint256::from(N_COINS), Uint256::from_u128(1));
 
     let sum_pools = offer_pool.checked_add(ask_pool)?;
     if sum_pools.is_zero() {
@@ -35,7 +35,10 @@
     }
 
     // ann = amp * n_coins
-    let ann = Decimal256::from_ratio(Uint256::from_u128((*amp).into()).checked_mul(N_COINS)?, 1u8);
+    let ann = Decimal256::from_ratio(
+        Uint256::from_u128((*amp).into()).checked_mul(Uint256::from(N_COINS))?,
+        1u8,
+    );
 
     // perform Newton-Raphson method
     let mut current_d = sum_pools;
@@ -99,7 +102,7 @@
     ask_precision: u8,
     direction: StableSwapDirection,
 ) -> Result<Uint128, ContractError> {
-    let ann = Uint256::from_u128((*amp).into()).checked_mul(N_COINS)?;
+    let ann = Uint256::from_u128((*amp).into()).checked_mul(Uint256::from(N_COINS))?;
 
     let d = calculate_stableswap_d(offer_pool, ask_pool, amp, ask_precision)?
         .to_uint256_with_precision(u32::from(ask_precision))?;
@@ -111,8 +114,8 @@
     .to_uint256_with_precision(u32::from(ask_precision))?;
 
     let c = d
-        .checked_multiply_ratio(d, pool_sum.checked_mul(N_COINS)?)?
-        .checked_multiply_ratio(d, ann.checked_mul(N_COINS)?)?;
+        .checked_multiply_ratio(d, pool_sum.checked_mul(Uint256::from(N_COINS))?)?
+        .checked_multiply_ratio(d, ann.checked_mul(Uint256::from(N_COINS))?)?;
 
     let b = pool_sum.checked_add(d.checked_div(ann)?)?;
 
@@ -474,8 +477,8 @@
 
 pub fn assert_slippage_tolerance(
     slippage_tolerance: &Option<Decimal>,
-    deposits: &[Uint128; 2],
-    pools: &[Coin; 2],
+    deposits: &[Coin],
+    pools: &[Coin],
     pool_type: PoolType,
     amount: Uint128,
     pool_token_supply: Uint128,
@@ -487,7 +490,7 @@
         }
 
         let one_minus_slippage_tolerance = Decimal256::one() - slippage_tolerance;
-        let deposits: [Uint256; 2] = [deposits[0].into(), deposits[1].into()];
+        let deposits: [Uint256; 2] = [deposits[0].amount.into(), deposits[1].amount.into()];
         let pools: [Uint256; 2] = [pools[0].amount.into(), pools[1].amount.into()];
 
         // Ensure each prices are not dropped as much as slippage tolerance rate
@@ -662,3 +665,120 @@
         ask_decimal,
     ))
 }
+
+// TODO: handle unwraps properly
+#[allow(clippy::unwrap_used)]
+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());
+
+    if sum_x == Uint128::zero() {
+        Some(Uint256::zero())
+    } else {
+        // do as below but for a generic number of assets
+        let amount_times_coins: Vec<Uint128> = deposits
+            .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 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, n_coins).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;
+                }
+            }
+        }
+        Some(d)
+    }
+}
+
+// TODO: handle unwraps properly
+#[allow(clippy::unwrap_used)]
+fn compute_next_d(
+    amp_factor: &u64,
+    d_init: Uint256,
+    d_prod: Uint256,
+    sum_x: Uint128,
+    n_coins: Uint128,
+) -> Option<Uint256> {
+    let ann = amp_factor.checked_mul(n_coins.u128() as u64)?;
+    let leverage = Uint256::from(sum_x).checked_mul(ann.into()).unwrap();
+    // d = (ann * sum_x + d_prod * n_coins) * d / ((ann - 1) * d + (n_coins + 1) * d_prod)
+    let numerator = d_init
+        .checked_mul(
+            d_prod
+                .checked_mul(n_coins.into())
+                .unwrap()
+                .checked_add(leverage)
+                .unwrap(),
+        )
+        .unwrap();
+    let denominator = d_init
+        .checked_mul(ann.checked_sub(1)?.into())
+        .unwrap()
+        .checked_add(
+            d_prod
+                .checked_mul((n_coins.checked_add(1u128.into()).unwrap()).into())
+                .unwrap(),
+        )
+        .unwrap();
+    Some(numerator.checked_div(denominator).unwrap())
+}
+
+/// Computes the amount of pool tokens to mint after a deposit.
+#[allow(clippy::unwrap_used, clippy::too_many_arguments)]
+pub fn compute_mint_amount_for_deposit(
+    amp_factor: &u64,
+    deposits: &Vec<Coin>,
+    pool_assets: &[Coin],
+    pool_token_supply: Uint128,
+) -> Option<Uint128> {
+    // Initial invariant
+    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)?;
+    if d_1 <= d_0 {
+        None
+    } else {
+        let amount = Uint256::from(pool_token_supply)
+            .checked_mul(d_1.checked_sub(d_0).unwrap())
+            .unwrap()
+            .checked_div(d_0)
+            .unwrap();
+        Some(Uint128::try_from(amount).unwrap())
+    }
+}

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

@@ -1,6 +1,6 @@
 use cosmwasm_std::{
     coin, coins, ensure, to_json_binary, wasm_execute, BankMsg, Coin, CosmosMsg, DepsMut, Env,
-    MessageInfo, Response, StdError, SubMsg,
+    MessageInfo, Response, SubMsg,
 };
 use cosmwasm_std::{Decimal, OverflowError, Uint128};
 
@@ -23,7 +23,7 @@
 // After writing create_pool I see this can get quite verbose so attempting to
 // break it down into smaller modules which house some things like swap, liquidity etc
 use crate::contract::SINGLE_SIDE_LIQUIDITY_PROVISION_REPLY_ID;
-use crate::helpers::aggregate_outgoing_fees;
+use crate::helpers::{aggregate_outgoing_fees, compute_d, compute_mint_amount_for_deposit};
 use crate::queries::query_simulation;
 use crate::state::{
     LiquidityProvisionData, SingleSideLiquidityProvisionBuffer,
@@ -251,22 +251,11 @@
                             .multiply_ratio(total_share, pool_assets[1].amount),
                     );
 
-                    let deposits_as: [Uint128; 2] = deposits
-                        .iter()
-                        .map(|coin| coin.amount)
-                        .collect::<Vec<_>>()
-                        .try_into()
-                        .map_err(|_| StdError::generic_err("Error converting vector to array"))?;
-                    let pools_as: [Coin; 2] = pool_assets
-                        .to_vec()
-                        .try_into()
-                        .map_err(|_| StdError::generic_err("Error converting vector to array"))?;
-
                     // assert slippage tolerance
                     helpers::assert_slippage_tolerance(
                         &slippage_tolerance,
-                        &deposits_as,
-                        &pools_as,
+                        &deposits,
+                        &pool_assets,
                         pool.pool_type.clone(),
                         amount,
                         total_share,
@@ -275,10 +264,43 @@
                     amount
                 }
             }
-            PoolType::StableSwap { amp: _ } => {
-                // TODO: Handle stableswap
+            PoolType::StableSwap { amp: amp_factor } => {
+                if total_share == Uint128::zero() {
+                    // 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).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() {
+                        return Err(ContractError::InvalidInitialLiquidityAmount(
+                            min_lp_token_amount,
+                        ));
+                    }
 
-                Uint128::one()
+                    share
+                } else {
+                    let amount = compute_mint_amount_for_deposit(
+                        amp_factor,
+                        &deposits,
+                        &pool_assets,
+                        total_share,
+                    )
+                    .unwrap();
+                    helpers::assert_slippage_tolerance(
+                        &slippage_tolerance,
+                        &deposits,
+                        &pool_assets,
+                        pool.pool_type.clone(),
+                        amount,
+                        total_share,
+                    )?;
+                    amount
+                }
             }
         };
 

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

@@ -95,3 +95,12 @@ pub const SWAP_ROUTES: Map<(&str, &str), SwapOperations> = Map::new("swap_routes
 
 pub const CONFIG: Item<Config> = Item::new("config");
 pub const POOL_COUNTER: Item<u64> = Item::new("pool_count");
+
+pub const STABLE_SWAP_PARAMS: Item<StableSwapParams> = Item::new("stable_swap_params");
+#[cw_serde]
+pub struct StableSwapParams {
+    pub initial_amp: u64,
+    pub future_amp: u64,
+    pub initial_amp_block: u64,
+    pub future_amp_block: u64,
+}