Initial StateAnnotation Implementation (#762)

* Initial StateAnnotation Implementation

* add function return type

Co-Authored-By: JoranHonig <JoranHonig@users.noreply.github.com>

* add return type to annotate

Co-Authored-By: JoranHonig <JoranHonig@users.noreply.github.com>

* Add return type to annotations

Co-Authored-By: JoranHonig <JoranHonig@users.noreply.github.com>

* add retyrn type to annotate

Co-Authored-By: JoranHonig <JoranHonig@users.noreply.github.com>

* add type hints to worldstate __init__

Co-Authored-By: JoranHonig <JoranHonig@users.noreply.github.com>

* add return type to annotations

Co-Authored-By: JoranHonig <JoranHonig@users.noreply.github.com>

* style fixes

* Add list

mythril/laser/ethereum/state/annotation.py

@@ -0,0 +1,31 @@
+class StateAnnotation:
+    """
+    The StateAnnotation class is used to persist information over traces. This allows modules to reason about traces
+    without the need to traverse the state space themselves.
+    """
+
+    @property
+    def persist_to_world_state(self) -> bool:
+        """
+        If this function returns true then laser will also annotate the world state.
+
+        If you want annotations to persist through different user initiated message call transactions
+        then this should be enabled.
+
+        The default is set to False
+        """
+        return False
+
+
+class NoCopyAnnotation(StateAnnotation):
+    """
+    This class provides a base annotation class for annotations that shouldn't be copied on every new state.
+    Rather the same object should be propagated.
+    This is very useful if you are looking to analyze a property over multiple substates
+    """
+
+    def __copy__(self):
+        return self
+
+    def __deepcopy__(self, _):
+        return self

mythril/laser/ethereum/state/global_state.py

@@ -1,4 +1,4 @@
-from typing import Dict, Union
+from typing import Dict, Union, List
 
 from copy import copy, deepcopy
 from z3 import BitVec
@@ -6,6 +6,7 @@ from z3 import BitVec
 from mythril.laser.ethereum.cfg import Node
 from mythril.laser.ethereum.state.environment import Environment
 from mythril.laser.ethereum.state.machine_state import MachineState
+from mythril.laser.ethereum.state.annotation import StateAnnotation
 
 
 class GlobalState:
@@ -33,7 +34,7 @@ class GlobalState:
         self.transaction_stack = transaction_stack if transaction_stack else []
         self.op_code = ""
         self.last_return_data = last_return_data
-        self.annotations = annotations or []
+        self._annotations = annotations or []
 
     def __copy__(self) -> "GlobalState":
         world_state = copy(self.world_state)
@@ -47,7 +48,7 @@ class GlobalState:
             mstate,
             transaction_stack=transaction_stack,
             last_return_data=self.last_return_data,
-            annotations=self.annotations,
+            annotations=[copy(a) for a in self._annotations],
         )
 
     @property
@@ -78,3 +79,13 @@ class GlobalState:
     def new_bitvec(self, name: str, size=256) -> BitVec:
         transaction_id = self.current_transaction.id
         return BitVec("{}_{}".format(transaction_id, name), size)
+
+    def annotate(self, annotation: StateAnnotation) -> None:
+        self._annotations.append(annotation)
+
+        if annotation.persist_to_world_state:
+            self.world_state.annotate(annotation)
+
+    @property
+    def annotations(self) -> List[StateAnnotation]:
+        return self._annotations

mythril/laser/ethereum/state/world_state.py

@@ -1,7 +1,9 @@
 from copy import copy
 from random import randint
+from typing import List
 
 from mythril.laser.ethereum.state.account import Account
+from mythril.laser.ethereum.state.annotation import StateAnnotation
 
 
 class WorldState:
@@ -9,13 +11,16 @@ class WorldState:
     The WorldState class represents the world state as described in the yellow paper
     """
 
-    def __init__(self, transaction_sequence=None):
+    def __init__(
+        self, transaction_sequence=None, annotations: List[StateAnnotation] = None
+    ) -> None:
         """
         Constructor for the world state. Initializes the accounts record
         """
         self.accounts = {}
         self.node = None
         self.transaction_sequence = transaction_sequence or []
+        self._annotations = annotations or []
 
     def __getitem__(self, item: str) -> Account:
         """
@@ -26,7 +31,11 @@ class WorldState:
         return self.accounts[item]
 
     def __copy__(self) -> "WorldState":
-        new_world_state = WorldState(transaction_sequence=self.transaction_sequence[:])
+        new_annotations = [copy(a) for a in self._annotations]
+        new_world_state = WorldState(
+            transaction_sequence=self.transaction_sequence[:],
+            annotations=new_annotations,
+        )
         new_world_state.accounts = copy(self.accounts)
         new_world_state.node = self.node
         return new_world_state
@@ -67,6 +76,13 @@ class WorldState:
         new_account.storage = storage
         self._put_account(new_account)
 
+    def annotate(self, annotation: StateAnnotation) -> None:
+        self._annotations.append(annotation)
+
+    @property
+    def annotations(self) -> List[StateAnnotation]:
+        return self._annotations
+
     def _generate_new_address(self) -> str:
         """ Generates a new address for the global state"""
         while True: