add pool limit

ref-exchange/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ref-exchange"
-version = "1.9.3"
+version = "1.9.4"
 authors = ["Illia Polosukhin <[email protected]>"]
 edition = "2018"
 publish = false

ref-exchange/release_notes.md

@@ -1,5 +1,12 @@
 # Release Notes
 
+### Version 1.9.4
+```
+EcvZHyadAraTekV5isk2yucMFQvpSHG7uciCNuacDiau
+```
+1. add pool limit
+2. add execute_actions_in_va
+
 ### Version 1.9.3
 ```
 1PW1wtYsciZKsaRNqNMpY3P1W2wD42PjZVraL142VN4

ref-exchange/src/custom_keys.rs

@@ -4,3 +4,6 @@ pub const RATE_STORAGE_KEY: &str = "custom_rate_key";
 // Key for degen token info
 pub const DEGEN_STORAGE_KEY: &str = "custom_degen_key";
 pub const DEGEN_ORACLE_CONFIG_STORAGE_KEY: &str = "custom_degen_oracle_config_key";
+
+// Key for pool limit
+pub const POOL_LIMIT: &str = "pl";

ref-exchange/src/degen_swap/math.rs

@@ -134,7 +134,7 @@ impl DegenSwap {
     }
 
     /// *
-    fn degen_balances(&self, amounts: &Vec<Balance>) -> Vec<Balance> {
+    pub fn degen_balances(&self, amounts: &Vec<Balance>) -> Vec<Balance> {
         amounts.iter().zip(self.degens.iter()).map(|(&amount, &degen)| {
             self.mul_degen(amount, degen)
         }).collect()

ref-exchange/src/degen_swap/mod.rs

@@ -201,6 +201,12 @@ impl DegenSwapPool {
             .as_u128()
     }
 
+    pub fn get_tvl(&self) -> u128 {
+        self.get_invariant_with_degens(&self.get_degens())
+            .degen_balances(&self.c_amounts).iter()
+            .map(|v| v / 10u128.pow(TARGET_DECIMAL.into())).sum()
+    }
+
     /// caculate mint share and related fee for adding liquidity
     /// return (share, fee_part)
     fn calc_add_liquidity_with_degens(

ref-exchange/src/lib.rs

@@ -28,11 +28,12 @@ use crate::rated_swap::{RatedSwapPool, rate::{RateTrait, global_get_rate, global
 use crate::utils::{check_token_duplicates, pair_rated_price_to_vec_u8, TokenCache};
 pub use crate::custom_keys::*;
 pub use crate::views::{PoolInfo, ShadowRecordInfo, RatedPoolInfo, StablePoolInfo, ContractMetadata, RatedTokenInfo, DegenTokenInfo, AddLiquidityPrediction, RefStorageState};
-pub use crate::token_receiver::AddLiquidityInfo;
+pub use crate::token_receiver::{AddLiquidityInfo, VIRTUAL_ACC};
 pub use crate::shadow_actions::*;
 pub use crate::unit_lpt_cumulative_infos::*;
 pub use crate::oracle::*;
 pub use crate::degen_swap::*;
+pub use crate::pool_limit_info::*;
 
 mod account_deposit;
 mod action;
@@ -54,6 +55,7 @@ mod views;
 mod custom_keys;
 mod shadow_actions;
 mod unit_lpt_cumulative_infos;
+mod pool_limit_info;
 
 near_sdk::setup_alloc!();
 
@@ -69,6 +71,7 @@ pub(crate) enum StorageKey {
     Referral,
     ShadowRecord {account_id: AccountId},
     UnitShareCumulativeInfo,
+    PoolLimit,
 }
 
 #[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Eq, PartialEq, Clone)]
@@ -226,6 +229,59 @@ impl Contract {
         )))
     }
 
+    #[payable]
+    pub fn execute_actions_in_va(
+        &mut self,
+        use_tokens: HashMap<AccountId, U128>,
+        actions: Vec<Action>,
+        referral_id: Option<ValidAccountId>,
+    ) -> HashMap<AccountId, U128> {
+        self.assert_contract_running();
+        assert_ne!(actions.len(), 0, "{}", ERR72_AT_LEAST_ONE_SWAP);
+        let sender_id = env::predecessor_account_id();
+        let mut account = self.internal_unwrap_account(&sender_id);
+        // Validate that all tokens are whitelisted if no deposit (e.g. trade with access key).
+        if env::attached_deposit() == 0 {
+            for action in &actions {
+                for token in action.tokens() {
+                    assert!(
+                        account.get_balance(&token).is_some() 
+                            || self.is_whitelisted_token(&token),
+                        "{}",
+                        // [AUDIT_05]
+                        ERR27_DEPOSIT_NEEDED
+                    );
+                }
+            }
+        }
+
+        let mut virtual_account: Account = Account::new(&String::from(VIRTUAL_ACC));
+        let referral_info :Option<(AccountId, u32)> = referral_id
+            .as_ref().and_then(|rid| self.referrals.get(rid.as_ref()))
+            .map(|fee| (referral_id.unwrap().into(), fee));
+        for (use_token, use_amount) in use_tokens.iter() {
+            account.withdraw(use_token, use_amount.0);
+            virtual_account.deposit(use_token, use_amount.0);
+        }
+        let _ = self.internal_execute_actions(
+            &mut virtual_account,
+            &referral_info,
+            &actions,
+            ActionResult::None,
+        );
+        let mut result = HashMap::new();
+        for (token, amount) in virtual_account.tokens.to_vec() {
+            if amount > 0 {
+                account.deposit(&token, amount);
+                result.insert(token, amount.into());
+            }
+        }
+
+        virtual_account.tokens.clear();
+        self.internal_save_account(&sender_id, account);
+        result
+    }
+
     /// [AUDIT_03_reject(NOPE action is allowed by design)]
     /// [AUDIT_04]
     /// Executes generic set of actions.
@@ -377,6 +433,11 @@ impl Contract {
             AdminFees::new(self.admin_fee_bps),
             false
         );
+        if pool.kind() == "DEGEN_SWAP" {
+            if let Some(degen_pool_limit) = read_pool_limit_from_storage().get(&pool_id).map(|v| v.get_degen_pool_limit()) {
+                assert!(pool.get_tvl() <= degen_pool_limit.tvl_limit, "Exceed Max TVL");
+            }
+        }
         // [AUDITION_AMENDMENT] 2.3.7 Code Optimization (I)
         let mut deposits = self.internal_unwrap_account(&sender_id);
         let tokens = pool.tokens();

ref-exchange/src/owner.rs

@@ -446,6 +446,37 @@ impl Contract {
         }
     }
 
+    #[payable]
+    pub fn add_degen_pool_limit(&mut self, pool_id: u64, degen_pool_limit_info: DegenPoolLimitInfo) {
+        assert_one_yocto();
+        assert!(self.is_owner_or_guardians(), "{}", ERR100_NOT_ALLOWED);
+        assert!(self.get_pool(pool_id).pool_kind == "DEGEN_SWAP");
+        let mut pool_limit = read_pool_limit_from_storage();
+        assert!(pool_limit.get(&pool_id).is_none(), "degen pool limit already exist");
+        pool_limit.insert(&pool_id, &VPoolLimitInfo::DegenPoolLimit(degen_pool_limit_info.into()));
+        write_pool_limit_to_storage(pool_limit);
+    }
+
+    #[payable]
+    pub fn update_degen_pool_limit(&mut self, pool_id: u64, degen_pool_limit_info: DegenPoolLimitInfo) {
+        assert_one_yocto();
+        self.assert_owner();
+        assert!(self.get_pool(pool_id).pool_kind == "DEGEN_SWAP");
+        let mut pool_limit = read_pool_limit_from_storage();
+        assert!(pool_limit.get(&pool_id).is_some(), "degen pool limit not exist");
+        pool_limit.insert(&pool_id, &VPoolLimitInfo::DegenPoolLimit(degen_pool_limit_info.into()));
+        write_pool_limit_to_storage(pool_limit);
+    }
+
+    #[payable]
+    pub fn remove_pool_limit(&mut self, pool_id: u64) {
+        assert_one_yocto();
+        self.assert_owner();
+        let mut pool_limit = read_pool_limit_from_storage();
+        assert!(pool_limit.remove(&pool_id).is_some(), "Invalid pool_id");
+        write_pool_limit_to_storage(pool_limit);
+    }
+
     pub(crate) fn assert_owner(&self) {
         assert_eq!(
             env::predecessor_account_id(),

ref-exchange/src/pool.rs

@@ -165,6 +165,15 @@ impl Pool {
         }
     }
 
+    pub fn get_tvl(&self) -> u128 {
+        match self {
+            Pool::SimplePool(_) => unimplemented!(),
+            Pool::StableSwapPool(_) => unimplemented!(),
+            Pool::RatedSwapPool(_) => unimplemented!(),
+            Pool::DegenSwapPool(pool) => pool.get_tvl(),
+        }
+    }
+
     /// Swaps given number of token_in for token_out and returns received amount.
     pub fn swap(
         &mut self,

ref-exchange/src/pool_limit_info.rs

@@ -0,0 +1,62 @@
+use crate::*;
+use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
+use crate::utils::u128_dec_format;
+
+#[derive(BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
+#[serde(crate = "near_sdk::serde")]
+#[cfg_attr(not(target_arch = "wasm32"), derive(Debug))]
+pub struct DegenPoolLimitInfo {
+    #[serde(with = "u128_dec_format")]
+    pub tvl_limit: u128
+}
+
+#[derive(BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
+#[serde(crate = "near_sdk::serde")]
+#[cfg_attr(not(target_arch = "wasm32"), derive(Debug))]
+pub enum VDegenPoolLimitInfo {
+    Current(DegenPoolLimitInfo),
+}
+
+impl From<VDegenPoolLimitInfo> for DegenPoolLimitInfo {
+    fn from(v: VDegenPoolLimitInfo) -> Self {
+        match v {
+            VDegenPoolLimitInfo::Current(c) => c,
+        }
+    }
+}
+
+impl From<DegenPoolLimitInfo> for VDegenPoolLimitInfo {
+    fn from(c: DegenPoolLimitInfo) -> Self {
+        VDegenPoolLimitInfo::Current(c)
+    }
+}
+
+#[derive(BorshSerialize, BorshDeserialize, Serialize, Deserialize)]
+#[serde(crate = "near_sdk::serde")]
+#[cfg_attr(not(target_arch = "wasm32"), derive(Debug))]
+pub enum VPoolLimitInfo {
+    DegenPoolLimit(VDegenPoolLimitInfo)
+}
+
+impl VPoolLimitInfo {
+    pub fn get_degen_pool_limit(self) -> DegenPoolLimitInfo {
+        match self {
+            VPoolLimitInfo::DegenPoolLimit(l) => l.into(),
+        }
+    }
+}
+
+pub fn read_pool_limit_from_storage() -> UnorderedMap<u64, VPoolLimitInfo> {
+    if let Some(content) = env::storage_read(POOL_LIMIT.as_bytes()) {
+        UnorderedMap::try_from_slice(&content).expect("deserialize pool limit info failed.")
+    } else {
+        UnorderedMap::new(StorageKey::PoolLimit)
+    }
+}
+
+pub fn write_pool_limit_to_storage(pool_limit: UnorderedMap<u64, VPoolLimitInfo>) {
+    env::storage_write(
+        POOL_LIMIT.as_bytes(), 
+        &pool_limit.try_to_vec().unwrap(),
+    );
+}

ref-exchange/src/views.rs

@@ -844,4 +844,25 @@ impl Contract {
 
         Some((add_liquidity_predictions, token_cache.into()))
     }
+
+    pub fn get_degen_pool_tvl(&self, pool_id: u64) -> U128 {
+        self.pools.get(pool_id).expect(ERR85_NO_POOL).get_tvl().into()
+    }
+
+    pub fn get_pool_limit_by_pool_id(&self, pool_id: u64) -> Option<VPoolLimitInfo> {
+        read_pool_limit_from_storage().get(&pool_id)
+    }
+
+    pub fn get_pool_limit_paged(&self, from_index: Option<u64>, limit: Option<u64>) -> HashMap<u64, VPoolLimitInfo> {
+        let pool_limit = read_pool_limit_from_storage();
+        let keys = pool_limit.keys_as_vector();
+        let from_index = from_index.unwrap_or(0);
+        let limit = limit.unwrap_or(keys.len() as u64);
+        (from_index..std::cmp::min(keys.len() as u64, from_index + limit))
+            .map(|idx| {
+                let key = keys.get(idx).unwrap();
+                (key.clone(), pool_limit.get(&key).unwrap())
+            })
+            .collect()
+    }
 }