Skip to content

Commit 82f7c92

Browse files
tcoratgerclaude
andauthored
refactor(ssz): consolidate to 4 files, modernize idioms, tighten tests (leanEthereum#782)
Source (src/lean_spec/subspecs/ssz/): - Delete pack.py and utils.py; inline their helpers as private members of hash.py and merkleization.py - Collapse the 4 pack functions into a single _pack_bytes using bytes.ljust - Rewrite _pack_bits using int.to_bytes instead of bytearray bit-twiddling - Inline hash_nodes at its 5 call sites - Remove dead code in _merkleize_efficient (unreachable second loop and zero-tree fallback); add an assertion that documents the invariant - Standardize ceiling division on math.ceil everywhere - Rewrite docstrings per .claude/rules/documentation.md (no backticks, no identifier names in prose, one sentence per line) - Document constants with WHY, not what - Add ASCII diagrams to merkleize and _pack_bits Tests (tests/lean_spec/subspecs/ssz/): - Delete fork-specific tests that were misplaced in the SSZ folder (test_block, test_state, test_signed_attestation) - Delete tests redundant with the public-surface coverage (test_pack, test_utils_ssz, test_boundary_values, test_edge_cases, test_malformed_data, test_nested_structures) - Rewrite test_hash.py with parametric coverage of every dispatch arm built from first-principles expected values - Rewrite test_merkleization.py absorbing _next_pow2 coverage - Every test function has a one-line docstring Result: 100% line + branch coverage on all 4 SSZ source files (107 tests). Source goes from 6 files / 401 lines to 4 files / 332 lines. Tests go from 9 files / 2,585 lines to 2 files / 877 lines. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent f49fa78 commit 82f7c92

16 files changed

Lines changed: 732 additions & 2767 deletions
Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
1-
"""Constants defined in the SSZ specification."""
1+
"""Constants used by SSZ merkleization."""
22

33
from __future__ import annotations
44

55
from typing import Final
66

77
BYTES_PER_CHUNK: Final = 32
8-
"""Number of bytes per Merkle chunk."""
8+
"""Width of a Merkle leaf chunk in bytes."""
99

10-
BITS_PER_BYTE: Final = 8
11-
"""Number of bits per byte."""
12-
13-
BITS_PER_CHUNK: Final = BYTES_PER_CHUNK * BITS_PER_BYTE
14-
"""Number of bits per Merkle chunk (256 bits)."""
10+
BITS_PER_CHUNK: Final = BYTES_PER_CHUNK * 8
11+
"""Width of a Merkle leaf chunk in bits."""

src/lean_spec/subspecs/ssz/hash.py

Lines changed: 67 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,8 @@
1-
"""SSZ Merkleization entry point.
2-
3-
Computes Merkle roots for SSZ types.
4-
5-
Handles:
6-
7-
- Basic types: pack into chunks
8-
- Composite types: merkleize child roots
9-
- Variable-size types: mix in length
10-
"""
1+
"""Hash tree root dispatch for SSZ values."""
112

123
from __future__ import annotations
134

5+
from collections.abc import Sequence
146
from functools import singledispatch
157
from math import ceil
168

@@ -19,133 +11,151 @@
1911
from lean_spec.types.bitfields import BaseBitlist, BaseBitvector
2012
from lean_spec.types.boolean import Boolean
2113
from lean_spec.types.byte_arrays import BaseByteList, BaseBytes, Bytes32
22-
from lean_spec.types.collections import (
23-
SSZList,
24-
SSZVector,
25-
)
14+
from lean_spec.types.collections import SSZList, SSZVector
2615
from lean_spec.types.container import Container
2716
from lean_spec.types.uint import BaseUint
2817

2918
from .merkleization import merkleize, mix_in_length
30-
from .pack import pack_bits, pack_bytes
19+
20+
21+
def _pack_bytes(data: bytes) -> list[Bytes32]:
22+
"""Right-pad serialized bytes to a chunk boundary and split into chunks.
23+
24+
Layout for a 5-byte payload:
25+
26+
bytes : 01 02 03 04 05
27+
padded : 01 02 03 04 05 00 00 ... 00 (zero-padded to 32 bytes)
28+
chunks : [ Bytes32(01 02 03 04 05 00 ...) ]
29+
30+
Inner chunks are already chunk-aligned; only the trailing chunk is padded.
31+
"""
32+
return [
33+
Bytes32(data[i : i + BYTES_PER_CHUNK].ljust(BYTES_PER_CHUNK, b"\x00"))
34+
for i in range(0, len(data), BYTES_PER_CHUNK)
35+
]
36+
37+
38+
def _pack_bits(bits: Sequence[Boolean]) -> list[Bytes32]:
39+
"""Pack a boolean sequence into bytes, then into chunks for merkleization.
40+
41+
The first input bit becomes the least significant bit of the first byte.
42+
Each next input bit moves up one position, wrapping to the next byte after eight.
43+
44+
Layout for [1, 0, 1, 1]:
45+
46+
bit position : 7 6 5 4 3 2 1 0
47+
byte 0 : 0 0 0 0 1 1 0 1
48+
^ ^ ^ ^
49+
3 2 1 0 <- input order
50+
51+
The SSZ serialization delimiter and the length-mix are separate steps,
52+
handled by the caller when needed.
53+
"""
54+
value = sum(1 << i for i, bit in enumerate(bits) if bit)
55+
return _pack_bytes(value.to_bytes(ceil(len(bits) / 8), "little"))
3156

3257

3358
@singledispatch
3459
def hash_tree_root(value: object) -> Bytes32:
35-
"""Compute the Merkle root for an SSZ value.
36-
37-
Dispatches to type-specific implementations.
60+
"""Compute the SSZ Merkle root of a value.
3861
3962
Raises:
40-
TypeError: If the value type has no registered implementation.
63+
TypeError: If the value's type has no registered handler.
4164
"""
4265
raise TypeError(f"hash_tree_root: unsupported value type {type(value).__name__}")
4366

4467

4568
@hash_tree_root.register
4669
def _htr_uint(value: BaseUint) -> Bytes32:
47-
"""Basic scalars: pack bytes into chunks and merkleize."""
48-
return merkleize(pack_bytes(value.encode_bytes()))
70+
return merkleize(_pack_bytes(value.encode_bytes()))
4971

5072

5173
@hash_tree_root.register
5274
def _htr_boolean(value: Boolean) -> Bytes32:
53-
return merkleize(pack_bytes(value.encode_bytes()))
75+
return merkleize(_pack_bytes(value.encode_bytes()))
5476

5577

5678
@hash_tree_root.register
5779
def _htr_fp(value: Fp) -> Bytes32:
58-
"""KoalaBear field elements: pack bytes into chunks and merkleize."""
59-
return merkleize(pack_bytes(value.encode_bytes()))
80+
return merkleize(_pack_bytes(value.encode_bytes()))
6081

6182

6283
@hash_tree_root.register
6384
def _htr_bytes(value: bytes) -> Bytes32:
64-
"""Treat raw bytes like ByteVector[N]."""
65-
return merkleize(pack_bytes(value))
85+
return merkleize(_pack_bytes(value))
6686

6787

6888
@hash_tree_root.register
6989
def _htr_bytearray(value: bytearray) -> Bytes32:
70-
return merkleize(pack_bytes(bytes(value)))
90+
return merkleize(_pack_bytes(bytes(value)))
7191

7292

7393
@hash_tree_root.register
7494
def _htr_memoryview(value: memoryview) -> Bytes32:
75-
return merkleize(pack_bytes(value.tobytes()))
95+
return merkleize(_pack_bytes(value.tobytes()))
7696

7797

7898
@hash_tree_root.register
7999
def _htr_bytevector(value: BaseBytes) -> Bytes32:
80-
return merkleize(pack_bytes(value.encode_bytes()))
100+
return merkleize(_pack_bytes(value.encode_bytes()))
81101

82102

83103
@hash_tree_root.register
84104
def _htr_bytelist(value: BaseByteList) -> Bytes32:
85105
data = value.encode_bytes()
86-
# Compute limit in chunks and merkleize the packed bytes
87106
limit_chunks = ceil(type(value).LIMIT / BYTES_PER_CHUNK)
88-
# Mix in the length of the data
89-
return mix_in_length(merkleize(pack_bytes(data), limit=limit_chunks), len(data))
107+
return mix_in_length(merkleize(_pack_bytes(data), limit=limit_chunks), len(data))
90108

91109

92110
@hash_tree_root.register
93111
def _htr_bitvector_base(value: BaseBitvector) -> Bytes32:
94-
# Compute limit in chunks using ceiling division
95-
limit = (type(value).LENGTH + BITS_PER_CHUNK - 1) // BITS_PER_CHUNK
96-
return merkleize(pack_bits(tuple(bool(b) for b in value.data)), limit=limit)
112+
limit = ceil(type(value).LENGTH / BITS_PER_CHUNK)
113+
return merkleize(_pack_bits(value.data), limit=limit)
97114

98115

99116
@hash_tree_root.register
100117
def _htr_bitlist_base(value: BaseBitlist) -> Bytes32:
101-
# Compute limit in chunks using ceiling division
102-
limit = (type(value).LIMIT + BITS_PER_CHUNK - 1) // BITS_PER_CHUNK
118+
limit = ceil(type(value).LIMIT / BITS_PER_CHUNK)
103119
return mix_in_length(
104-
merkleize(pack_bits(tuple(bool(b) for b in value.data)), limit=limit),
120+
merkleize(_pack_bits(value.data), limit=limit),
105121
len(value.data),
106122
)
107123

108124

109125
@hash_tree_root.register
110126
def _htr_vector(value: SSZVector) -> Bytes32:
111-
elem_t, length = type(value).ELEMENT_TYPE, type(value).LENGTH
112-
127+
cls = type(value)
128+
elem_t, length = cls.ELEMENT_TYPE, cls.LENGTH
113129
if issubclass(elem_t, (BaseUint, Boolean, Fp)):
114-
# BASIC elements: pack serialized bytes
130+
# Basic elements pack their serialized bytes into a single byte stream before chunking.
115131
elem_size = elem_t.get_byte_length()
116-
# Compute limit in chunks: ceil((length * elem_size) / BYTES_PER_CHUNK)
117-
limit_chunks = (length * elem_size + BYTES_PER_CHUNK - 1) // BYTES_PER_CHUNK
132+
limit_chunks = ceil(length * elem_size / BYTES_PER_CHUNK)
118133
return merkleize(
119-
pack_bytes(b"".join(e.encode_bytes() for e in value)),
134+
_pack_bytes(b"".join(e.encode_bytes() for e in value)),
120135
limit=limit_chunks,
121136
)
122-
123-
# COMPOSITE elements: merkleize child roots with limit = length
137+
# Composite elements each contribute their own hash tree root as a leaf.
124138
return merkleize([hash_tree_root(e) for e in value], limit=length)
125139

126140

127141
@hash_tree_root.register
128142
def _htr_list(value: SSZList) -> Bytes32:
129-
elem_t, limit = type(value).ELEMENT_TYPE, type(value).LIMIT
130-
143+
cls = type(value)
144+
elem_t, limit = cls.ELEMENT_TYPE, cls.LIMIT
131145
if issubclass(elem_t, (BaseUint, Boolean, Fp)):
132-
# BASIC elements: pack serialized bytes
133146
elem_size = elem_t.get_byte_length()
134-
# Compute limit in chunks: ceil((limit * elem_size) / BYTES_PER_CHUNK)
135-
limit_chunks = (limit * elem_size + BYTES_PER_CHUNK - 1) // BYTES_PER_CHUNK
147+
limit_chunks = ceil(limit * elem_size / BYTES_PER_CHUNK)
136148
root = merkleize(
137-
pack_bytes(b"".join(e.encode_bytes() for e in value)),
149+
_pack_bytes(b"".join(e.encode_bytes() for e in value)),
138150
limit=limit_chunks,
139151
)
140152
else:
141-
# COMPOSITE elements: merkleize child roots
142153
root = merkleize([hash_tree_root(e) for e in value], limit=limit)
143-
144-
# Mix in the length for both cases
145154
return mix_in_length(root, len(value))
146155

147156

148157
@hash_tree_root.register
149158
def _htr_container(value: Container) -> Bytes32:
150-
# Preserve declared field order from the Pydantic model
151-
return merkleize([hash_tree_root(getattr(value, fname)) for fname in type(value).model_fields])
159+
# Pydantic preserves declaration order, which is the canonical SSZ field order.
160+
cls = type(value)
161+
return merkleize([hash_tree_root(getattr(value, name)) for name in cls.model_fields])

0 commit comments

Comments
 (0)