create transient storage snapshots for process_create_message

src/ethereum/cancun/fork.py

@@ -51,6 +51,7 @@
 )
 from .state import (
     State,
+    TransientStorage,
     account_exists_and_is_empty,
     destroy_account,
     destroy_touched_empty_accounts,
@@ -556,6 +557,7 @@ def apply_body(
         traces=[],
         excess_blob_gas=excess_blob_gas,
         blob_versioned_hashes=(),
+        transient_storage=TransientStorage(),
     )
 
     system_tx_output = process_message_call(system_tx_message, system_tx_env)
@@ -591,6 +593,7 @@ def apply_body(
             traces=[],
             excess_blob_gas=excess_blob_gas,
             blob_versioned_hashes=blob_versioned_hashes,
+            transient_storage=TransientStorage(),
         )
 
         gas_used, logs, error = process_transaction(env, tx)

src/ethereum/cancun/state.py

@@ -71,7 +71,7 @@ def close_state(state: State) -> None:
 
 
 def begin_transaction(
-    state: State, transient_storage: Optional[TransientStorage] = None
+    state: State, transient_storage: TransientStorage
 ) -> None:
     """
     Start a state transaction.
@@ -92,14 +92,13 @@ def begin_transaction(
             {k: copy_trie(t) for (k, t) in state._storage_tries.items()},
         )
     )
-    if transient_storage is not None:
-        transient_storage._snapshots.append(
-            {k: copy_trie(t) for (k, t) in transient_storage._tries.items()}
-        )
+    transient_storage._snapshots.append(
+        {k: copy_trie(t) for (k, t) in transient_storage._tries.items()}
+    )
 
 
 def commit_transaction(
-    state: State, transient_storage: Optional[TransientStorage] = None
+    state: State, transient_storage: TransientStorage
 ) -> None:
     """
     Commit a state transaction.
@@ -115,12 +114,11 @@ def commit_transaction(
     if not state._snapshots:
         state.created_accounts.clear()
 
-    if transient_storage and transient_storage._snapshots:
-        transient_storage._snapshots.pop()
+    transient_storage._snapshots.pop()
 
 
 def rollback_transaction(
-    state: State, transient_storage: Optional[TransientStorage] = None
+    state: State, transient_storage: TransientStorage
 ) -> None:
     """
     Rollback a state transaction, resetting the state to the point when the
@@ -137,8 +135,7 @@ def rollback_transaction(
     if not state._snapshots:
         state.created_accounts.clear()
 
-    if transient_storage and transient_storage._snapshots:
-        transient_storage._tries = transient_storage._snapshots.pop()
+    transient_storage._tries = transient_storage._snapshots.pop()
 
 
 def get_account(state: State, address: Address) -> Account:

src/ethereum/cancun/vm/__init__.py

@@ -47,6 +47,7 @@ class Environment:
     traces: List[dict]
     excess_blob_gas: U64
     blob_versioned_hashes: Tuple[VersionedHash, ...]
+    transient_storage: TransientStorage
 
 
 @dataclass
@@ -93,7 +94,6 @@ class Evm:
     error: Optional[Exception]
     accessed_addresses: Set[Address]
     accessed_storage_keys: Set[Tuple[Address, Bytes32]]
-    transient_storage: TransientStorage
 
 
 def incorporate_child_on_success(evm: Evm, child_evm: Evm) -> None:

src/ethereum/cancun/vm/instructions/storage.py

@@ -149,7 +149,7 @@ def tload(evm: Evm) -> None:
 
     # OPERATION
     value = get_transient_storage(
-        evm.transient_storage, evm.message.current_target, key
+        evm.env.transient_storage, evm.message.current_target, key
     )
     push(evm.stack, value)
 
@@ -175,7 +175,7 @@ def tstore(evm: Evm) -> None:
     # OPERATION
     ensure(not evm.message.is_static, WriteInStaticContext)
     set_transient_storage(
-        evm.transient_storage, evm.message.current_target, key, new_value
+        evm.env.transient_storage, evm.message.current_target, key, new_value
     )
 
     # PROGRAM COUNTER

src/ethereum/cancun/vm/interpreter.py

@@ -29,7 +29,6 @@
 
 from ..fork_types import Address, Log
 from ..state import (
-    TransientStorage,
     account_exists_and_is_empty,
     account_has_code_or_nonce,
     begin_transaction,
@@ -161,7 +160,7 @@ def process_create_message(message: Message, env: Environment) -> Evm:
         Items containing execution specific objects.
     """
     # take snapshot of state before processing the message
-    begin_transaction(env.state)
+    begin_transaction(env.state, env.transient_storage)
 
     # If the address where the account is being created has storage, it is
     # destroyed. This can only happen in the following highly unlikely
@@ -188,15 +187,15 @@ def process_create_message(message: Message, env: Environment) -> Evm:
             charge_gas(evm, contract_code_gas)
             ensure(len(contract_code) <= MAX_CODE_SIZE, OutOfGasError)
         except ExceptionalHalt as error:
-            rollback_transaction(env.state)
+            rollback_transaction(env.state, env.transient_storage)
             evm.gas_left = Uint(0)
             evm.output = b""
             evm.error = error
         else:
             set_code(env.state, message.current_target, contract_code)
-            commit_transaction(env.state)
+            commit_transaction(env.state, env.transient_storage)
     else:
-        rollback_transaction(env.state)
+        rollback_transaction(env.state, env.transient_storage)
     return evm
 
 
@@ -220,10 +219,7 @@ def process_message(message: Message, env: Environment) -> Evm:
         raise StackDepthLimitError("Stack depth limit reached")
 
     # take snapshot of state before processing the message
-    transient_storage = (
-        message.parent_evm.transient_storage if message.parent_evm else None
-    )
-    begin_transaction(env.state, transient_storage)
+    begin_transaction(env.state, env.transient_storage)
 
     touch_account(env.state, message.current_target)
 
@@ -236,9 +232,9 @@ def process_message(message: Message, env: Environment) -> Evm:
     if evm.error:
         # revert state to the last saved checkpoint
         # since the message call resulted in an error
-        rollback_transaction(env.state, evm.transient_storage)
+        rollback_transaction(env.state, env.transient_storage)
     else:
-        commit_transaction(env.state, evm.transient_storage)
+        commit_transaction(env.state, env.transient_storage)
     return evm
 
 
@@ -260,10 +256,6 @@ def execute_code(message: Message, env: Environment) -> Evm:
     """
     code = message.code
     valid_jump_destinations = get_valid_jump_destinations(code)
-    if message.parent_evm:
-        transient_storage = message.parent_evm.transient_storage
-    else:
-        transient_storage = TransientStorage()
 
     evm = Evm(
         pc=Uint(0),
@@ -284,7 +276,6 @@ def execute_code(message: Message, env: Environment) -> Evm:
         error=None,
         accessed_addresses=message.accessed_addresses,
         accessed_storage_keys=message.accessed_storage_keys,
-        transient_storage=transient_storage,
     )
     try:
         if evm.message.code_address in PRE_COMPILED_CONTRACTS:

src/ethereum_spec_tools/evm_tools/loaders/fork_loader.py

@@ -182,6 +182,11 @@ def State(self) -> Any:
         """State class of the fork"""
         return self._module("state").State
 
+    @property
+    def TransientStorage(self) -> Any:
+        """Transient storage class of the fork"""
+        return self._module("state").TransientStorage
+
     @property
     def get_account(self) -> Any:
         """get_account function of the fork"""

src/ethereum_spec_tools/evm_tools/t8n/__init__.py

@@ -224,6 +224,7 @@ def environment(self, tx: Any, gas_available: Any) -> Any:
 
         if self.fork.is_after_fork("ethereum.cancun"):
             kw_arguments["excess_blob_gas"] = self.env.excess_blob_gas
+            kw_arguments["transient_storage"] = self.fork.TransientStorage()
 
         return self.fork.Environment(**kw_arguments)
 
@@ -356,6 +357,7 @@ def apply_body(self) -> None:
                 traces=[],
                 excess_blob_gas=self.env.excess_blob_gas,
                 blob_versioned_hashes=(),
+                transient_storage=self.fork.TransientStorage(),
             )
 
             system_tx_output = self.fork.process_message_call(