Merge pull request #1107 from ConsenSys/pruning_fix_2

Simplify bounded loops extension

mythril/laser/ethereum/strategy/extensions/bounded_loops.py

@@ -14,11 +14,11 @@ class JumpdestCountAnnotation(StateAnnotation):
     """State annotation that counts the number of jumps per destination."""
 
     def __init__(self) -> None:
-        self._jumpdest_count = {}  # type: Dict[int, int]
+        self._reached_count = {}  # type: Dict[int, int]
 
     def __copy__(self):
         result = JumpdestCountAnnotation()
-        result._jumpdest_count = copy(self._jumpdest_count)
+        result._reached_count = copy(self._reached_count)
         return result
 
 
@@ -31,11 +31,11 @@ class BoundedLoopsStrategy(BasicSearchStrategy):
         """"""
 
         self.super_strategy = super_strategy
-        self.jumpdest_limit = args[0][0]
+        self.bound = args[0][0]
 
         log.info(
             "Loaded search strategy extension: Loop bounds (limit = {})".format(
-                self.jumpdest_limit
+                self.bound
             )
         )
 
@@ -44,17 +44,14 @@ class BoundedLoopsStrategy(BasicSearchStrategy):
         )
 
     def get_strategic_global_state(self) -> GlobalState:
-        """
-        :return:
+        """ Returns the next state
+
+        :return: Global state
         """
 
         while True:
 
             state = self.super_strategy.get_strategic_global_state()
-            opcode = state.get_current_instruction()["opcode"]
-
-            if opcode != "JUMPI":
-                return state
 
             annotations = cast(
                 List[JumpdestCountAnnotation],
@@ -67,18 +64,15 @@ class BoundedLoopsStrategy(BasicSearchStrategy):
             else:
                 annotation = annotations[0]
 
-            try:
-                target = util.get_concrete_int(state.mstate.stack[-1])
-            except TypeError:
-                return state
+            address = state.get_current_instruction()["address"]
 
-            try:
-                annotation._jumpdest_count[target] += 1
-            except KeyError:
-                annotation._jumpdest_count[target] = 1
+            if address in annotation._reached_count:
+                annotation._reached_count[address] += 1
+            else:
+                annotation._reached_count[address] = 1
 
-            if annotation._jumpdest_count[target] > self.jumpdest_limit:
-                log.debug("JUMPDEST limit reached, skipping JUMPI")
+            if annotation._reached_count[address] > self.bound:
+                log.debug("Loop bound reached, skipping state")
                 continue
 
             return state