near and storage tracing

ref-exchange/src/account_deposit.rs

@@ -5,29 +5,52 @@ use std::convert::TryInto;
 
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::json_types::{ValidAccountId, U128};
-use near_sdk::{assert_one_yocto, env, near_bindgen, AccountId, Balance, PromiseResult};
+use near_sdk::{
+    assert_one_yocto, env, near_bindgen, AccountId, Balance, PromiseResult, StorageUsage,
+};
 
 use crate::utils::{ext_fungible_token, ext_self, GAS_FOR_FT_TRANSFER};
 use crate::*;
 
-const MAX_ACCOUNT_LENGTH: u128 = 64;
-const MIN_ACCOUNT_DEPOSIT_LENGTH: u128 = MAX_ACCOUNT_LENGTH + 16 + 4;
+/// bytes length of u64 values
+const U64_STORAGE: StorageUsage = 8;
+/// bytes length of u32 values. Used in length operations
+const U32_STORAGE: StorageUsage = 4;
+// 64 = max account name length
+pub const INIT_ACCOUNT_STORAGE: StorageUsage = 64 + 2 * U64_STORAGE + U32_STORAGE;
+
+// NEAR native token. This is not a valid token ID. HACK: NEAR is a native token, we use the
+// empty string we use it to reference not existing near account.
+// pub const NEAR: AccountId = "".to_string();
 
 /// Account deposits information and storage cost.
 #[derive(BorshSerialize, BorshDeserialize, Default)]
 #[cfg_attr(feature = "test", derive(Clone))]
 pub struct AccountDeposit {
-    /// Native amount sent to the exchange.
-    /// Used for storage now, but in future can be used for trading as well.
-    pub amount: Balance,
+    /// NEAR sent to the exchange.
+    /// Used for storage and trading.
+    pub near_amount: Balance,
     /// Amounts of various tokens in this account.
     pub tokens: HashMap<AccountId, Balance>,
+    pub storage_used: StorageUsage,
 }
 
 impl AccountDeposit {
+    pub fn new(account_id: &AccountId, near_amount: Balance) -> Self {
+        Self {
+            near_amount,
+            tokens: HashMap::default(),
+            // Here we manually compute the initial storage size of account deposit.
+            storage_used: account_id.len() as StorageUsage + U64_STORAGE + U32_STORAGE,
+        }
+    }
+
     /// Adds amount to the balance of given token while checking that storage is covered.
     pub(crate) fn add(&mut self, token: &AccountId, amount: Balance) {
-        if let Some(x) = self.tokens.get_mut(token) {
+        if *token == "" {
+            // We use empty string to represent NEAR
+            self.near_amount += amount;
+        } else if let Some(x) = self.tokens.get_mut(token) {
             *x = *x + amount;
         } else {
             self.tokens.insert(token.clone(), amount);
@@ -38,38 +61,52 @@ impl AccountDeposit {
     /// Subtract from `token` balance.
     /// Panics if `amount` is bigger than the current balance.
     pub(crate) fn sub(&mut self, token: &AccountId, amount: Balance) {
+        if *token == "" {
+            // We use empty string to represent NEAR
+            self.near_amount -= amount;
+            self.assert_storage_usage();
+            return;
+        }
         let value = *self.tokens.get(token).expect(ERR21_TOKEN_NOT_REG);
         assert!(value >= amount, "{}", ERR22_NOT_ENOUGH_TOKENS);
         self.tokens.insert(token.clone(), value - amount);
     }
 
-    /// Returns amount of $NEAR necessary to cover storage used by this data structure.
+    /// Returns amount of $NEAR necessary to cover storage used by account referenced to this structure.
+    #[inline]
     pub fn storage_usage(&self) -> Balance {
-        (MIN_ACCOUNT_DEPOSIT_LENGTH + self.tokens.len() as u128 * (MAX_ACCOUNT_LENGTH + 16))
+        (if self.storage_used < INIT_ACCOUNT_STORAGE {
+            self.storage_used
+        } else {
+            INIT_ACCOUNT_STORAGE
+        }) as Balance
             * env::storage_byte_cost()
     }
 
     /// Returns how much NEAR is available for storage.
-    pub fn storage_available(&self) -> Balance {
-        self.amount - self.storage_usage()
+    #[inline]
+    pub(crate) fn storage_available(&self) -> Balance {
+        self.near_amount - self.storage_usage()
     }
 
     /// Asserts there is sufficient amount of $NEAR to cover storage usage.
     pub fn assert_storage_usage(&self) {
         assert!(
-            self.storage_usage() <= self.amount,
+            self.storage_usage() <= self.near_amount,
             "{}",
             ERR11_INSUFFICIENT_STORAGE
         );
     }
 
-    /// Returns minimal account deposit storage usage possible.
-    pub fn min_storage_usage() -> Balance {
-        MIN_ACCOUNT_DEPOSIT_LENGTH * env::storage_byte_cost()
+    /// Updates the account storage usage and sets.
+    pub(crate) fn update_storage(&mut self, tx_start_storage: StorageUsage) {
+        let s = env::storage_usage();
+        self.storage_used += s - tx_start_storage;
+        self.assert_storage_usage();
     }
 
-    /// Registers given token and set balance to 0.
-    /// Fails if not enough amount to cover new storage usage.
+    /// Registers given `token_id` and set balance to 0.
+    /// Fails if not enough NEAR is in deposit to cover new storage usage.
     pub(crate) fn register(&mut self, token_id: &AccountId) {
         self.tokens.insert(token_id.clone(), 0);
         self.assert_storage_usage();
@@ -89,22 +126,22 @@ impl Contract {
     /// Fails if not enough balance on this account to cover storage.
     pub fn register_tokens(&mut self, token_ids: Vec<ValidAccountId>) {
         let sender_id = env::predecessor_account_id();
-        let mut deposits = self.get_account_deposits(&sender_id);
+        let mut deposits = self.get_account(&sender_id);
         for token_id in token_ids {
             deposits.register(token_id.as_ref());
         }
-        self.deposited_amounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &deposits);
     }
 
     /// Unregister given token from user's account deposit.
     /// Panics if the balance of any given token is non 0.
     pub fn unregister_tokens(&mut self, token_ids: Vec<ValidAccountId>) {
         let sender_id = env::predecessor_account_id();
-        let mut deposits = self.get_account_deposits(&sender_id);
+        let mut deposits = self.get_account(&sender_id);
         for token_id in token_ids {
             deposits.unregister(token_id.as_ref());
         }
-        self.deposited_amounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &deposits);
     }
 
     /// Withdraws given token from the deposits of given user.
@@ -116,13 +153,13 @@ impl Contract {
         let token_id: AccountId = token_id.into();
         let amount: u128 = amount.into();
         let sender_id = env::predecessor_account_id();
-        let mut deposits = self.get_account_deposits(&sender_id);
+        let mut deposits = self.get_account(&sender_id);
         // Note: subtraction and deregistration will be reverted if the promise fails.
         deposits.sub(&token_id, amount);
         if unregister == Some(true) {
             deposits.unregister(&token_id);
         }
-        self.deposited_amounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &deposits);
         ext_fungible_token::ft_transfer(
             sender_id.clone().try_into().unwrap(),
             amount.into(),
@@ -159,9 +196,9 @@ impl Contract {
             PromiseResult::Successful(_) => {}
             PromiseResult::Failed => {
                 // This reverts the changes from withdraw function.
-                let mut deposits = self.get_account_deposits(&sender_id);
+                let mut deposits = self.get_account(&sender_id);
                 deposits.add(&token_id, amount.0);
-                self.deposited_amounts.insert(&sender_id, &deposits);
+                self.accounts.insert(&sender_id, &deposits);
             }
         };
     }
@@ -171,10 +208,14 @@ impl Contract {
     /// Registers account in deposited amounts with given amount of $NEAR.
     /// If account already exists, adds amount to it.
     /// This should be used when it's known that storage is prepaid.
-    pub(crate) fn internal_register_account(&mut self, account_id: &AccountId, amount: Balance) {
-        let mut deposit_amount = self.deposited_amounts.get(&account_id).unwrap_or_default();
-        deposit_amount.amount += amount;
-        self.deposited_amounts.insert(&account_id, &deposit_amount);
+    pub(crate) fn register_account(&mut self, account_id: &AccountId, amount: Balance) {
+        let acc = if let Some(mut account_deposit) = self.accounts.get(&account_id) {
+            account_deposit.near_amount += amount;
+            account_deposit
+        } else {
+            AccountDeposit::new(account_id, amount)
+        };
+        self.accounts.insert(&account_id, &acc);
     }
 
     /// Record deposit of some number of tokens to this contract.
@@ -185,32 +226,29 @@ impl Contract {
         token_id: &AccountId,
         amount: Balance,
     ) {
-        let mut account_deposit = self.get_account_deposits(sender_id);
+        let mut acc = self.get_account(sender_id);
         assert!(
-            self.whitelisted_tokens.contains(token_id)
-                || account_deposit.tokens.contains_key(token_id),
+            self.whitelisted_tokens.contains(token_id) || acc.tokens.contains_key(token_id),
             "{}",
             ERR12_TOKEN_NOT_WHITELISTED
         );
-        account_deposit.add(token_id, amount);
-        self.deposited_amounts.insert(sender_id, &account_deposit);
+        acc.add(token_id, amount);
+        self.accounts.insert(sender_id, &acc);
     }
 
     // Returns `from` AccountDeposit.
     #[inline]
-    pub(crate) fn get_account_deposits(&self, from: &AccountId) -> AccountDeposit {
-        self.deposited_amounts
-            .get(from)
-            .expect(ERR10_ACC_NOT_REGISTERED)
+    pub(crate) fn get_account(&self, from: &AccountId) -> AccountDeposit {
+        self.accounts.get(from).expect(ERR10_ACC_NOT_REGISTERED)
     }
 
     /// Returns current balance of given token for given user. If there is nothing recorded, returns 0.
-    pub(crate) fn internal_get_deposit(
+    pub(crate) fn get_deposit_balance(
         &self,
         sender_id: &AccountId,
         token_id: &AccountId,
     ) -> Balance {
-        self.deposited_amounts
+        self.accounts
             .get(sender_id)
             .and_then(|d| d.tokens.get(token_id).cloned())
             .unwrap_or_default()

ref-exchange/src/errors.rs

@@ -5,7 +5,7 @@ pub const ERR11_INSUFFICIENT_STORAGE: &str = "E11: insufficient $NEAR storage de
 pub const ERR12_TOKEN_NOT_WHITELISTED: &str = "E12: token not whitelisted";
 
 // Account Deposits //
-
+// TODO: change type to avoid assert!(..., "{}", ERR...). Maybe using raw strings? r#""#
 pub const ERR21_TOKEN_NOT_REG: &str = "E21: token not registered";
 pub const ERR22_NOT_ENOUGH_TOKENS: &str = "E22: not enough tokens in deposit";
 // pub const ERR23_NOT_ENOUGH_NEAR: &str = "E23: not enough NEAR in deposit";

ref-exchange/src/lib.rs

@@ -8,7 +8,7 @@ use near_sdk::collections::{LookupMap, UnorderedSet, Vector};
 use near_sdk::json_types::{ValidAccountId, U128};
 use near_sdk::{assert_one_yocto, env, log, near_bindgen, AccountId, PanicOnDefault, Promise};
 
-use crate::account_deposit::AccountDeposit;
+use crate::account_deposit::{AccountDeposit, INIT_ACCOUNT_STORAGE};
 pub use crate::action::*;
 use crate::errors::*;
 use crate::pool::Pool;
@@ -42,7 +42,7 @@ pub struct Contract {
     /// List of all the pools.
     pools: Vector<Pool>,
     /// Balances of deposited tokens for each account.
-    deposited_amounts: LookupMap<AccountId, AccountDeposit>,
+    accounts: LookupMap<AccountId, AccountDeposit>,
     /// Set of whitelisted tokens by "owner".
     whitelisted_tokens: UnorderedSet<AccountId>,
 }
@@ -56,15 +56,16 @@ impl Contract {
             exchange_fee,
             referral_fee,
             pools: Vector::new(b"p".to_vec()),
-            deposited_amounts: LookupMap::new(b"d".to_vec()),
+            accounts: LookupMap::new(b"d".to_vec()),
             whitelisted_tokens: UnorderedSet::new(b"w".to_vec()),
         }
     }
 
     /// Adds new "Simple Pool" with given tokens and given fee.
-    /// Attached NEAR should be enough to cover the added storage.
+    /// Deposited NEAR must be enough to cover the added storage.
     #[payable]
     pub fn add_simple_pool(&mut self, tokens: Vec<ValidAccountId>, fee: u32) -> u64 {
+        assert_one_yocto();
         check_token_duplicates(&tokens);
         self.internal_add_pool(Pool::SimplePool(SimplePool::new(
             self.pools.len() as u32,
@@ -121,13 +122,13 @@ impl Contract {
                 assert!(amount >= &min_amount.0, "ERR_MIN_AMOUNT");
             }
         }
-        let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
+        let mut acc = self.get_account(&sender_id);
         let tokens = pool.tokens();
         // Subtract updated amounts from deposits. This will fail if there is not enough funds for any of the tokens.
         for i in 0..tokens.len() {
-            deposits.sub(&tokens[i], amounts[i]);
+            acc.sub(&tokens[i], amounts[i]);
         }
-        self.deposited_amounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &acc);
         self.pools.replace(pool_id, &pool);
     }
 
@@ -147,11 +148,11 @@ impl Contract {
         );
         self.pools.replace(pool_id, &pool);
         let tokens = pool.tokens();
-        let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
+        let mut acc = self.get_account(&&sender_id);
         for i in 0..tokens.len() {
-            deposits.add(&tokens[i], amounts[i]);
+            acc.add(&tokens[i], amounts[i]);
         }
-        self.deposited_amounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &acc);
     }
 }
 
@@ -161,20 +162,12 @@ impl Contract {
     /// If there is not enough attached balance to cover storage, fails.
     /// If too much attached - refunds it back.
     fn internal_add_pool(&mut self, pool: Pool) -> u64 {
-        let prev_storage = env::storage_usage();
+        let start_storage = env::storage_usage();
         let id = self.pools.len() as u64;
         self.pools.push(&pool);
 
-        // Check how much storage cost and refund the left over back.
-        let storage_cost = (env::storage_usage() - prev_storage) as u128 * env::storage_byte_cost();
-        assert!(
-            storage_cost <= env::attached_deposit(),
-            "ERR_STORAGE_DEPOSIT"
-        );
-        let refund = env::attached_deposit() - storage_cost;
-        if refund > 0 {
-            Promise::new(env::predecessor_account_id()).transfer(refund);
-        }
+        // TODO: update deposit handling
+
         id
     }
 
@@ -190,8 +183,9 @@ impl Contract {
         min_amount_out: u128,
         referral_id: &Option<AccountId>,
     ) -> u128 {
-        let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
-        deposits.sub(token_in, amount_in);
+        let start_storage = env::storage_usage();
+        let mut acc = self.accounts.get(&sender_id).unwrap_or_default();
+        acc.sub(token_in, amount_in);
         let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
         let amount_out = pool.swap(
             token_in,
@@ -201,9 +195,16 @@ impl Contract {
             &self.owner_id,
             referral_id,
         );
-        deposits.add(token_out, amount_out);
-        self.deposited_amounts.insert(&sender_id, &deposits);
+        acc.add(token_out, amount_out);
+        self.accounts.insert(&sender_id, &acc);
         self.pools.replace(pool_id, &pool);
+        // TODO:
+        // we need to insert 2 times: once in line 194 (prev) and second time here.
+        // This is because we won't trace properly new storage until we insert the record into the storage tree.
+        // Alternative would be to refactor the AccountDeposit and move ynear to another, top-level map.
+        // Better solution: compute the storage consumed by AccountDeposit on the fly.
+        acc.update_storage(start_storage);
+
         amount_out
     }
 }

ref-exchange/src/owner.rs

@@ -64,7 +64,7 @@ impl Contract {
             exchange_fee: contract_v1.exchange_fee,
             referral_fee: contract_v1.referral_fee,
             pools: contract_v1.pools,
-            deposited_amounts: contract_v1.deposited_amounts,
+            accounts: contract_v1.deposited_amounts,
             whitelisted_tokens: UnorderedSet::new(b"w".to_vec()),
         }
     }