Merge branch 'master' into transactionControllerRefractorPt3

3 files changed, 130 insertions, 3 deletions

diff --git a/app/scripts/controllers/balance.js b/app/scripts/controllers/balance.js
new file mode 100644
index 000000000..964dff0df
--- /dev/null
+++ b/app/scripts/controllers/balance.js
@@ -0,0 +1,61 @@
+const ObservableStore = require('obs-store')
+const PendingBalanceCalculator = require('../lib/pending-balance-calculator')
+const BN = require('ethereumjs-util').BN
+
+class BalanceController {
+
+  constructor (opts = {}) {
+    const { address, accountTracker, txController, blockTracker } = opts
+    this.address = address
+    this.accountTracker = accountTracker
+    this.txController = txController
+    this.blockTracker = blockTracker
+
+    const initState = {
+      ethBalance: undefined,
+    }
+    this.store = new ObservableStore(initState)
+
+    this.balanceCalc = new PendingBalanceCalculator({
+      getBalance: () => this._getBalance(),
+      getPendingTransactions: this._getPendingTransactions.bind(this),
+    })
+
+    this._registerUpdates()
+  }
+
+  async updateBalance () {
+    const balance = await this.balanceCalc.getBalance()
+    this.store.updateState({
+      ethBalance: balance,
+    })
+  }
+
+  _registerUpdates () {
+    const update = this.updateBalance.bind(this)
+    this.txController.on('submitted', update)
+    this.txController.on('confirmed', update)
+    this.txController.on('failed', update)
+    this.accountTracker.store.subscribe(update)
+    this.blockTracker.on('block', update)
+  }
+
+  async _getBalance () {
+    const { accounts } = this.accountTracker.store.getState()
+    const entry = accounts[this.address]
+    const balance = entry.balance
+    return balance ? new BN(balance.substring(2), 16) : undefined
+  }
+
+  async _getPendingTransactions () {
+    const pending = this.txController.getFilteredTxList({
+      from: this.address,
+      status: 'submitted',
+      err: undefined,
+    })
+    return pending
+  }
+
+}
+
+module.exports = BalanceController
diff --git a/app/scripts/controllers/computed-balances.js b/app/scripts/controllers/computed-balances.js
new file mode 100644
index 000000000..2479e1b3a
--- /dev/null
+++ b/app/scripts/controllers/computed-balances.js
@@ -0,0 +1,66 @@
+const ObservableStore = require('obs-store')
+const extend = require('xtend')
+const BalanceController = require('./balance')
+
+class ComputedbalancesController {
+
+  constructor (opts = {}) {
+    const { accountTracker, txController, blockTracker } = opts
+    this.accountTracker = accountTracker
+    this.txController = txController
+    this.blockTracker = blockTracker
+
+    const initState = extend({
+      computedBalances: {},
+    }, opts.initState)
+    this.store = new ObservableStore(initState)
+    this.balances = {}
+
+    this._initBalanceUpdating()
+  }
+
+  updateAllBalances () {
+    for (let address in this.accountTracker.store.getState().accounts) {
+      this.balances[address].updateBalance()
+    }
+  }
+
+  _initBalanceUpdating () {
+    const store = this.accountTracker.store.getState()
+    this.addAnyAccountsFromStore(store)
+    this.accountTracker.store.subscribe(this.addAnyAccountsFromStore.bind(this))
+  }
+
+  addAnyAccountsFromStore(store) {
+    const balances = store.accounts
+
+    for (let address in balances) {
+      this.trackAddressIfNotAlready(address)
+    }
+  }
+
+  trackAddressIfNotAlready (address) {
+    const state = this.store.getState()
+    if (!(address in state.computedBalances)) {
+      this.trackAddress(address)
+    }
+  }
+
+  trackAddress (address) {
+    let updater = new BalanceController({
+      address,
+      accountTracker: this.accountTracker,
+      txController: this.txController,
+      blockTracker: this.blockTracker,
+    })
+    updater.store.subscribe((accountBalance) => {
+      let newState = this.store.getState()
+      newState.computedBalances[address] = accountBalance
+      this.store.updateState(newState)
+    })
+    this.balances[address] = updater
+    updater.updateBalance()
+  }
+}
+
+module.exports = ComputedbalancesController
diff --git a/app/scripts/controllers/transactions.js b/app/scripts/controllers/transactions.js
index 3ee1c22aa..de5fa5b20 100644
--- a/app/scripts/controllers/transactions.js
+++ b/app/scripts/controllers/transactions.js
@@ -32,7 +32,7 @@ module.exports = class TransactionController extends EventEmitter {
     this.provider = opts.provider
     this.blockTracker = opts.blockTracker
     this.signEthTx = opts.signTransaction
-    this.ethStore = opts.ethStore
+    this.accountTracker = opts.accountTracker
 
     this.memStore = new ObservableStore({})
     this.query = new EthQuery(this.provider)
@@ -61,7 +61,7 @@ module.exports = class TransactionController extends EventEmitter {
       nonceTracker: this.nonceTracker,
       retryLimit: 3500, // Retry 3500 blocks, or about 1 day.
       getBalance: (address) => {
-        const account = this.ethStore.getState().accounts[address]
+        const account = this.accountTracker.getState().accounts[address]
         if (!account) return
         return account.balance
       },
@@ -78,7 +78,7 @@ module.exports = class TransactionController extends EventEmitter {
     this.blockTracker.on('rawBlock', this.pendingTxTracker.checkForTxInBlock.bind(this.pendingTxTracker))
     // this is a little messy but until ethstore has been either
     // removed or redone this is to guard against the race condition
-    // where ethStore hasent been populated by the results yet
+    // where accountTracker hasent been populated by the results yet
     this.blockTracker.once('latest', () => {
       this.blockTracker.on('latest', this.pendingTxTracker.resubmitPendingTxs.bind(this.pendingTxTracker))
     })