From 6f21b9279649b1fa1bc571f2ddfa40ac88b4ea33 Mon Sep 17 00:00:00 2001
From: Eric Ngo <eric.ngo@consensys.net>
Date: Fri, 4 Oct 2019 21:23:45 -0700
Subject: [PATCH] Added 'context' to loop annotations.

---
 mythril/laser/ethereum/state/machine_state.py | 24 ++++++++++++++++---
 .../strategy/extensions/bounded_loops.py      | 20 +++++++++++-----
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/mythril/laser/ethereum/state/machine_state.py b/mythril/laser/ethereum/state/machine_state.py
index e7bca3c0..a30c8afa 100644
--- a/mythril/laser/ethereum/state/machine_state.py
+++ b/mythril/laser/ethereum/state/machine_state.py
@@ -95,6 +95,7 @@ class MachineState:
         depth=0,
         max_gas_used=0,
         min_gas_used=0,
+        prev_pc=-1,
     ) -> None:
         """Constructor for machineState.
 
@@ -106,8 +107,9 @@ class MachineState:
         :param depth:
         :param max_gas_used:
         :param min_gas_used:
+        :param prev_pc:
         """
-        self.pc = pc
+        self._pc = pc
         self.stack = MachineStack(stack)
         self.memory = memory or Memory()
         self.gas_limit = gas_limit
@@ -115,6 +117,7 @@ class MachineState:
         self.max_gas_used = max_gas_used  # upper gas usage bound
         self.constraints = constraints or Constraints()
         self.depth = depth
+        self.prev_pc = prev_pc  # holds context of current pc
 
     def calculate_extension_size(self, start: int, size: int) -> int:
         """
@@ -210,11 +213,12 @@ class MachineState:
             gas_limit=self.gas_limit,
             max_gas_used=self.max_gas_used,
             min_gas_used=self.min_gas_used,
-            pc=self.pc,
+            pc=self._pc,
             stack=copy(self.stack),
             memory=copy(self.memory),
             constraints=copy(self.constraints),
             depth=self.depth,
+            prev_pc=self.prev_pc,
         )
 
     def __str__(self):
@@ -224,6 +228,19 @@ class MachineState:
         """
         return str(self.as_dict)
 
+    @property
+    def pc(self) -> int:
+        """
+
+        :return:
+        """
+        return self._pc
+
+    @pc.setter
+    def pc(self, value):
+        self.prev_pc = self._pc
+        self._pc = value
+
     @property
     def memory_size(self) -> int:
         """
@@ -239,11 +256,12 @@ class MachineState:
         :return:
         """
         return dict(
-            pc=self.pc,
+            pc=self._pc,
             stack=self.stack,
             memory=self.memory,
             memsize=self.memory_size,
             gas=self.gas_limit,
             max_gas_used=self.max_gas_used,
             min_gas_used=self.min_gas_used,
+            prev_pc=self.prev_pc,
         )
diff --git a/mythril/laser/ethereum/strategy/extensions/bounded_loops.py b/mythril/laser/ethereum/strategy/extensions/bounded_loops.py
index 151e60c1..94532a32 100644
--- a/mythril/laser/ethereum/strategy/extensions/bounded_loops.py
+++ b/mythril/laser/ethereum/strategy/extensions/bounded_loops.py
@@ -64,22 +64,30 @@ class BoundedLoopsStrategy(BasicSearchStrategy):
             else:
                 annotation = annotations[0]
 
-            address = state.get_current_instruction()["address"]
+            cur_instr = state.get_current_instruction()
 
-            if address in annotation._reached_count:
-                annotation._reached_count[address] += 1
+            if cur_instr["opcode"].upper() != "JUMPDEST":
+                return state
+
+            # create unique instruction identifier
+            key = "{}{}{}".format(
+                cur_instr["opcode"], cur_instr["address"], state.mstate.prev_pc
+            )
+
+            if key in annotation._reached_count:
+                annotation._reached_count[key] += 1
             else:
-                annotation._reached_count[address] = 1
+                annotation._reached_count[key] = 1
 
             # The creation transaction gets a higher loop bound to give it a better chance of success.
             # TODO: There's probably a nicer way to do this
 
             if isinstance(
                 state.current_transaction, ContractCreationTransaction
-            ) and annotation._reached_count[address] < max(8, self.bound):
+            ) and annotation._reached_count[key] < max(8, self.bound):
                 return state
 
-            elif annotation._reached_count[address] > self.bound:
+            elif annotation._reached_count[key] > self.bound:
                 log.debug("Loop bound reached, skipping state")
                 continue