dev: update storage trie type

ethereum · Jan 6, 2025 · 338b3cb · 338b3cb
1 parent 4c7eaa8
commit 338b3cb
Show file tree

Hide file tree

Showing 17 changed files with 113 additions and 97 deletions.
diff --git a/src/ethereum/arrow_glacier/state.py b/src/ethereum/arrow_glacier/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, Set, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -36,12 +36,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
     created_accounts: Set[Address] = field(default_factory=set)
@@ -228,7 +229,7 @@ def mark_account_created(state: State, address: Address) -> None:
     state.created_accounts.add(address)
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -258,7 +259,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes
@@ -580,7 +581,7 @@ def increase_balance(account: Account) -> None:
     modify_state(state, address, increase_balance)
 
 
-def get_storage_original(state: State, address: Address, key: Bytes) -> U256:
+def get_storage_original(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get the original value in a storage slot i.e. the value before the current
     transaction began. This function reads the value from the snapshots taken

diff --git a/src/ethereum/berlin/state.py b/src/ethereum/berlin/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, Set, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -36,12 +36,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
     created_accounts: Set[Address] = field(default_factory=set)
@@ -228,7 +229,7 @@ def mark_account_created(state: State, address: Address) -> None:
     state.created_accounts.add(address)
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -258,7 +259,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes
@@ -580,7 +581,7 @@ def increase_balance(account: Account) -> None:
     modify_state(state, address, increase_balance)
 
 
-def get_storage_original(state: State, address: Address, key: Bytes) -> U256:
+def get_storage_original(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get the original value in a storage slot i.e. the value before the current
     transaction began. This function reads the value from the snapshots taken

diff --git a/src/ethereum/byzantium/state.py b/src/ethereum/byzantium/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -36,12 +36,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
 
@@ -202,7 +203,7 @@ def destroy_storage(state: State, address: Address) -> None:
         del state._storage_tries[address]
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -232,7 +233,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes

diff --git a/src/ethereum/cancun/state.py b/src/ethereum/cancun/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -37,12 +37,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
     created_accounts: Set[Address] = field(default_factory=set)
@@ -55,8 +56,8 @@ class TransientStorage:
     within a transaction.
     """
 
-    _tries: Dict[Address, Trie[Bytes, U256]] = field(default_factory=dict)
-    _snapshots: List[Dict[Address, Trie[Bytes, U256]]] = field(
+    _tries: Dict[Address, Trie[Bytes32, U256]] = field(default_factory=dict)
+    _snapshots: List[Dict[Address, Trie[Bytes32, U256]]] = field(
         default_factory=list
     )
 
@@ -261,7 +262,7 @@ def mark_account_created(state: State, address: Address) -> None:
     state.created_accounts.add(address)
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -291,7 +292,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes
@@ -607,7 +608,7 @@ def write_code(sender: Account) -> None:
     modify_state(state, address, write_code)
 
 
-def get_storage_original(state: State, address: Address, key: Bytes) -> U256:
+def get_storage_original(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get the original value in a storage slot i.e. the value before the current
     transaction began. This function reads the value from the snapshots taken
@@ -641,7 +642,7 @@ def get_storage_original(state: State, address: Address, key: Bytes) -> U256:
 
 
 def get_transient_storage(
-    transient_storage: TransientStorage, address: Address, key: Bytes
+    transient_storage: TransientStorage, address: Address, key: Bytes32
 ) -> U256:
     """
     Get a value at a storage key on an account from transient storage.
@@ -672,7 +673,7 @@ def get_transient_storage(
 def set_transient_storage(
     transient_storage: TransientStorage,
     address: Address,
-    key: Bytes,
+    key: Bytes32,
     value: U256,
 ) -> None:
     """

diff --git a/src/ethereum/constantinople/state.py b/src/ethereum/constantinople/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -36,12 +36,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
 
@@ -202,7 +203,7 @@ def destroy_storage(state: State, address: Address) -> None:
         del state._storage_tries[address]
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -232,7 +233,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes

diff --git a/src/ethereum/dao_fork/state.py b/src/ethereum/dao_fork/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -36,12 +36,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
 
@@ -202,7 +203,7 @@ def destroy_storage(state: State, address: Address) -> None:
         del state._storage_tries[address]
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -232,7 +233,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes

diff --git a/src/ethereum/frontier/state.py b/src/ethereum/frontier/state.py
@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field
 from typing import Callable, Dict, List, Optional, Tuple
 
-from ethereum_types.bytes import Bytes
+from ethereum_types.bytes import Bytes, Bytes32
 from ethereum_types.frozen import modify
 from ethereum_types.numeric import U256, Uint
 
@@ -36,12 +36,13 @@ class State:
     _main_trie: Trie[Address, Optional[Account]] = field(
         default_factory=lambda: Trie(secured=True, default=None)
     )
-    _storage_tries: Dict[Address, Trie[Bytes, U256]] = field(
+    _storage_tries: Dict[Address, Trie[Bytes32, U256]] = field(
         default_factory=dict
     )
     _snapshots: List[
         Tuple[
-            Trie[Address, Optional[Account]], Dict[Address, Trie[Bytes, U256]]
+            Trie[Address, Optional[Account]],
+            Dict[Address, Trie[Bytes32, U256]],
         ]
     ] = field(default_factory=list)
 
@@ -202,7 +203,7 @@ def destroy_storage(state: State, address: Address) -> None:
         del state._storage_tries[address]
 
 
-def get_storage(state: State, address: Address, key: Bytes) -> U256:
+def get_storage(state: State, address: Address, key: Bytes32) -> U256:
     """
     Get a value at a storage key on an account. Returns `U256(0)` if the
     storage key has not been set previously.
@@ -232,7 +233,7 @@ def get_storage(state: State, address: Address, key: Bytes) -> U256:
 
 
 def set_storage(
-    state: State, address: Address, key: Bytes, value: U256
+    state: State, address: Address, key: Bytes32, value: U256
 ) -> None:
     """
     Set a value at a storage key on an account. Setting to `U256(0)` deletes