From 856547fa97273d7cf8543022b5e32976cab296d9 Mon Sep 17 00:00:00 2001
From: Nikhil <nikhilparasaram@gmail.com>
Date: Mon, 21 Oct 2019 21:53:10 +0100
Subject: [PATCH] Fix tests

---
 mythril/laser/ethereum/keccak_function_manager.py |  5 +++--
 tests/laser/evm_testsuite/evm_test.py             | 11 ++++++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/mythril/laser/ethereum/keccak_function_manager.py b/mythril/laser/ethereum/keccak_function_manager.py
index d0f69d94..f79838a8 100644
--- a/mythril/laser/ethereum/keccak_function_manager.py
+++ b/mythril/laser/ethereum/keccak_function_manager.py
@@ -23,6 +23,7 @@ class KeccakFunctionManager:
         self.store_function = {}  # type: Dict[int, Tuple[Function, Function]]
         self.interval_hook_for_size = {}  # type: Dict[int, int]
         self._index_counter = TOTAL_PARTS - 34534
+        self.quick_inverse = {}  # type: Dict[BitVec, BitVec]  # This is for VMTests
 
     def find_keccak(self, data: BitVec) -> BitVec:
         """
@@ -72,8 +73,8 @@ class KeccakFunctionManager:
         func, inverse = self.get_function(length)
 
         condition = self._create_condition(func_input=data)
-        output = func(data)
-        return output, condition
+        self.quick_inverse[func(data)] = data
+        return func(data), condition
 
     def _create_condition(self, func_input: BitVec) -> Bool:
         """
diff --git a/tests/laser/evm_testsuite/evm_test.py b/tests/laser/evm_testsuite/evm_test.py
index 9fbb9dfa..9b9b913f 100644
--- a/tests/laser/evm_testsuite/evm_test.py
+++ b/tests/laser/evm_testsuite/evm_test.py
@@ -1,6 +1,7 @@
 from mythril.laser.ethereum.svm import LaserEVM
 from mythril.laser.ethereum.state.account import Account
 from mythril.laser.ethereum.state.world_state import WorldState
+from mythril.laser.ethereum.keccak_function_manager import keccak_function_manager
 from mythril.disassembler.disassembly import Disassembly
 from mythril.laser.ethereum.transaction.concolic import execute_message_call
 from mythril.laser.smt import Expression, BitVec, symbol_factory
@@ -176,7 +177,15 @@ def test_vmtest(
                 expected = int(value, 16)
                 actual = account.storage[symbol_factory.BitVecVal(int(index, 16), 256)]
                 if isinstance(actual, Expression):
-                    actual = actual.value
+                    if (
+                        actual.symbolic
+                        and actual in keccak_function_manager.quick_inverse
+                    ):
+                        actual = keccak_function_manager.find_keccak(
+                            keccak_function_manager.quick_inverse[actual]
+                        )
+                    else:
+                        actual = actual.value
                     actual = 1 if actual is True else 0 if actual is False else actual
                 else:
                     if type(actual) == bytes: