Refactor computation children to BaseComputation class

- Define computation children in the ``BaseComputation`` class and narrow the type as needed for subclasses - e.g. for ``MessageComputation`` narrow the children to be ``MessageComputationAPI``. - Some minor cleanup along the way.
ethereum · Apr 19, 2023 · c640856 · c640856
1 parent 0c4c4da
commit c640856
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 16 deletions.
diff --git a/eth/abc.py b/eth/abc.py
@@ -1862,7 +1862,7 @@ def base_fee_per_gas(self) -> Optional[int]:
 
 
 class ComputationAPI(
-    ContextManager['ComputationAPI'],
+    ContextManager["ComputationAPI"],
     StackManipulationAPI,
 ):
     """
@@ -2111,12 +2111,11 @@ class MessageComputationAPI(
 
     msg: MessageAPI
     transaction_context: TransactionContextAPI
-    children: List["MessageComputationAPI"]
 
     @abstractmethod
     def __init__(
         self,
-        state: 'StateAPI',
+        state: "StateAPI",
         message: MessageAPI,
         transaction_context: TransactionContextAPI,
     ) -> None:
@@ -2154,7 +2153,7 @@ def prepare_child_message(self,
     def apply_child_message_computation(
         self,
         child_msg: MessageAPI,
-    ) -> 'MessageComputationAPI':
+    ) -> "MessageComputationAPI":
         """
         Apply the vm message ``child_msg`` as a child message computation.
         """
@@ -2164,7 +2163,7 @@ def apply_child_message_computation(
     def generate_child_message_computation(
         self,
         child_msg: MessageAPI,
-    ) -> 'MessageComputationAPI':
+    ) -> "MessageComputationAPI":
         """
         Generate a child message computation from the given ``child_msg``.
         """
@@ -2173,7 +2172,7 @@ def generate_child_message_computation(
     @abstractmethod
     def add_child_message_computation(
         self,
-        child_message_computation: 'MessageComputationAPI',
+        child_message_computation: "MessageComputationAPI",
     ) -> None:
         """
         Add the given ``child_computation``.
@@ -2232,10 +2231,11 @@ def get_log_entries(self) -> Tuple[Tuple[bytes, Tuple[int, ...], bytes], ...]:
     @classmethod
     @abstractmethod
     def apply_message(
-            cls,
-            state: 'StateAPI',
-            message: MessageAPI,
-            transaction_context: TransactionContextAPI) -> 'MessageComputationAPI':
+        cls,
+        state: "StateAPI",
+        message: MessageAPI,
+        transaction_context: TransactionContextAPI,
+    ) -> "MessageComputationAPI":
         """
         Execute a VM message. This is where the VM-specific call logic exists.
         """

diff --git a/eth/vm/computation/base_computation.py b/eth/vm/computation/base_computation.py
@@ -3,9 +3,12 @@
     Any,
     Callable,
     Dict,
+    Generic,
+    List,
     Optional,
     Tuple,
     Type,
+    TypeVar,
     Union,
 )
 
@@ -73,7 +76,10 @@ def memory_gas_cost(size_in_bytes: int) -> int:
     return total_cost
 
 
-class BaseComputation(Configurable, ComputationAPI):
+C = TypeVar("C", bound="ComputationAPI")
+
+
+class BaseComputation(ComputationAPI, Configurable, Generic[C]):
     """
     The base class for all execution computations.
 
@@ -92,8 +98,8 @@ class BaseComputation(Configurable, ComputationAPI):
 
     state: StateAPI = None
     code: CodeStreamAPI = None
+    children: List[C] = None
     return_data: bytes = b''
-    accounts_to_delete: Dict[Address, Address] = None
 
     _memory: MemoryAPI = None
     _stack: StackAPI = None
@@ -107,6 +113,8 @@ class BaseComputation(Configurable, ComputationAPI):
 
     def __init__(self, state: StateAPI) -> None:
         self.state = state
+        self.children = []
+
         self._memory = Memory()
         self._stack = Stack()
 

diff --git a/eth/vm/computation/message_computation.py b/eth/vm/computation/message_computation.py
@@ -12,7 +12,8 @@
 )
 
 from eth.vm.computation.base_computation import (
-    BaseComputation, NO_RESULT,
+    BaseComputation,
+    NO_RESULT,
 )
 from eth_typing import (
     Address,
@@ -54,7 +55,10 @@
 )
 
 
-class MessageComputation(BaseComputation, MessageComputationAPI):
+class MessageComputation(
+    MessageComputationAPI,
+    BaseComputation[MessageComputationAPI],
+):
     """
     A class for executing message computations.
     """
@@ -74,14 +78,13 @@ def __init__(
         message: MessageAPI,
         transaction_context: TransactionContextAPI,
     ) -> None:
-        super().__init__(state)
+        BaseComputation.__init__(self, state)
 
         self.msg = message
         self.transaction_context = transaction_context
         self.code = CodeStream(message.code)
         self._gas_meter = self._configure_gas_meter()
 
-        self.children = []
         self.accounts_to_delete = {}
         self._log_entries = []