Skip to content

Commit

Permalink
dev: change type of branch node subnodes (#1066)
Browse files Browse the repository at this point in the history 
* dev: change type of branch node subnodes

* Annotate tuple with 16

* add ellipsis to tuple

* explicit tuple

* use explicit 16-field tuple

* fmt

* remove Annotated import

* fix types

* Use assert_type alongside cast (#1)

* Import from typing_extensions instead of typing (#2)

---------

Co-authored-by: Sam Wilson <57262657+SamWilsn@users.noreply.github.com>
  • Loading branch information
@enitrat @SamWilsn
enitrat and SamWilsn authored Jan 25, 2025
1 parent e4d7b40 commit 26ab1bc
Show file tree
Hide file tree
Showing 17 changed files with 481 additions and 97 deletions.
2 changes: 1 addition & 1 deletion setup.cfg
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ python_requires = >=3.10
install_requires =
pycryptodome>=3,<4
coincurve>=20,<21
typing_extensions>=4
typing_extensions>=4.2
py_ecc @ git+https://github.com/petertdavies/py_ecc.git@127184f4c57b1812da959586d0fe8f43bb1a2389
ethereum-types>=0.2.1,<0.3
ethereum-rlp>=0.1.1,<0.2
Expand Down
36 changes: 30 additions & 6 deletions src/ethereum/arrow_glacier/trie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,17 @@
MutableMapping,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
cast,
)

from ethereum_rlp import rlp
from ethereum_types.bytes import Bytes
from ethereum_types.frozen import slotted_freezable
from ethereum_types.numeric import U256, Uint
from typing_extensions import assert_type

from ethereum.crypto.hash import keccak256
from ethereum.london import trie as previous_trie
Expand Down Expand Up @@ -92,12 +95,32 @@ class ExtensionNode:
subnode: rlp.Extended


BranchSubnodes = Tuple[
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
]


@slotted_freezable
@dataclass
class BranchNode:
"""Branch node in the Merkle Trie"""

subnodes: List[rlp.Extended]
subnodes: BranchSubnodes
value: rlp.Extended


Expand Down Expand Up @@ -137,7 +160,7 @@ def encode_internal_node(node: Optional[InternalNode]) -> rlp.Extended:
node.subnode,
)
elif isinstance(node, BranchNode):
unencoded = node.subnodes + [node.value]
unencoded = list(node.subnodes) + [node.value]
else:
raise AssertionError(f"Invalid internal node type {type(node)}!")

Expand Down Expand Up @@ -458,10 +481,11 @@ def patricialize(
else:
branches[key[level]][key] = obj[key]

subnodes = tuple(
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
)
return BranchNode(
[
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
],
cast(BranchSubnodes, assert_type(subnodes, Tuple[rlp.Extended, ...])),
value,
)
36 changes: 30 additions & 6 deletions src/ethereum/berlin/trie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,17 @@
MutableMapping,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
cast,
)

from ethereum_rlp import rlp
from ethereum_types.bytes import Bytes
from ethereum_types.frozen import slotted_freezable
from ethereum_types.numeric import U256, Uint
from typing_extensions import assert_type

from ethereum.crypto.hash import keccak256
from ethereum.muir_glacier import trie as previous_trie
Expand Down Expand Up @@ -92,12 +95,32 @@ class ExtensionNode:
subnode: rlp.Extended


BranchSubnodes = Tuple[
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
]


@slotted_freezable
@dataclass
class BranchNode:
"""Branch node in the Merkle Trie"""

subnodes: List[rlp.Extended]
subnodes: BranchSubnodes
value: rlp.Extended


Expand Down Expand Up @@ -137,7 +160,7 @@ def encode_internal_node(node: Optional[InternalNode]) -> rlp.Extended:
node.subnode,
)
elif isinstance(node, BranchNode):
unencoded = node.subnodes + [node.value]
unencoded = list(node.subnodes) + [node.value]
else:
raise AssertionError(f"Invalid internal node type {type(node)}!")

Expand Down Expand Up @@ -458,10 +481,11 @@ def patricialize(
else:
branches[key[level]][key] = obj[key]

subnodes = tuple(
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
)
return BranchNode(
[
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
],
cast(BranchSubnodes, assert_type(subnodes, Tuple[rlp.Extended, ...])),
value,
)
36 changes: 30 additions & 6 deletions src/ethereum/byzantium/trie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,17 @@
MutableMapping,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
cast,
)

