Adding storage used field

ref-exchange/src/account_deposit.rs

@@ -5,32 +5,45 @@ 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};
+use near_sdk::{assert_one_yocto, env, near_bindgen, AccountId, Balance, StorageUsage};
 
 use crate::utils::{ext_fungible_token, GAS_FOR_FT_TRANSFER};
 use crate::*;
 
-const MAX_ACCOUNT_LENGTH: u128 = 64;
-const MIN_ACCOUNT_DEPOSIT_LENGTH: u128 = MAX_ACCOUNT_LENGTH + 16 + 4;
+const INT_LENGTH: StorageUsage = 8;
+const SIZE_LENGTH: StorageUsage = 4;
+const MIN_ACCOUNT_DEPOSIT_LENGTH: StorageUsage = 64 + INT_LENGTH + SIZE_LENGTH;
 
 /// Account deposits information and storage cost.
-#[derive(BorshSerialize, BorshDeserialize, Default)]
+#[derive(BorshSerialize, BorshDeserialize)]
 #[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,
     /// Amounts of various tokens in this account.
     pub tokens: HashMap<AccountId, Balance>,
+    /// Storage used across this struct and all the pools.
+    pub storage_used: StorageUsage,
 }
 
 impl AccountDeposit {
+    pub fn new(account_id: &AccountId, amount: Balance) -> Self {
+        Self {
+            amount,
+            tokens: HashMap::default(),
+            // Empty AccountDeposit stores length of account id + amount + length of tokens.
+            storage_used: account_id.len() as StorageUsage + INT_LENGTH + SIZE_LENGTH,
+        }
+    }
+
     /// Adds amount to the balance of given token while checking that storage is covered.
     pub fn add(&mut self, token: AccountId, amount: Balance) {
         if let Some(x) = self.tokens.get_mut(&token) {
             *x = *x + amount;
         } else {
             self.tokens.insert(token.clone(), amount);
+            self.storage_used += token.len() as StorageUsage + INT_LENGTH;
             self.assert_storage_usage();
         }
     }
@@ -39,14 +52,13 @@ impl AccountDeposit {
     /// Panics if `amount` is bigger than the current balance.
     pub fn sub(&mut self, token: AccountId, amount: Balance) {
         let value = *self.tokens.get(&token).expect(ERR21_TOKEN_NOT_REG);
-        assert!(value >= amount, ERR22_NOT_ENOUGH_TOKENS);
+        assert!(value >= amount, "{}", ERR22_NOT_ENOUGH_TOKENS);
         self.tokens.insert(token, value - amount);
     }
 
     /// Returns amount of $NEAR necessary to cover storage used by this data structure.
     pub fn storage_usage(&self) -> Balance {
-        (MIN_ACCOUNT_DEPOSIT_LENGTH + self.tokens.len() as u128 * (MAX_ACCOUNT_LENGTH + 16))
-            * env::storage_byte_cost()
+        self.storage_used as Balance * env::storage_byte_cost()
     }
 
     /// Returns how much NEAR is available for storage.
@@ -58,26 +70,29 @@ impl AccountDeposit {
     pub fn assert_storage_usage(&self) {
         assert!(
             self.storage_usage() <= self.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()
+        MIN_ACCOUNT_DEPOSIT_LENGTH as Balance * env::storage_byte_cost()
     }
 
     /// Registers given token and set balance to 0.
     /// Fails if not enough amount to cover new storage usage.
     pub fn register(&mut self, token_id: &AccountId) {
         self.tokens.insert(token_id.clone(), 0);
+        self.storage_used += token_id.len() as StorageUsage + INT_LENGTH;
         self.assert_storage_usage();
     }
 
     /// Unregisters `token_id` from this account balance.
     /// Panics if the `token_id` balance is not 0.
     pub fn unregister(&mut self, token_id: &AccountId) {
         let amount = self.tokens.remove(token_id).unwrap_or_default();
+        self.storage_used -= token_id.len() as StorageUsage + INT_LENGTH;
         assert_eq!(amount, 0, "{}", ERR24_NON_ZERO_TOKEN_BALANCE);
     }
 }
@@ -136,9 +151,14 @@ impl Contract {
     /// 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);
+        let account_deposit =
+            if let Some(mut account_deposit) = self.deposited_amounts.get(&account_id) {
+                account_deposit.amount += amount;
+                account_deposit
+            } else {
+                AccountDeposit::new(account_id, amount)
+            };
+        self.deposited_amounts.insert(&account_id, &account_deposit);
     }
 
     /// Record deposit of some number of tokens to this contract.
@@ -153,6 +173,7 @@ impl Contract {
         assert!(
             self.whitelisted_tokens.contains(token_id)
                 || account_deposit.tokens.contains_key(token_id),
+            "{}",
             ERR12_TOKEN_NOT_WHITELISTED
         );
         account_deposit.add(token_id.clone(), amount);

ref-exchange/src/lib.rs

@@ -109,7 +109,10 @@ impl Contract {
         let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
         // Add amounts given to liquidity first. It will return the balanced amounts.
         pool.add_liquidity(&sender_id, &mut amounts);
-        let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
+        let mut deposits = self
+            .deposited_amounts
+            .get(&sender_id)
+            .expect(ERR10_ACC_NOT_REGISTERED);
         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() {
@@ -135,7 +138,10 @@ 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 deposits = self
+            .deposited_amounts
+            .get(&sender_id)
+            .expect(ERR10_ACC_NOT_REGISTERED);
         for i in 0..tokens.len() {
             deposits.add(tokens[i].clone(), amounts[i]);
         }
@@ -178,7 +184,10 @@ impl Contract {
         min_amount_out: U128,
         referral_id: Option<AccountId>,
     ) -> U128 {
-        let mut deposits = self.deposited_amounts.get(&sender_id).unwrap_or_default();
+        let mut deposits = self
+            .deposited_amounts
+            .get(&sender_id)
+            .expect(ERR10_ACC_NOT_REGISTERED);
         let amount_in: u128 = amount_in.into();
         deposits.sub(token_in.as_ref().clone(), amount_in);
         let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");

test-token/src/lib.rs

@@ -4,7 +4,7 @@ use near_contract_standards::fungible_token::metadata::{
 use near_contract_standards::fungible_token::FungibleToken;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::json_types::{ValidAccountId, U128};
-use near_sdk::{env, near_bindgen, AccountId, PanicOnDefault, PromiseOrValue};
+use near_sdk::{near_bindgen, AccountId, PanicOnDefault, PromiseOrValue};
 
 near_sdk::setup_alloc!();