Router v2 stable liquidity

amm/contracts/router_v2/lib.rs

@@ -50,7 +50,7 @@ pub mod router_v2 {
         /// Returns Pool for `pool_id` if it exists.
         /// Adds the Pool to the cache.
         #[inline]
-        fn get_pool(&mut self, pool_id: AccountId) -> Result<Pool, RouterV2Error> {
+        fn get_and_cache_pool(&mut self, pool_id: AccountId) -> Result<Pool, RouterV2Error> {
             match self.cached_pools.get(pool_id) {
                 Some(pool) => Ok(pool),
                 None => {
@@ -62,11 +62,12 @@ pub mod router_v2 {
         }
 
         /// Returns StablePool for `pool_id`.
+        /// Adds the StablePool to the cache.
         #[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),
+        fn get_and_cache_stable_pool(&mut self, pool_id: AccountId) -> Result<StablePool, RouterV2Error> {
+            match self.get_and_cache_pool(pool_id)? {
+                Pool::StablePool(pool) => Ok(pool),
+                Pool::Pair(_) => Err(RouterV2Error::InvalidPoolAddress),
             }
         }
 
@@ -76,7 +77,7 @@ pub mod router_v2 {
         /// exist in the pair Factory.
         /// Adds the Pair to the cache.
         #[inline]
-        fn get_pair(
+        fn get_and_cache_pair(
             &mut self,
             pool_id: Option<AccountId>,
             token_0: AccountId,
@@ -86,9 +87,9 @@ pub mod router_v2 {
                 Some(pool_id) => pool_id,
                 None => self.pair_factory_ref().create_pair(token_0, token_1)?,
             };
-            match self.get_pool(pool_id)? {
+            match self.get_and_cache_pool(pool_id)? {
                 Pool::Pair(pair) => Ok(pair),
-                Pool::StablePool(_) => Err(RouterV2Error::PairNotFound),
+                Pool::StablePool(_) => Err(RouterV2Error::InvalidPoolAddress),
             }
         }
 
@@ -101,15 +102,15 @@ pub mod router_v2 {
         ) -> Result<(), RouterV2Error> {
             let n_pools = path.len();
             for i in 0..n_pools - 1 {
-                self.get_pool(path[i].pool_id)?.swap(
+                self.get_and_cache_pool(path[i].pool_id)?.swap(
                     path[i].token_in,
                     path[i + 1].token_in,
                     amounts[i + 1],
                     path[i + 1].pool_id,
                 )?;
             }
             // If last pool in the path, transfer tokens to the `to` recipient.
-            self.get_pool(path[n_pools - 1].pool_id)?.swap(
+            self.get_and_cache_pool(path[n_pools - 1].pool_id)?.swap(
                 path[n_pools - 1].token_in,
                 token_out,
                 amounts[n_pools],
@@ -130,13 +131,13 @@ pub mod router_v2 {
             let n_pools = path.len();
             let mut amounts = vec![0; n_pools + 1];
             amounts[n_pools] = amount_out;
-            amounts[n_pools - 1] = self.get_pool(path[n_pools - 1].pool_id)?.get_amount_in(
+            amounts[n_pools - 1] = self.get_and_cache_pool(path[n_pools - 1].pool_id)?.get_amount_in(
                 path[n_pools - 1].token_in,
                 token_out,
                 amount_out,
             )?;
             for i in (0..n_pools - 1).rev() {
-                amounts[i] = self.get_pool(path[i].pool_id)?.get_amount_in(
+                amounts[i] = self.get_and_cache_pool(path[i].pool_id)?.get_amount_in(
                     path[i].token_in,
                     path[i + 1].token_in,
                     amounts[i + 1],
@@ -159,13 +160,13 @@ pub mod router_v2 {
             let mut amounts = vec![0; n_pools + 1];
             amounts[0] = amount_in;
             for i in 0..n_pools - 1 {
-                amounts[i + 1] = self.get_pool(path[i].pool_id)?.get_amount_out(
+                amounts[i + 1] = self.get_and_cache_pool(path[i].pool_id)?.get_amount_out(
                     path[i].token_in,
                     path[i + 1].token_in,
                     amounts[i],
                 )?;
             }
-            amounts[n_pools] = self.get_pool(path[n_pools - 1].pool_id)?.get_amount_out(
+            amounts[n_pools] = self.get_and_cache_pool(path[n_pools - 1].pool_id)?.get_amount_out(
                 path[n_pools - 1].token_in,
                 token_out,
                 amounts[n_pools - 1],
@@ -392,7 +393,7 @@ pub mod router_v2 {
             to: AccountId,
             deadline: u64,
         ) -> Result<(u128, u128, u128), RouterV2Error> {
-            let pair = self.get_pair(pair, token_0, token_1)?;
+            let pair = self.get_and_cache_pair(pair, token_0, token_1)?;
             pair.add_liquidity(
                 token_0,
                 token_1,
@@ -417,7 +418,7 @@ pub mod router_v2 {
             deadline: u64,
         ) -> Result<(u128, Balance, u128), RouterV2Error> {
             let wnative = self.wnative;
-            let pair = self.get_pair(pair, token, wnative)?;
+            let pair = self.get_and_cache_pair(pair, token, wnative)?;
             pair.add_liquidity_native(
                 token,
                 wnative,
@@ -441,7 +442,7 @@ pub mod router_v2 {
             to: AccountId,
             deadline: u64,
         ) -> Result<(u128, u128), RouterV2Error> {
-            let pair = self.get_pair(Some(pair), token_0, token_1)?;
+            let pair = self.get_and_cache_pair(Some(pair), token_0, token_1)?;
             pair.remove_liquidity(
                 token_0,
                 token_1,
@@ -465,7 +466,7 @@ pub mod router_v2 {
             deadline: u64,
         ) -> Result<(u128, Balance), RouterV2Error> {
             let wnative = self.wnative;
-            let pair = self.get_pair(Some(pair), token, wnative)?;
+            let pair = self.get_and_cache_pair(Some(pair), token, wnative)?;
             pair.remove_liquidity_native(
                 token,
                 wnative,
@@ -490,7 +491,7 @@ pub mod router_v2 {
             native: bool,
         ) -> Result<(u128, u128), RouterV2Error> {
             let wnative = if native { Some(self.wnative) } else { None };
-            self.get_stable_pool(pool)?.add_liquidity(
+            self.get_and_cache_stable_pool(pool)?.add_liquidity(
                 min_share_amount,
                 amounts,
                 to,
@@ -510,7 +511,7 @@ pub mod router_v2 {
             native: bool,
         ) -> Result<(u128, u128), RouterV2Error> {
             let wnative = if native { Some(self.wnative) } else { None };
-            self.get_stable_pool(pool)?.remove_liquidity(
+            self.get_and_cache_stable_pool(pool)?.remove_liquidity(
                 max_share_amount,
                 amounts,
                 to,
@@ -530,7 +531,7 @@ pub mod router_v2 {
             native: bool,
         ) -> Result<Vec<u128>, RouterV2Error> {
             let wnative = if native { Some(self.wnative) } else { None };
-            self.get_stable_pool(pool)?.remove_liquidity_by_share(
+            self.get_and_cache_stable_pool(pool)?.remove_liquidity_by_share(
                 share_amount,
                 min_amounts,
                 to,

amm/contracts/router_v2/stable_pool.rs

@@ -1,11 +1,12 @@
+use amm_helpers::ensure;
 use ink::{
     codegen::TraitCallBuilder,
     contract_ref,
     env::{account_id, caller, transferred_value, DefaultEnvironment as Env},
     prelude::vec::Vec,
     primitives::AccountId,
 };
-use traits::{RouterV2Error, StablePool as StablePoolTrait, StablePoolError};
+use traits::{RouterV2Error, StablePool as StablePoolTrait};
 
 use crate::utils::{
     check_timestamp, psp22_approve, psp22_transfer, psp22_transfer_from, transfer_native, withdraw,
@@ -68,10 +69,6 @@ impl StablePool {
         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) => {
@@ -115,10 +112,6 @@ impl StablePool {
         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>(),
@@ -133,11 +126,7 @@ impl StablePool {
                     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])?;
-                }
+                self.handle_token_transfers(&amounts, to, wnative_idx)?;
                 res
             }
             None => {
@@ -164,10 +153,6 @@ impl StablePool {
         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) => {
@@ -177,11 +162,7 @@ impl StablePool {
                     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])?;
-                }
+                self.handle_token_transfers(&amounts, to, wnative_idx)?;
                 Ok(amounts)
             }
             None => {
@@ -234,4 +215,18 @@ impl StablePool {
             .position(|&token| wnative == token)
             .ok_or(RouterV2Error::InvalidToken)
     }
+
+    fn handle_token_transfers(
+        &self,
+        amounts: &[u128],
+        to: AccountId,
+        wnative_idx: usize,
+    ) -> Result<(), RouterV2Error> {
+        withdraw(self.tokens[wnative_idx], 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(())
+    }
 }

amm/contracts/router_v2/utils.rs

@@ -43,15 +43,11 @@ pub fn psp22_approve(token: AccountId, spender: AccountId, value: u128) -> Resul
 #[inline]
 pub fn wrap(wnative: AccountId, value: Balance) -> Result<(), RouterV2Error> {
     let mut wnative_ref: contract_ref!(WrappedAZERO, Env) = wnative.into();
-    wnative_ref
+    Ok(wnative_ref
         .call_mut()
         .deposit()
         .transferred_value(value)
-        .try_invoke()
-        .map_err(|_| {
-            RouterV2Error::CrossContractCallFailed(String::from("Wrapped AZERO: deposit"))
-        })???;
-    Ok(())
+        .invoke()?)
 }
 
 #[inline]

amm/traits/router_v2.rs

@@ -1,6 +1,6 @@
 use crate::{Balance, FactoryError, MathError, PairError, StablePoolError};
 use ink::{
-    prelude::{string::String, vec::Vec},
+    prelude::vec::Vec,
     primitives::AccountId,
     LangError,
 };
@@ -289,12 +289,9 @@ pub enum RouterV2Error {
     MathError(MathError),
     StablePoolError(StablePoolError),
 
-    CrossContractCallFailed(String),
     Expired,
     InvalidPoolAddress,
     InvalidToken,
-    PairNotFound,
-    PoolNotFound,
     TransferError,
 
     ExcessiveInputAmount,