add migration

ref-exchange/src/account_deposit.rs

@@ -20,17 +20,55 @@ const U32_STORAGE: StorageUsage = 4;
 /// max length of account id * 8 (1 byte)
 const ACC_ID_STORAGE: StorageUsage = 64 * 8;
 // 64 = max account name length
+
+// ACC_ID: the Contract accounts map
+// + U128_STORAGE: near_amount storage
+// + U32_STORAGE: tokens HashMap length
+// + U64_STORAGE: storage_used
+// + 2: version
 pub const INIT_ACCOUNT_STORAGE: StorageUsage =
-    ACC_ID_STORAGE + U128_STORAGE + U32_STORAGE + 2 * U64_STORAGE;
+    ACC_ID_STORAGE + U128_STORAGE + U32_STORAGE + 2 * U64_STORAGE + 2;
 
 // 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();
 
+#[derive(BorshDeserialize, BorshSerialize)]
+pub enum AccountDeposit {
+    V2(AccountDepositV2),
+}
+
+impl From<AccountDeposit> for AccountDepositV2 {
+    fn from(account: AccountDeposit) -> Self {
+        match account {
+            AccountDeposit::V2(a) => {
+                if a.storage_used > 0 {
+                    return a;
+                }
+                // migrate from V1
+                a.storage_used = U64_STORAGE;
+                a
+            }
+        }
+    }
+}
+
+/// Account deposits information and storage cost.
+/// Legacy version
+#[derive(BorshSerialize, BorshDeserialize, Default)]
+#[cfg_attr(feature = "test", derive(Clone))]
+pub struct AccountDepositV1 {
+    /// 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>,
+}
+
 /// Account deposits information and storage cost.
 #[derive(BorshSerialize, BorshDeserialize, Default)]
 #[cfg_attr(feature = "test", derive(Clone))]
-pub struct AccountDeposit {
+pub struct AccountDepositV2 {
     /// NEAR sent to the exchange.
     /// Used for storage and trading.
     pub near_amount: Balance,
@@ -39,13 +77,19 @@ pub struct AccountDeposit {
     pub storage_used: StorageUsage,
 }
 
-impl AccountDeposit {
+impl From<AccountDepositV2> for AccountDeposit {
+    fn from(a: AccountDepositV2) -> Self {
+        AccountDeposit::V2(a)
+    }
+}
+
+impl AccountDepositV2 {
     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,
+            storage_used: U64_STORAGE,
         }
     }
 
@@ -137,16 +181,16 @@ impl Contract {
         let sender_id = env::predecessor_account_id();
         let mut acc = self.get_account(&sender_id);
         acc.near_amount += env::attached_deposit();
-        self.accounts.insert(&sender_id, &acc);
+        self.accounts.insert(&sender_id, &acc.into());
     }
 
     /// Registers given token in the user's account deposit.
     /// 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(&sender_id);
-        deposits.register(&token_ids);
-        self.accounts.insert(&sender_id, &deposits);
+        let mut acc = self.get_account(&sender_id);
+        acc.register(&token_ids);
+        self.accounts.insert(&sender_id, &acc.into());
     }
 
     /// Unregister given token from user's account deposit.
@@ -157,7 +201,7 @@ impl Contract {
         for token_id in token_ids {
             deposits.unregister(token_id.as_ref());
         }
-        self.accounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &deposits.into());
     }
 
     /// Withdraws given token from the deposits of given user.
@@ -175,7 +219,7 @@ impl Contract {
         if unregister == Some(true) {
             deposits.unregister(&token_id);
         }
-        self.accounts.insert(&sender_id, &deposits);
+        self.accounts.insert(&sender_id, &deposits.into());
         ext_fungible_token::ft_transfer(
             sender_id.clone().try_into().unwrap(),
             amount.into(),
@@ -214,7 +258,7 @@ impl Contract {
                 // This reverts the changes from withdraw function.
                 let mut deposits = self.get_account(&sender_id);
                 deposits.add(&token_id, amount.0);
-                self.accounts.insert(&sender_id, &deposits);
+                self.accounts.insert(&sender_id, &deposits.into());
             }
         };
     }
@@ -229,7 +273,7 @@ impl Contract {
             account_deposit.near_amount += amount;
             account_deposit
         } else {
-            AccountDeposit::new(account_id, amount)
+            AccountDepositV2::new(account_id, amount)
         };
         self.accounts.insert(&account_id, &acc);
     }
@@ -249,13 +293,29 @@ impl Contract {
             ERR12_TOKEN_NOT_WHITELISTED
         );
         acc.add(token_id, amount);
