From 4f00e1506bec27e24863ace4e503ded6f70bbd2b Mon Sep 17 00:00:00 2001
From: yann300 <yann.levreau@gmail.com>
Date: Tue, 13 Nov 2018 17:34:58 +0100
Subject: [PATCH] check RETURN and REVERT instruction

---
 remix-lib/src/helpers/traceHelper.js | 4 ++++
 remix-lib/src/trace/traceAnalyser.js | 5 ++++-
 remix-lib/src/trace/traceCache.js    | 5 +++++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/remix-lib/src/helpers/traceHelper.js b/remix-lib/src/helpers/traceHelper.js
index a4fb9695eb..7c19322681 100644
--- a/remix-lib/src/helpers/traceHelper.js
+++ b/remix-lib/src/helpers/traceHelper.js
@@ -33,6 +33,10 @@ module.exports = {
     return step.op === 'STOP'
   },
 
+  isRevertInstruction: function (step) {
+    return step.op === 'REVERT'
+  },
+
   isSSTOREInstruction: function (step) {
     return step.op === 'SSTORE'
   },
diff --git a/remix-lib/src/trace/traceAnalyser.js b/remix-lib/src/trace/traceAnalyser.js
index 66aefa2c51..33aa7c73f7 100644
--- a/remix-lib/src/trace/traceAnalyser.js
+++ b/remix-lib/src/trace/traceAnalyser.js
@@ -82,9 +82,12 @@ TraceAnalyser.prototype.buildStorage = function (index, step, context) {
     this.traceCache.pushStoreChanges(index + 1, context.storageContext[context.storageContext.length - 1])
   } else if (traceHelper.isSSTOREInstruction(step)) {
     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)) {
+  } else if (traceHelper.isReturnInstruction(step) || traceHelper.isStopInstruction(step)) {
     context.storageContext.pop()
     this.traceCache.pushStoreChanges(index + 1, context.storageContext[context.storageContext.length - 1])
+  } else if (traceHelper.isRevertInstruction(step)) {
+    context.storageContext.pop()
+    this.traceCache.resetStoreChanges()
   }
   return context
 }
diff --git a/remix-lib/src/trace/traceCache.js b/remix-lib/src/trace/traceCache.js
index 6448a12fd0..6cfe28d58f 100644
--- a/remix-lib/src/trace/traceCache.js
+++ b/remix-lib/src/trace/traceCache.js
@@ -82,6 +82,11 @@ TraceCache.prototype.pushContractCreation = function (token, code) {
   this.contractCreation[token] = code
 }
 
+TraceCache.prototype.resetStoreChanges = function (index, address, key, value) {
+  this.sstore = {}
+  this.storageChanges = []
+}
+
 TraceCache.prototype.pushStoreChanges = function (index, address, key, value) {
   this.sstore[index] = {
     'address': address,