fix storage retrieval

src/trace/traceAnalyser.js

@@ -10,8 +10,7 @@ TraceAnalyser.prototype.analyse = function (trace, tx, callback) {
   this.trace = trace
   this.traceCache.pushStoreChanges(0, tx.to)
   var context = {
-    currentStorageAddress: tx.to,
-    previousStorageAddress: tx.to,
+    storageContext: [tx.to],
     currentCallIndex: 0,
     lastCallIndex: 0
   }
@@ -80,16 +79,16 @@ TraceAnalyser.prototype.buildStorage = function (index, step, context) {
   if (traceHelper.newContextStorage(step) && !traceHelper.isCallToPrecompiledContract(index, this.trace)) {
     var calledAddress = traceHelper.resolveCalledAddress(index, this.trace)
     if (calledAddress) {
-      context.currentStorageAddress = calledAddress
+      context.storageContext.push(calledAddress)
     } else {
       console.log('unable to build storage changes. ' + index + ' does not match with a CALL. storage changes will be corrupted')
     }
-    this.traceCache.pushStoreChanges(index + 1, context.currentStorageAddress)
+    this.traceCache.pushStoreChanges(index + 1, context.storageContext[context.storageContext.length - 1])
   } else if (traceHelper.isSSTOREInstruction(step)) {
-    this.traceCache.pushStoreChanges(index + 1, context.currentStorageAddress, step.stack[step.stack.length - 1], step.stack[step.stack.length - 2])
+    this.traceCache.pushStoreChanges(index + 1, context.storageContext[context.storageContext.length - 1], step.stack[step.stack.length - 1], step.stack[step.stack.length - 2])
   } else if (traceHelper.isReturnInstruction(step)) {
-    context.currentStorageAddress = context.previousStorageAddress
-    this.traceCache.pushStoreChanges(index + 1, context.currentStorageAddress)
+    context.storageContext.pop()
+    this.traceCache.pushStoreChanges(index + 1, context.storageContext[context.storageContext.length - 1])
   }
   return context
 }

src/trace/traceManager.js

@@ -84,14 +84,8 @@ TraceManager.prototype.getStorageAt = function (stepIndex, tx, callback, address
     if (stoChange === null) return callback('no storage found', null)
     address = this.traceCache.sstore[stoChange].address
   }
-  var storage = {}
-  storage = this.traceCache.rebuildStorage(address, storage, stepIndex)
-  callback(null, storage)
-  /*
-  // TODO: use it if we need the full storage to be loaded
   var self = this
   if (this.traceRetriever.debugStorageAtAvailable()) {
-    var address = this.traceCache.sstore[stoChange].address
     this.traceRetriever.getStorage(tx, address, function (error, result) {
       if (error) {
         console.log(error)
@@ -104,7 +98,6 @@ TraceManager.prototype.getStorageAt = function (stepIndex, tx, callback, address
   } else {
     callback(null, this.trace[stoChange].storage)
   }
-  */
 }
 
 TraceManager.prototype.getAddresses = function (callback) {

src/trace/traceRetriever.js

@@ -3,7 +3,6 @@ var traceHelper = require('../helpers/traceHelper')
 var util = require('../helpers/global')
 
 function TraceRetriever () {
-  this.storages = {} // contains all intial storage (by addresses)
 }
 
 TraceRetriever.prototype.getTrace = function (txHash, callback) {
@@ -21,23 +20,15 @@ TraceRetriever.prototype.getTrace = function (txHash, callback) {
 TraceRetriever.prototype.getStorage = function (tx, address, callback) {
   if (traceHelper.isContractCreation(address)) {
     callback(null, {})
-  } else if (this.storages[address]) {
-    callback(null, this.storages[address])
   } else {
-    // we always return an empty storage ... storage changes will be displayed instead of the full contract storage
-    callback(null, {})
-    /*
-    var self = this
-    util.web3.debug.storageAt(tx.blockNumber.toString(), tx.transactionIndex, address, function (error, result) {
-      self.storages[address] = result
+    util.web3.debug.storageAt(null, tx.hash, address, function (error, result) {
       callback(error, result)
     })
-    */
   }
 }
 
 TraceRetriever.prototype.debugStorageAtAvailable = function () {
-  return false  // util.web3.version.node.toLowerCase().indexOf('geth') === -1 // storageAt not available if using geth
+  return true
 }
 
 module.exports = TraceRetriever

src/ui/SolidityState.js

@@ -29,7 +29,7 @@ SolidityState.prototype.init = function () {
       return
     }
 
-    self.traceManager.getStorageAt(index, null, function (error, storage) {
+    self.traceManager.getStorageAt(index, this.parent.tx, function (error, storage) {
       if (error) {
         self.basicPanel.update({ info: error })
       } else {

src/web3Provider/web3VmProvider.js

@@ -1,4 +1,5 @@
 var util = require('../helpers/util')
+var traceHelper = require('../helpers/traceHelper')
 var Web3 = require('web3')
 
 function web3VmProvider () {
@@ -19,6 +20,7 @@ function web3VmProvider () {
   this.debug.storageAt = function (blockNumber, txIndex, address, cb) { return self.storageAt(blockNumber, txIndex, address, cb) }
   this.providers = { 'HttpProvider': function (url) {} }
   this.currentProvider = {'host': 'vm provider'}
+  this.storageCache = {}
 }
 
 web3VmProvider.prototype.setVM = function (vm) {
@@ -62,6 +64,12 @@ web3VmProvider.prototype.txWillProcess = function (self, data) {
     tx.value = util.hexConvert(data.value)
   }
   self.txs[self.processingHash] = tx
+  self.storageCache[self.processingHash] = {}
+  if (tx.to) {
+    self.vm.stateManager.dumpStorage(tx.to, function (storage) {
+      self.storageCache[self.processingHash][tx.to] = storage
+    })
+  }
 }
 
 web3VmProvider.prototype.txProcessed = function (self, data) {
@@ -88,6 +96,14 @@ web3VmProvider.prototype.pushTrace = function (self, data) {
     gas: data.gasLeft.toString()
   }
   self.vmTraces[self.processingHash].structLogs.push(step)
+  if (traceHelper.newContextStorage(step)) {
+    if (!self.storageCache[self.processingHash][address]) {
+      var address = step.stack[step.stack.length - 2]
+      self.vm.stateManager.dumpStorage(address, function (storage) {
+        self.storageCache[self.processingHash][address] = storage
+      })
+    }
+  }
 }
 
 web3VmProvider.prototype.getCode = function (address, cb) {
@@ -111,7 +127,14 @@ web3VmProvider.prototype.traceTransaction = function (txHash, options, cb) {
   }
 }
 
-web3VmProvider.prototype.storageAt = function (blockNumber, txIndex, address, cb) { cb(null, {}) }
+web3VmProvider.prototype.storageAt = function (blockNumber, txIndex, address, cb) { // txIndex is the hash in the case of the VM
+  if (this.storageCache[txIndex] && this.storageCache[txIndex][address]) {
+    var storage = this.storageCache[txIndex][address]
+    return cb(null, JSON.parse(JSON.stringify(storage))) // copy creation...
+  } else {
+    cb('unable to retrieve storage ' + txIndex + ' ' + address)
+  }
+}
 
 web3VmProvider.prototype.getBlockNumber = function (cb) { cb(null, 'vm provider') }