from ethereum_rlp import rlp
from ethereum_types.bytes import Bytes
from ethereum_types.frozen import slotted_freezable
from ethereum_types.numeric import U256, Uint
from typing_extensions import assert_type

from ethereum.crypto.hash import keccak256
from ethereum.spurious_dragon import trie as previous_trie
Expand Down Expand Up @@ -92,12 +95,32 @@ class ExtensionNode:
subnode: rlp.Extended


BranchSubnodes = Tuple[
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
]


@slotted_freezable
@dataclass
class BranchNode:
"""Branch node in the Merkle Trie"""

subnodes: List[rlp.Extended]
subnodes: BranchSubnodes
value: rlp.Extended


Expand Down Expand Up @@ -137,7 +160,7 @@ def encode_internal_node(node: Optional[InternalNode]) -> rlp.Extended:
node.subnode,
)
elif isinstance(node, BranchNode):
unencoded = node.subnodes + [node.value]
unencoded = list(node.subnodes) + [node.value]
else:
raise AssertionError(f"Invalid internal node type {type(node)}!")

Expand Down Expand Up @@ -458,10 +481,11 @@ def patricialize(
else:
branches[key[level]][key] = obj[key]

subnodes = tuple(
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
)
return BranchNode(
[
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
],
cast(BranchSubnodes, assert_type(subnodes, Tuple[rlp.Extended, ...])),
value,
)
36 changes: 30 additions & 6 deletions src/ethereum/cancun/trie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,17 @@
MutableMapping,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
cast,
)

from ethereum_rlp import rlp
from ethereum_types.bytes import Bytes
from ethereum_types.frozen import slotted_freezable
from ethereum_types.numeric import U256, Uint
from typing_extensions import assert_type

from ethereum.crypto.hash import keccak256
from ethereum.shanghai import trie as previous_trie
Expand Down Expand Up @@ -95,12 +98,32 @@ class ExtensionNode:
subnode: rlp.Extended


BranchSubnodes = Tuple[
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
]


@slotted_freezable
@dataclass
class BranchNode:
"""Branch node in the Merkle Trie"""

subnodes: List[rlp.Extended]
subnodes: BranchSubnodes
value: rlp.Extended


Expand Down Expand Up @@ -140,7 +163,7 @@ def encode_internal_node(node: Optional[InternalNode]) -> rlp.Extended:
node.subnode,
)
elif isinstance(node, BranchNode):
unencoded = node.subnodes + [node.value]
unencoded = list(node.subnodes) + [node.value]
else:
raise AssertionError(f"Invalid internal node type {type(node)}!")

Expand Down Expand Up @@ -461,10 +484,11 @@ def patricialize(
else:
branches[key[level]][key] = obj[key]

subnodes = tuple(
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
)
return BranchNode(
[
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
],
cast(BranchSubnodes, assert_type(subnodes, Tuple[rlp.Extended, ...])),
value,
)
36 changes: 30 additions & 6 deletions src/ethereum/constantinople/trie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,17 @@
MutableMapping,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
cast,
)

from ethereum_rlp import rlp
from ethereum_types.bytes import Bytes
from ethereum_types.frozen import slotted_freezable
from ethereum_types.numeric import U256, Uint
from typing_extensions import assert_type

from ethereum.byzantium import trie as previous_trie
from ethereum.crypto.hash import keccak256
Expand Down Expand Up @@ -92,12 +95,32 @@ class ExtensionNode:
subnode: rlp.Extended


BranchSubnodes = Tuple[
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
rlp.Extended,
]


@slotted_freezable
@dataclass
class BranchNode:
"""Branch node in the Merkle Trie"""

subnodes: List[rlp.Extended]
subnodes: BranchSubnodes
value: rlp.Extended


Expand Down Expand Up @@ -137,7 +160,7 @@ def encode_internal_node(node: Optional[InternalNode]) -> rlp.Extended:
node.subnode,
)
elif isinstance(node, BranchNode):
unencoded = node.subnodes + [node.value]
unencoded = list(node.subnodes) + [node.value]
else:
raise AssertionError(f"Invalid internal node type {type(node)}!")

Expand Down Expand Up @@ -458,10 +481,11 @@ def patricialize(
else:
branches[key[level]][key] = obj[key]

subnodes = tuple(
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
)
return BranchNode(
[
encode_internal_node(patricialize(branches[k], level + Uint(1)))
for k in range(16)
],
cast(BranchSubnodes, assert_type(subnodes, Tuple[rlp.Extended, ...])),
value,
)
Loading

0 comments on commit 26ab1bc

Please sign in to comment.