Router v2 stable liquidity

amm/contracts/router_v2/lib.rs

@@ -12,7 +12,7 @@ pub struct CallerIsNotOwner;
 #[ink::contract]
 pub mod router_v2 {
     use crate::{
-        pool::{Pair, Pool},
+        pool::{Pair, Pool, StablePool},
         utils::*,
     };
     use amm_helpers::ensure;
@@ -60,6 +60,16 @@ pub mod router_v2 {
                 }
             }
         }
+
+        /// Returns StablePool for `pool_id`.
+        #[inline]
+        fn get_stable_pool(&mut self, pool_id: AccountId) -> Result<StablePool, RouterV2Error> {
+            match self.get_pool(pool_id) {
+                Ok(Pool::StablePool(pool)) => Ok(pool),
+                _ => Err(RouterV2Error::PoolNotFound),
+            }
+        }
+
         /// Returns Pair for `pool_id`.
         /// If `pool_id` is `None`, it creates a new Pair for
         /// `(token_0, token_1)` tokens if the Pair does not
@@ -284,9 +294,7 @@ pub mod router_v2 {
             self.swap(&amounts, &path, wnative, self.env().account_id())?;
             let native_out = amounts[amounts.len() - 1];
             withdraw(wnative, native_out)?;
-            self.env()
-                .transfer(to, native_out)
-                .map_err(|_| RouterV2Error::TransferError)?;
+            transfer_native(to, native_out)?;
             Ok(amounts)
         }
 
@@ -315,9 +323,7 @@ pub mod router_v2 {
             )?;
             self.swap(&amounts, &path, wnative, self.env().account_id())?;
             withdraw(wnative, native_out)?;
-            self.env()
-                .transfer(to, native_out)
-                .map_err(|_| RouterV2Error::TransferError)?;
+            transfer_native(to, native_out)?;
             Ok(amounts)
         }
 
@@ -470,6 +476,68 @@ pub mod router_v2 {
                 deadline,
             )
         }
+
+        // ----------- STABLE POOL LIQUIDITY METHODS ----------- //
+
+        #[ink(message, payable)]
+        fn add_stable_pool_liquidity(
+            &mut self,
+            pool: AccountId,
+            min_share_amount: u128,
+            amounts: Vec<u128>,
+            to: AccountId,
+            deadline: u64,
+            native: bool,
+        ) -> Result<(u128, u128), RouterV2Error> {
+            let wnative = if native { Some(self.wnative) } else { None };
+            self.get_stable_pool(pool)?.add_liquidity(
+                min_share_amount,
+                amounts,
+                to,
+                deadline,
+                wnative,
+            )
+        }
+
+        #[ink(message)]
+        fn remove_stable_pool_liquidity(
+            &mut self,
+            pool: AccountId,
+            max_share_amount: u128,
+            amounts: Vec<u128>,
+            to: AccountId,
+            deadline: u64,
+            native: bool,
+        ) -> Result<(u128, u128), RouterV2Error> {
+            let wnative = if native { Some(self.wnative) } else { None };
+            self.get_stable_pool(pool)?.remove_liquidity(
+                max_share_amount,
+                amounts,
+                to,
+                deadline,
+                wnative,
+            )
+        }
+
+        #[ink(message)]
+        fn remove_stable_pool_liquidity_by_share(
+            &mut self,
+            pool: AccountId,
+            share_amount: u128,
+            min_amounts: Vec<u128>,
+            to: AccountId,
+            deadline: u64,
+            native: bool,
+        ) -> Result<Vec<u128>, RouterV2Error> {
+            let wnative = if native { Some(self.wnative) } else { None };
+            self.get_stable_pool(pool)?.remove_liquidity_by_share(
+                share_amount,
+                min_amounts,
+                to,
+                deadline,
+                wnative,
+            )
+        }
     }
 
     #[cfg(test)]

amm/contracts/router_v2/pair.rs

@@ -3,7 +3,7 @@ use amm_helpers::{ensure, math::casted_mul};
 use ink::{
     codegen::TraitCallBuilder,
     contract_ref,
-    env::{account_id, caller, transfer, transferred_value, DefaultEnvironment as Env},
+    env::{account_id, caller, transferred_value, DefaultEnvironment as Env},
     primitives::AccountId,
 };
 use traits::{Balance, MathError, Pair as PairTrait, RouterV2Error};
