|
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.""" |
11 | 2 |
|
12 | 3 | from __future__ import annotations |
13 | 4 |
|
| 5 | +from collections.abc import Sequence |
14 | 6 | from functools import singledispatch |
15 | 7 | from math import ceil |
16 | 8 |
|
|
19 | 11 | from lean_spec.types.bitfields import BaseBitlist, BaseBitvector |
20 | 12 | from lean_spec.types.boolean import Boolean |
21 | 13 | 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 |
26 | 15 | from lean_spec.types.container import Container |
27 | 16 | from lean_spec.types.uint import BaseUint |
28 | 17 |
|
29 | 18 | 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")) |
31 | 56 |
|
32 | 57 |
|
33 | 58 | @singledispatch |
34 | 59 | 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. |
38 | 61 |
|
39 | 62 | Raises: |
40 | | - TypeError: If the value type has no registered implementation. |
| 63 | + TypeError: If the value's type has no registered handler. |
41 | 64 | """ |
42 | 65 | raise TypeError(f"hash_tree_root: unsupported value type {type(value).__name__}") |
43 | 66 |
|
44 | 67 |
|
45 | 68 | @hash_tree_root.register |
46 | 69 | 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())) |
49 | 71 |
|
50 | 72 |
|
51 | 73 | @hash_tree_root.register |
52 | 74 | def _htr_boolean(value: Boolean) -> Bytes32: |
53 | | - return merkleize(pack_bytes(value.encode_bytes())) |
| 75 | + return merkleize(_pack_bytes(value.encode_bytes())) |
54 | 76 |
|
55 | 77 |
|
56 | 78 | @hash_tree_root.register |
57 | 79 | 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())) |
60 | 81 |
|
61 | 82 |
|
62 | 83 | @hash_tree_root.register |
63 | 84 | 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)) |
66 | 86 |
|
67 | 87 |
|
68 | 88 | @hash_tree_root.register |
69 | 89 | def _htr_bytearray(value: bytearray) -> Bytes32: |
70 | | - return merkleize(pack_bytes(bytes(value))) |
| 90 | + return merkleize(_pack_bytes(bytes(value))) |
71 | 91 |
|
72 | 92 |
|
73 | 93 | @hash_tree_root.register |
74 | 94 | def _htr_memoryview(value: memoryview) -> Bytes32: |
75 | | - return merkleize(pack_bytes(value.tobytes())) |
| 95 | + return merkleize(_pack_bytes(value.tobytes())) |
76 | 96 |
|
77 | 97 |
|
78 | 98 | @hash_tree_root.register |
79 | 99 | def _htr_bytevector(value: BaseBytes) -> Bytes32: |
80 | | - return merkleize(pack_bytes(value.encode_bytes())) |
| 100 | + return merkleize(_pack_bytes(value.encode_bytes())) |
81 | 101 |
|
82 | 102 |
|
83 | 103 | @hash_tree_root.register |
84 | 104 | def _htr_bytelist(value: BaseByteList) -> Bytes32: |
85 | 105 | data = value.encode_bytes() |
86 | | - # Compute limit in chunks and merkleize the packed bytes |
87 | 106 | 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)) |
90 | 108 |
|
91 | 109 |
|
92 | 110 | @hash_tree_root.register |
93 | 111 | 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) |
97 | 114 |
|
98 | 115 |
|
99 | 116 | @hash_tree_root.register |
100 | 117 | 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) |
103 | 119 | 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), |
105 | 121 | len(value.data), |
106 | 122 | ) |
107 | 123 |
|
108 | 124 |
|
109 | 125 | @hash_tree_root.register |
110 | 126 | 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 |
113 | 129 | 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. |
115 | 131 | 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) |
118 | 133 | 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)), |
120 | 135 | limit=limit_chunks, |
121 | 136 | ) |
122 | | - |
123 | | - # COMPOSITE elements: merkleize child roots with limit = length |
| 137 | + # Composite elements each contribute their own hash tree root as a leaf. |
124 | 138 | return merkleize([hash_tree_root(e) for e in value], limit=length) |
125 | 139 |
|
126 | 140 |
|
127 | 141 | @hash_tree_root.register |
128 | 142 | 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 |
131 | 145 | if issubclass(elem_t, (BaseUint, Boolean, Fp)): |
132 | | - # BASIC elements: pack serialized bytes |
133 | 146 | 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) |
136 | 148 | 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)), |
138 | 150 | limit=limit_chunks, |
139 | 151 | ) |
140 | 152 | else: |
141 | | - # COMPOSITE elements: merkleize child roots |
142 | 153 | root = merkleize([hash_tree_root(e) for e in value], limit=limit) |
143 | | - |
144 | | - # Mix in the length for both cases |
145 | 154 | return mix_in_length(root, len(value)) |
146 | 155 |
|
147 | 156 |
|
148 | 157 | @hash_tree_root.register |
149 | 158 | 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