xmss: some type simplifications (leanEthereum#206)

Author: tcoratger

src/lean_spec/subspecs/xmss/containers.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Annotated, List
+from typing import TYPE_CHECKING, Annotated, List, cast
 
 from pydantic import Field
 
@@ -179,7 +179,7 @@ class PublicKey(StrictBaseModel):
     All field elements are serialized in little-endian byte order.
     """
 
-    root: List[Fp]
+    root: HashDigestVector
     """The Merkle root, which commits to all one-time keys for the key's lifetime."""
     parameter: Parameter
     """The public parameter `P` that personalizes the hash function."""
@@ -196,7 +196,9 @@ def __bytes__(self) -> bytes:
             >>> isinstance(data, bytes)
             True
         """
-        return Fp.serialize_list(self.root) + Fp.serialize_list(self.parameter)
+        return Fp.serialize_list(cast(List[Fp], list(self.root.data))) + Fp.serialize_list(
+            self.parameter
+        )
 
     def to_bytes(self, config: XmssConfig) -> bytes:
         """
@@ -262,7 +264,7 @@ def from_bytes(cls, data: bytes, config: XmssConfig) -> PublicKey:
         root = Fp.deserialize_list(data[:root_len], config.HASH_LEN_FE)
         parameter = Fp.deserialize_list(data[root_len:], config.PARAMETER_LEN)
 
-        return cls(root=root, parameter=parameter)
+        return cls(root=HashDigestVector(data=root), parameter=parameter)
 
 
 class Signature(StrictBaseModel):

src/lean_spec/subspecs/xmss/interface.py

@@ -191,7 +191,7 @@ def key_gen(
         )
 
         # Collect roots for building the top tree.
-        bottom_tree_roots: List[List[Fp]] = [
+        bottom_tree_roots: List[HashDigestVector] = [
             left_bottom_tree.root(),
             right_bottom_tree.root(),
         ]

src/lean_spec/subspecs/xmss/merkle_tree.py

@@ -168,17 +168,15 @@ def build(
         # A full tree is represented as a HashSubTree with lowest_layer=0
         return HashSubTree(depth=Uint64(depth), lowest_layer=Uint64(0), layers=layers)
 
-    def root(self, tree: HashSubTree) -> List[Fp]:
+    def root(self, tree: HashSubTree) -> HashDigestVector:
         """
         Extracts the root digest from a constructed Merkle tree.
 
         The root is the single node in the final, highest layer of the `HashSubTree`
         and serves as the primary component of the master public key.
         """
         # The root is the single node in the final layer.
-        root_node = cast(HashDigestVector, tree.layers[-1].nodes.data[0])
-        root_data = cast("Tuple[Fp, ...]", root_node.data)
-        return list(root_data)
+        return cast(HashDigestVector, tree.layers[-1].nodes[0])
 
     def path(self, tree: HashSubTree, position: Uint64) -> HashTreeOpening:
         """
@@ -236,7 +234,7 @@ def path(self, tree: HashSubTree, position: Uint64) -> HashTreeOpening:
     def verify_path(
         self,
         parameter: Parameter,
-        root: List[Fp],
+        root: HashDigestVector,
         position: Uint64,
         leaf_parts: List[List[Fp]],
         opening: HashTreeOpening,
@@ -316,7 +314,7 @@ def verify_path(
 
         # After iterating through the entire path, the final computed node
         # should be the root of the tree.
-        return current_node == root
+        return current_node == list(root.data)
 
 
 PROD_MERKLE_TREE = MerkleTree(config=PROD_CONFIG, hasher=PROD_TWEAK_HASHER, rand=PROD_RAND)

src/lean_spec/subspecs/xmss/subtree.py

@@ -189,7 +189,7 @@ def new_top_tree(
         depth: int,
         start_bottom_tree_index: Uint64,
         parameter: Parameter,
-        bottom_tree_roots: List[List[Fp]],
+        bottom_tree_roots: List[HashDigestVector],
     ) -> HashSubTree:
         """
         Constructs a top tree from the roots of bottom trees.
@@ -232,6 +232,9 @@ def new_top_tree(
         # The top tree starts at the middle layer.
         lowest_layer = depth // 2
 
+        # Convert HashDigestVector roots to List[Fp] for building
+        roots_as_lists = [cast(List[Fp], list(root.data)) for root in bottom_tree_roots]
+
         # Build the top tree using the bottom tree roots as the lowest layer.
         return cls.new(
             hasher=hasher,
@@ -240,7 +243,7 @@ def new_top_tree(
             depth=depth,
             start_index=start_bottom_tree_index,
             parameter=parameter,
-            lowest_layer_nodes=bottom_tree_roots,
+            lowest_layer_nodes=roots_as_lists,
         )
 
     @classmethod
@@ -344,7 +347,7 @@ def new_bottom_tree(
 
         return cls(depth=Uint64(depth), lowest_layer=Uint64(0), layers=truncated_layers)
 
-    def root(self) -> List[Fp]:
+    def root(self) -> HashDigestVector:
         """
         Extracts the root digest from this subtree.
 
@@ -366,9 +369,7 @@ def root(self) -> List[Fp]:
 
         # The root is the only node in the highest layer for proper subtrees.
         # For top trees and bottom trees, the highest layer should have exactly one node.
-        root_node = cast(HashDigestVector, highest_layer.nodes.data[0])
-        root_data = cast("Tuple[Fp, ...]", root_node.data)
-        return list(root_data)
+        return cast(HashDigestVector, highest_layer.nodes[0])
 
     def path(self, position: Uint64) -> HashTreeOpening:
         """
@@ -400,7 +401,8 @@ def path(self, position: Uint64) -> HashTreeOpening:
         if position >= lowest_layer.start_index + Uint64(len(lowest_layer.nodes)):
             raise ValueError("Position is beyond the subtree's range.")
 
-        co_path: List[List[Fp]] = []
+        # Build the co-path directly with SSZ types
+        siblings = HashDigestList(data=[])
         current_position = position
 
         # Iterate through layers from lowest to highest, EXCLUDING the final root layer.
@@ -412,23 +414,20 @@ def path(self, position: Uint64) -> HashTreeOpening:
             sibling_index = sibling_position - layer.start_index
 
             # Ensure the sibling exists in this layer
-            if sibling_index < Uint64(0) or sibling_index >= Uint64(len(layer.nodes.data)):
+            if sibling_index < Uint64(0) or sibling_index >= Uint64(len(layer.nodes)):
                 raise ValueError(
                     f"Sibling index {sibling_index} out of bounds for layer "
-                    f"with {len(layer.nodes.data)} nodes"
+                    f"with {len(layer.nodes)} nodes"
                 )
 
-            # Add the sibling's hash to the co-path.
-            sibling_node = cast(HashDigestVector, layer.nodes.data[int(sibling_index)])
-            sibling_data = cast("Tuple[Fp, ...]", sibling_node.data)
-            co_path.append(list(sibling_data))
+            # Access the sibling directly from the SSZ list and add to path
+            siblings = siblings + [layer.nodes[int(sibling_index)]]
 
             # Move to the parent's position for the next iteration.
             current_position = current_position // Uint64(2)
 
-        # Wrap in SSZ types
-        ssz_siblings = [HashDigestVector(data=sibling) for sibling in co_path]
-        return HashTreeOpening(siblings=HashDigestList(data=ssz_siblings))
+        # Return the opening with SSZ-typed siblings
+        return HashTreeOpening(siblings=siblings)
 
 
 def combined_path(

tests/lean_spec/subspecs/xmss/test_containers.py

@@ -20,7 +20,7 @@ class TestPublicKey:
 
     def test_bytes_protocol(self) -> None:
         """Test that PublicKey implements Python's bytes protocol."""
-        root = [Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+        root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=i + 100) for i in range(TEST_CONFIG.PARAMETER_LEN)]
         pk = PublicKey(root=root, parameter=parameter)
 
@@ -31,18 +31,17 @@ def test_bytes_protocol(self) -> None:
 
     def test_to_bytes_with_validation(self) -> None:
         """Test that to_bytes validates field lengths."""
-        root = [Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+        root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=i) for i in range(TEST_CONFIG.PARAMETER_LEN)]
         pk = PublicKey(root=root, parameter=parameter)
 
         # Valid serialization
         data = pk.to_bytes(TEST_CONFIG)
         assert len(data) == TEST_CONFIG.PUBLIC_KEY_LEN_BYTES
 
-        # Invalid root length
-        invalid_pk = PublicKey(root=[Fp(value=0)] * 5, parameter=parameter)
-        with pytest.raises(ValueError, match="Invalid root length"):
-            invalid_pk.to_bytes(TEST_CONFIG)
+        # Invalid root length - HashDigestVector validates length at construction
+        with pytest.raises(ValueError, match="requires exactly"):
+            HashDigestVector(data=[Fp(value=0)] * 5)
 
         # Invalid parameter length
         invalid_pk = PublicKey(root=root, parameter=[Fp(value=0)] * 3)
@@ -51,7 +50,7 @@ def test_to_bytes_with_validation(self) -> None:
 
     def test_roundtrip_test_config(self) -> None:
         """Test serialization round-trip with TEST_CONFIG."""
-        root = [Fp(value=i * 10) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+        root = HashDigestVector(data=[Fp(value=i * 10) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=i * 20) for i in range(TEST_CONFIG.PARAMETER_LEN)]
         original = PublicKey(root=root, parameter=parameter)
 
@@ -65,7 +64,7 @@ def test_roundtrip_test_config(self) -> None:
 
     def test_roundtrip_prod_config(self) -> None:
         """Test serialization round-trip with PROD_CONFIG."""
-        root = [Fp(value=i) for i in range(PROD_CONFIG.HASH_LEN_FE)]
+        root = HashDigestVector(data=[Fp(value=i) for i in range(PROD_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=i + 1000) for i in range(PROD_CONFIG.PARAMETER_LEN)]
         original = PublicKey(root=root, parameter=parameter)
 
@@ -86,14 +85,16 @@ def test_from_bytes_invalid_length(self) -> None:
 
     def test_serialization_format(self) -> None:
         """Test that serialization follows the documented format: root || parameter."""
-        root = [Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+        root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=i + 100) for i in range(TEST_CONFIG.PARAMETER_LEN)]
         pk = PublicKey(root=root, parameter=parameter)
 
         data = bytes(pk)
 
         # Check that root comes first
-        root_data = Fp.serialize_list(root)
+        from typing import List, cast
+
+        root_data = Fp.serialize_list(cast(List[Fp], list(root.data)))
         parameter_data = Fp.serialize_list(parameter)
 
         assert data == root_data + parameter_data
@@ -253,7 +254,7 @@ class TestSerializationProperties:
 
     def test_public_key_deterministic(self) -> None:
         """Test that serialization is deterministic."""
-        root = [Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+        root = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=i) for i in range(TEST_CONFIG.PARAMETER_LEN)]
         pk = PublicKey(root=root, parameter=parameter)
 
@@ -288,8 +289,8 @@ def test_signature_deterministic(self) -> None:
 
     def test_different_values_produce_different_bytes(self) -> None:
         """Test that different values produce different serializations."""
-        root1 = [Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)]
-        root2 = [Fp(value=i + 1) for i in range(TEST_CONFIG.HASH_LEN_FE)]
+        root1 = HashDigestVector(data=[Fp(value=i) for i in range(TEST_CONFIG.HASH_LEN_FE)])
+        root2 = HashDigestVector(data=[Fp(value=i + 1) for i in range(TEST_CONFIG.HASH_LEN_FE)])
         parameter = [Fp(value=0)] * TEST_CONFIG.PARAMETER_LEN
 
         pk1 = PublicKey(root=root1, parameter=parameter)