@@ -63,6 +63,8 @@ impl Pair {
 
     /// Makes a cross-contract call to fetch the Pair reserves.
     /// Returns reserves `(reserve_0, reserve_1)` in order of `token_0` and `token_1`
+    /// NOTE: before calling this method ensure that `token_0` and `token_1` belong to
+    /// this `Pair` pool
     fn get_reserves(&self, token_0: &AccountId, token_1: &AccountId) -> (u128, u128) {
         let (reserve_0, reserve_1, _) = self.contract_ref().get_reserves();
         if token_0 < token_1 {
@@ -189,8 +191,7 @@ impl Pair {
         let liquidity = self.contract_ref().mint(to)?;
 
         if received_value > amount_native {
-            transfer::<Env>(caller, received_value - amount_native)
-                .map_err(|_| RouterV2Error::TransferError)?;
+            transfer_native(caller, received_value - amount_native)?;
         }
 
         Ok((amount_0, amount_native, liquidity))
@@ -244,7 +245,7 @@ impl Pair {
         )?;
         psp22_transfer(token, to, amount_token)?;
         withdraw(wnative, amount_native)?;
-        transfer::<Env>(to, amount_native).map_err(|_| RouterV2Error::TransferError)?;
+        transfer_native(to, amount_native)?;
         Ok((amount_token, amount_native))
     }
 

amm/contracts/router_v2/stable_pool.rs

@@ -1,9 +1,17 @@
 use amm_helpers::ensure;
 use ink::{
-    codegen::TraitCallBuilder, contract_ref, env::DefaultEnvironment as Env, prelude::vec::Vec,
+    codegen::TraitCallBuilder,
+    contract_ref,
+    env::{account_id, caller, transferred_value, DefaultEnvironment as Env},
+    prelude::vec::Vec,
     primitives::AccountId,
 };
-use traits::{RouterV2Error, StablePool as StablePoolTrait};
+use traits::{RouterV2Error, StablePool as StablePoolTrait, StablePoolError};
+
+use crate::utils::{
+    check_timestamp, psp22_approve, psp22_transfer, psp22_transfer_from, transfer_native, withdraw,
+    wrap,
+};
 
 #[derive(scale::Decode, scale::Encode)]
 #[cfg_attr(
@@ -31,6 +39,13 @@ impl StablePool {
             },
             Err(_) => return None,
         };
+        // set spending allowance of each token for the pool to `u128::MAX`
+        // required for adding liquidity
+        for &token in &tokens {
+            if psp22_approve(token, pool_id, u128::MAX).is_err() {
+                return None;
+            };
+        }
         Some(Self {
             id: pool_id,
             tokens,
@@ -41,6 +56,143 @@ impl StablePool {
         self.id.into()
     }
 
+    /// Adds liquidity to the pool.
+    ///
+    /// If `wnative` is specified, it attemps to wrap the transferred native token
+    /// and use it instead of transferring the wrapped version.
+    pub fn add_liquidity(
+        &self,
+        min_share_amount: u128,
+        amounts: Vec<u128>,
+        to: AccountId,
+        deadline: u64,
+        wnative: Option<AccountId>,
+    ) -> Result<(u128, u128), RouterV2Error> {
+        check_timestamp(deadline)?;
+        ensure!(
+            self.tokens.len() == amounts.len(),
+            RouterV2Error::StablePoolError(StablePoolError::IncorrectAmountsCount)
+        );
+        let native_received = transferred_value::<Env>();
+        let (wnative_idx, native_surplus) = match wnative {
+            Some(wnative) => {
+                let wnative_idx = self.wnative_idx(wnative)?;
+                let wnative_amount = amounts[wnative_idx];
+                ensure!(
+                    native_received >= wnative_amount,
+                    RouterV2Error::InsufficientTransferredAmount
+                );
+                wrap(wnative, wnative_amount)?;
+                (wnative_idx, native_received.saturating_sub(wnative_amount))
+            }
+            None => (self.tokens.len(), native_received),
+        };
+        if native_surplus > 0 {
+            transfer_native(caller::<Env>(), native_surplus)?;
+        }
+        for i in (0..self.tokens.len()).filter(|&idx| idx != wnative_idx) {
+            psp22_transfer_from(
+                self.tokens[i],
+                caller::<Env>(),
+                account_id::<Env>(),
+                amounts[i],
+            )?;
+        }
+        Ok(self
+            .contract_ref()
+            .add_liquidity(min_share_amount, amounts, to)?)
+    }
+
+    /// Withdraws liquidity from the pool by the specified amounts.
+    ///
+    /// If `wnative` is specified, it attemps to unwrap the wrapped native token
+    /// and withdraw it to the `to` account.
+    pub fn remove_liquidity(
+        &self,
+        max_share_amount: u128,
+        amounts: Vec<u128>,
+        to: AccountId,
+        deadline: u64,
+        wnative: Option<AccountId>,
+    ) -> Result<(u128, u128), RouterV2Error> {
+        check_timestamp(deadline)?;
+        ensure!(
+            self.tokens.len() == amounts.len(),
+            RouterV2Error::StablePoolError(StablePoolError::IncorrectAmountsCount)
+        );
+        psp22_transfer_from(
+            self.id,
+            caller::<Env>(),
+            account_id::<Env>(),
+            max_share_amount,
+        )?;
+        let (lp_burned, fee_part) = match wnative {
+            Some(wnative) => {
+                let wnative_idx = self.wnative_idx(wnative)?;
+                let res = self.contract_ref().remove_liquidity_by_amounts(
+                    max_share_amount,
+                    amounts.clone(),
+                    account_id::<Env>(),
+                )?;
+                withdraw(wnative, amounts[wnative_idx])?;
+                transfer_native(to, amounts[wnative_idx])?;
+                for i in (0..self.tokens.len()).filter(|&idx| idx != wnative_idx) {
+                    psp22_transfer(self.tokens[i], to, amounts[i])?;
+                }
+                res
+            }
+            None => {
+                self.contract_ref()
+                    .remove_liquidity_by_amounts(max_share_amount, amounts, to)?
+            }
+        };
+        if max_share_amount > lp_burned {
+            psp22_transfer(self.id, caller::<Env>(), max_share_amount - lp_burned)?;
+        }
+        Ok((lp_burned, fee_part))
+    }
+
+    /// Withdraws liquidity from the pool in balanced propotions.
+    ///
+    /// If `wnative` is specified, it attemps to unwrap the wrapped native token
+    /// and withdraw it to the `to` account.
+    pub fn remove_liquidity_by_share(
+        &self,
+        share_amount: u128,
+        min_amounts: Vec<u128>,
+        to: AccountId,
+        deadline: u64,
+        wnative: Option<AccountId>,
+    ) -> Result<Vec<u128>, RouterV2Error> {
+        check_timestamp(deadline)?;
+        ensure!(
+            self.tokens.len() == min_amounts.len(),
+            RouterV2Error::StablePoolError(StablePoolError::IncorrectAmountsCount)
+        );
+        psp22_transfer_from(self.id, caller::<Env>(), account_id::<Env>(), share_amount)?;
+        match wnative {
+            Some(wnative) => {
+                let wnative_idx = self.wnative_idx(wnative)?;
+                let amounts = self.contract_ref().remove_liquidity_by_shares(
+                    share_amount,
+                    min_amounts,
+                    account_id::<Env>(),
+                )?;
+                withdraw(wnative, amounts[wnative_idx])?;
+                transfer_native(to, amounts[wnative_idx])?;
+                for i in (0..self.tokens.len()).filter(|&idx| idx != wnative_idx) {
+                    psp22_transfer(self.tokens[i], to, amounts[i])?;
+                }
+                Ok(amounts)
+            }
+            None => {
+                Ok(self
+                    .contract_ref()
+                    .remove_liquidity_by_shares(share_amount, min_amounts, to)?)
+            }
+        }
+    }
+
     pub fn swap(
         &self,
         token_in: AccountId,
@@ -94,4 +246,11 @@ impl StablePool {
             Err(err) => Err(err.into()),
         }
     }
+
+    fn wnative_idx(&self, wnative: AccountId) -> Result<usize, RouterV2Error> {
+        self.tokens
+            .iter()
+            .position(|&token| wnative == token)
+            .ok_or(RouterV2Error::InvalidToken)
+    }
 }

amm/contracts/router_v2/utils.rs

@@ -2,7 +2,7 @@ use amm_helpers::ensure;
 use ink::{
     codegen::TraitCallBuilder,
     contract_ref,
-    env::{block_timestamp, DefaultEnvironment as Env},
+    env::{block_timestamp, transfer, DefaultEnvironment as Env},
     prelude::{string::String, vec::Vec},
     primitives::AccountId,
 };
@@ -34,6 +34,12 @@ pub fn psp22_transfer_from(
     token.transfer_from(from, to, value, Vec::new())
 }
 
+#[inline]
+pub fn psp22_approve(token: AccountId, spender: AccountId, value: u128) -> Result<(), PSP22Error> {
+    let mut token: contract_ref!(PSP22, Env) = token.into();
+    token.approve(spender, value)
+}
+
 #[inline]
 pub fn wrap(wnative: AccountId, value: Balance) -> Result<(), RouterV2Error> {
     let mut wnative_ref: contract_ref!(WrappedAZERO, Env) = wnative.into();
@@ -53,3 +59,8 @@ pub fn withdraw(wnative: AccountId, value: Balance) -> Result<(), RouterV2Error>
     let mut wnative_ref: contract_ref!(WrappedAZERO, Env) = wnative.into();
     Ok(wnative_ref.withdraw(value)?)
 }
+
+#[inline]
+pub fn transfer_native(to: AccountId, amount: u128) -> Result<(), RouterV2Error> {
+    transfer::<Env>(to, amount).map_err(|_| RouterV2Error::TransferError)
+}