From 75177ce34cac589be26fb8089aac04feccdbae81 Mon Sep 17 00:00:00 2001
From: Dan Finlay <dan@danfinlay.com>
Date: Wed, 18 Oct 2017 15:08:34 -0700
Subject: [PATCH] Make account tracking more reactive

We were doing a lot of conditional observation & updating.
Pulled out a bunch of that for generic observer/syncers.
---
 app/scripts/controllers/computed-balances.js | 22 +++++++++++++-------
 app/scripts/lib/account-tracker.js           | 18 ++++++++++++++++
 app/scripts/metamask-controller.js           | 22 +-------------------
 3 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/app/scripts/controllers/computed-balances.js b/app/scripts/controllers/computed-balances.js
index 3479eae2b..009405d29 100644
--- a/app/scripts/controllers/computed-balances.js
+++ b/app/scripts/controllers/computed-balances.js
@@ -25,22 +25,28 @@ class ComputedbalancesController {
     }
   }
 
-  forgetAllBalances () {
-    this.balances = {}
-  }
-
   _initBalanceUpdating () {
     const store = this.accountTracker.store.getState()
-    this.addAnyAccountsFromStore(store)
-    this.accountTracker.store.subscribe(this.addAnyAccountsFromStore.bind(this))
+    this.syncAllAccountsFromStore(store)
+    this.accountTracker.store.subscribe(this.syncAllAccountsFromStore.bind(this))
   }
 
-  addAnyAccountsFromStore(store) {
-    const balances = store.accounts
+  syncAllAccountsFromStore(store) {
+    const upstream = Object.keys(store.accounts)
+    const balances = Object.keys(this.balances)
+    .map(address => this.balances[address])
 
+    // Follow new addresses
     for (let address in balances) {
       this.trackAddressIfNotAlready(address)
     }
+
+    // Unfollow old ones
+    balances.forEach(({ address }) => {
+      if (!upstream.includes(address)) {
+        delete this.balances[address]
+      }
+    })
   }
 
   trackAddressIfNotAlready (address) {
diff --git a/app/scripts/lib/account-tracker.js b/app/scripts/lib/account-tracker.js
index cdc21282d..13dea918f 100644
--- a/app/scripts/lib/account-tracker.js
+++ b/app/scripts/lib/account-tracker.js
@@ -38,6 +38,24 @@ class AccountTracker extends EventEmitter {
   // public
   //
 
+  syncWithAddresses (addresses) {
+    const accounts = this.store.getState().accounts
+    const locals = Object.keys(accounts)
+    .map(account => accounts[account.address])
+
+    addresses.forEach((upstream) => {
+      if (!locals.includes(upstream)) {
+        this.addAccount(upstream)
+      }
+    })
+
+    locals.forEach((local) => {
+      if (!addresses.includes(local)) {
+        this.removeAccount(local)
+      }
+    })
+  }
+
   addAccount (address) {
     const accounts = this.store.getState().accounts
     accounts[address] = {}
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index b312106dd..eae4478b5 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -123,13 +123,7 @@ module.exports = class MetamaskController extends EventEmitter {
         const address = addresses[0]
         this.preferencesController.setSelectedAddress(address)
       }
-    })
-    this.keyringController.on('newAccount', (address) => {
-      this.preferencesController.setSelectedAddress(address)
-      this.accountTracker.addAccount(address)
-    })
-    this.keyringController.on('removedAccount', (address) => {
-      this.accountTracker.removeAccount(address)
+      this.accountTracker.syncWithAddresses(addresses)
     })
 
     // address book controller
@@ -459,31 +453,17 @@ module.exports = class MetamaskController extends EventEmitter {
   //
 
   async createNewVaultAndKeychain (password, cb) {
-    this.forgetOldAccounts()
     const vault = await this.keyringController.createNewVaultAndKeychain(password)
     this.selectFirstIdentity(vault)
     return vault
   }
 
   async createNewVaultAndRestore (password, seed, cb) {
-    this.forgetOldAccounts()
     const vault = await this.keyringController.createNewVaultAndRestore(password, seed)
     this.selectFirstIdentity(vault)
     return vault
   }
 
-  forgetOldAccounts () {
-    const { accountTracker, balancesController } = this
-    let oldAccounts = []
-    try {
-      oldAccounts = Object.keys(accountTracker.store.getState().accounts)
-    } catch (e) {
-      log.warn('Could not load old accounts to forget', e)
-    }
-    oldAccounts.forEach(addr => accountTracker.removeAccount(addr))
-    balancesController.forgetAllBalances()
-  }
-
   selectFirstIdentity (vault) {
     const { identities } = vault
     const address = Object.keys(identities)[0]