1 files changed, 68 insertions, 14 deletions

diff --git a/app/scripts/controllers/preferences.js b/app/scripts/controllers/preferences.js
index f6250dc16..88b6fa998 100644
--- a/app/scripts/controllers/preferences.js
+++ b/app/scripts/controllers/preferences.js
@@ -13,6 +13,7 @@ class PreferencesController {
 	 * @property {array} store.frequentRpcList A list of custom rpcs to provide the user
    * @property {string} store.currentAccountTab Indicates the selected tab in the ui
    * @property {array} store.tokens The tokens the user wants display in their token lists
+   * @property {object} store.accountTokens The tokens stored per account and then per network type
    * @property {boolean} store.useBlockie The users preference for blockie identicons within the UI
    * @property {object} store.featureFlags A key-boolean map, where keys refer to features and booleans to whether the
    * user wishes to see that feature
@@ -24,6 +25,7 @@ class PreferencesController {
     const initState = extend({
       frequentRpcList: [],
       currentAccountTab: 'history',
+      accountTokens: {},
       tokens: [],
       useBlockie: false,
       featureFlags: {},
@@ -33,8 +35,10 @@ class PreferencesController {
     }, opts.initState)
 
     this.diagnostics = opts.diagnostics
-
+    this.network = opts.network
     this.store = new ObservableStore(initState)
+    this._defineTokens()
+    this._subscribeProviderType()
   }
 // PUBLIC METHODS
 
@@ -77,12 +81,17 @@ class PreferencesController {
    */
   setAddresses (addresses) {
     const oldIdentities = this.store.getState().identities
+    const accountTokens = this.store.getState().accountTokens
+
     const identities = addresses.reduce((ids, address, index) => {
       const oldId = oldIdentities[address] || {}
       ids[address] = {name: `Account ${index + 1}`, address, ...oldId}
       return ids
     }, {})
-    this.store.updateState({ identities })
+    for (const address in identities) {
+      if (!(address in accountTokens)) accountTokens[address] = {}
+    }
+    this.store.updateState({ identities, accountTokens })
   }
 
   /**
@@ -93,11 +102,13 @@ class PreferencesController {
    */
   removeAddress (address) {
     const identities = this.store.getState().identities
+    const accountTokens = this.store.getState().accountTokens
     if (!identities[address]) {
       throw new Error(`${address} can't be deleted cause it was not found`)
     }
     delete identities[address]
-    this.store.updateState({ identities })
+    delete accountTokens[address]
+    this.store.updateState({ identities, accountTokens })
 
     // If the selected account is no longer valid,
     // select an arbitrary other account:
@@ -117,14 +128,17 @@ class PreferencesController {
    */
   addAddresses (addresses) {
     const identities = this.store.getState().identities
+    const accountTokens = this.store.getState().accountTokens
     addresses.forEach((address) => {
       // skip if already exists
       if (identities[address]) return
       // add missing identity
       const identityCount = Object.keys(identities).length
+
+      if (!(address in accountTokens)) accountTokens[address] = {}
       identities[address] = { name: `Account ${identityCount + 1}`, address }
     })
-    this.store.updateState({ identities })
+    this.store.updateState({ identities, accountTokens })
   }
 
   /*
@@ -179,11 +193,16 @@ class PreferencesController {
    *
    */
   setSelectedAddress (_address) {
-    return new Promise((resolve, reject) => {
-      const address = normalizeAddress(_address)
-      this.store.updateState({ selectedAddress: address })
-      resolve()
-    })
+    const address = normalizeAddress(_address)
+    const accountTokens = this.store.getState().accountTokens
+    const providerType = this.network.providerStore.getState().type
+
+    if (!(address in accountTokens)) accountTokens[address] = {}
+    if (!(providerType in accountTokens[address])) accountTokens[address][providerType] = []
+    const tokens = accountTokens[address][providerType]
+    this.store.updateState({ selectedAddress: address, tokens })
+
+    return Promise.resolve(tokens)
   }
 
   /**
@@ -233,7 +252,11 @@ class PreferencesController {
       tokens.push(newEntry)
     }
 
-    this.store.updateState({ tokens })
+    const selectedAddress = this.store.getState().selectedAddress
+    const accountTokens = this.store.getState().accountTokens
+    const providerType = this.network.providerStore.getState().type
+    accountTokens[selectedAddress][providerType] = tokens
+    this.store.updateState({ accountTokens, tokens })
 
     return Promise.resolve(tokens)
   }
@@ -246,11 +269,13 @@ class PreferencesController {
    *
    */
   removeToken (rawAddress) {
-    const tokens = this.store.getState().tokens
+    const accountTokens = this.store.getState().accountTokens
+    const selectedAddress = this.store.getState().selectedAddress
+    const providerType = this.network.providerStore.getState().type
+    const updatedTokens = accountTokens[selectedAddress][providerType].filter(token => token.address !== rawAddress)
+    accountTokens[selectedAddress][providerType] = updatedTokens
+    this.store.updateState({ accountTokens, tokens: updatedTokens })
 
-    const updatedTokens = tokens.filter(token => token.address !== rawAddress)
-
-    this.store.updateState({ tokens: updatedTokens })
     return Promise.resolve(updatedTokens)
   }
 
@@ -376,6 +401,35 @@ class PreferencesController {
   //
   // PRIVATE METHODS
   //
+
+  /**
+   * Getter definition for the `tokens` property of store when controller is initialized
+   *
+   *
+   */
+  _defineTokens () {
+    const selectedAddress = this.store.getState().selectedAddress
+    const accountTokens = this.store.getState().accountTokens
+    const providerType = this.network.providerStore.getState().type
+    if (!(selectedAddress in accountTokens)) accountTokens[selectedAddress] = {}
+    if (!(providerType in accountTokens[selectedAddress])) return []
+    this.tokens = accountTokens[selectedAddress][providerType]
+  }
+
+  /**
+   * Subscription to network provider type
+   *
+   *
+   */
+  _subscribeProviderType () {
+    this.network.providerStore.subscribe(({ type }) => {
+      const selectedAddress = this.store.getState().selectedAddress
+      const accountTokens = this.store.getState().accountTokens
+      if (!(type in accountTokens[selectedAddress])) accountTokens[selectedAddress][type] = []
+      const tokens = accountTokens[selectedAddress][type]
+      this.store.updateState({ tokens })
+    })
+  }
 }
 
 module.exports = PreferencesController