-        self.accounts.insert(sender_id, &acc);
+        self.accounts.insert(sender_id, &acc.into());
     }
 
     // Returns `from` AccountDeposit.
     #[inline]
-    pub(crate) fn get_account(&self, from: &AccountId) -> AccountDeposit {
-        self.accounts.get(from).expect(ERR10_ACC_NOT_REGISTERED)
+    pub(crate) fn get_account(&self, from: &AccountId) -> AccountDepositV2 {
+        self.accounts
+            .get(from)
+            .expect(ERR10_ACC_NOT_REGISTERED)
+            .into()
+    }
+
+    pub(crate) fn get_account_option(&self, from: &AccountId) -> Option<AccountDepositV2> {
+        // let key = ("d".to_owned() + from).into_bytes();
+        // let data = env::storage_read(&key);
+        // if data == None {
+        //     return None;
+        // }
+        // let Some(data) = data;
+        // AccountDepositV1::Dese
+        // borsh::de::
+
+        self.accounts.get(from).and_then(|a| a.into())
     }
 
     /// Returns current balance of given token for given user. If token_id == "" then returns NEAR (native)

ref-exchange/src/legacy.rs

@@ -4,7 +4,7 @@ use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::{LookupMap, Vector};
 use near_sdk::AccountId;
 
-use crate::account_deposit::AccountDeposit;
+use crate::account_deposit::AccountDepositV2;
 use crate::pool::Pool;
 
 /// Version before whitelisted tokens collection.
@@ -19,5 +19,5 @@ pub struct ContractV1 {
     /// List of all the pools.
     pub pools: Vector<Pool>,
     /// Balances of deposited tokens for each account.
-    pub deposited_amounts: LookupMap<AccountId, AccountDeposit>,
+    pub deposited_amounts: LookupMap<AccountId, AccountDepositV2>,
 }

ref-exchange/src/lib.rs

@@ -10,7 +10,7 @@ use near_sdk::{
     assert_one_yocto, env, log, near_bindgen, AccountId, PanicOnDefault, Promise, StorageUsage,
 };
 
-use crate::account_deposit::{AccountDeposit, INIT_ACCOUNT_STORAGE};
+use crate::account_deposit::{AccountDeposit, AccountDepositV1, INIT_ACCOUNT_STORAGE};
 pub use crate::action::*;
 use crate::errors::*;
 use crate::pool::Pool;
@@ -137,7 +137,7 @@ impl Contract {
         self.pools.replace(pool_id, &pool);
         // Can create a new shares record in a pool
         acc.update_storage(start_storage);
-        self.accounts.insert(&sender_id, &acc);
+        self.accounts.insert(&sender_id, &acc.into());
     }
 
     /// Remove liquidity from the pool and transfer it into account deposit.
@@ -163,7 +163,7 @@ impl Contract {
         }
         // Can remove shares record in a pool
         acc.update_storage(start_storage);
-        self.accounts.insert(&sender_id, &acc);
+        self.accounts.insert(&sender_id, &acc.into());
     }
 }
 
@@ -175,7 +175,7 @@ impl Contract {
         let from = env::predecessor_account_id();
         let mut acc = self.get_account(&from);
         acc.update_storage(tx_start_storage);
-        self.accounts.insert(&from, &acc);
+        self.accounts.insert(&from, &acc.into());
     }
 
     /// Adds given pool to the list and returns it's id.
@@ -214,14 +214,14 @@ impl Contract {
             referral_id,
         );
         acc.add(token_out, amount_out);
-        self.accounts.insert(&sender_id, &acc);
+        self.accounts.insert(&sender_id, &acc.into());
         self.pools.replace(pool_id, &pool);
         // NOTE: this can cause changes in the deposits which increases an account storage (eg,
         // if user doesn't have `token_out` in AccountDepoist, then a new record will be created).
         // This is not a problem, because we compute the `AccountDepoist` storage consumption
         // separaterly, hence we must do this update.
         acc.update_storage(start_storage);
-        self.accounts.insert(&sender_id, &acc);
+        self.accounts.insert(&sender_id, &acc.into());
         amount_out
     }
 }

ref-exchange/src/owner.rs

@@ -6,7 +6,7 @@ use crate::legacy::ContractV1;
 use crate::utils::{GAS_FOR_DEPLOY_CALL, GAS_FOR_UPGRADE_CALL};
 
 #[near_bindgen]
-impl Contract {
+impl ContractV1 {
     /// Change owner. Only can be called by owner.
     pub fn set_owner(&mut self, owner_id: ValidAccountId) {
         self.assert_owner();