Refactor to use StateAnnotation

mythril/laser/ethereum/strategy/custom.py

@@ -1,7 +1,9 @@
 from mythril.laser.ethereum.state.global_state import GlobalState
 from mythril.laser.ethereum.strategy.basic import BreadthFirstSearchStrategy
+from mythril.laser.ethereum.state.annotation import StateAnnotation
 from mythril.laser.ethereum import util
-from typing import Dict
+from typing import Dict, cast, List
+from copy import copy
 import logging
 
 
@@ -9,6 +11,18 @@ JUMPDEST_LIMIT = 4
 log = logging.getLogger(__name__)
 
 
+class JumpdestCountAnnotation(StateAnnotation):
+    """State annotation used when a path is chosen based on a predictable variable."""
+
+    def __init__(self) -> None:
+        self._jumpdest_count = {}  # type: Dict[object, dict]
+
+    def __copy__(self):
+        result = JumpdestCountAnnotation()
+        result.call_offsets = copy(self._jumpdest_count)
+        return result
+
+
 class BFSBoundedLoopsStrategy(BreadthFirstSearchStrategy):
     """Implements a breadth first search strategy that prunes loops.
     JUMPI instructions are skipped after a jump destination has been
@@ -31,21 +45,33 @@ class BFSBoundedLoopsStrategy(BreadthFirstSearchStrategy):
         if opcode != "JUMPI":
             return state
 
+        annotations = cast(
+            List[JumpdestCountAnnotation],
+            list(state.get_annotations(JumpdestCountAnnotation)),
+        )
+
+        if len(annotations) == 0:
+            annotation = JumpdestCountAnnotation()
+            state.annotate(annotation)
+        else:
+            annotation = annotations[0]
+
         transaction = state.current_transaction
         target = util.get_concrete_int(state.mstate.stack[-1])
 
         if transaction in self._jumpdest_count:
 
             try:
-                if self._jumpdest_count[transaction][target] == JUMPDEST_LIMIT:
+                if annotation._jumpdest_count[transaction][target] == JUMPDEST_LIMIT:
+                    log.info("JUMPDEST limit reached, skipping JUMPI")
                     return self.work_list.pop(0)
             except KeyError:
-                self._jumpdest_count[transaction][target] = 0
+                annotation._jumpdest_count[transaction][target] = 0
 
-            self._jumpdest_count[transaction][target] += 1
+            annotation._jumpdest_count[transaction][target] += 1
 
         else:
-            self._jumpdest_count[transaction] = {}
-            self._jumpdest_count[transaction][target] = 0
+            annotation._jumpdest_count[transaction] = {}
+            annotation._jumpdest_count[transaction][target] = 0
 
         return state