xmss: more ssz containers (#209)

* xmss: more ssz containers

* more ssz

Author: tcoratger

src/lean_spec/subspecs/xmss/containers.py

@@ -7,6 +7,7 @@
 from ...types import StrictBaseModel, Uint64
 from ...types.byte_arrays import BaseBytes
 from ...types.collections import SSZList, SSZVector
+from ...types.container import Container
 from ..koalabear import P_BYTES, Fp
 from .constants import PRF_KEY_LENGTH, PROD_CONFIG
 
@@ -90,14 +91,19 @@ class Parameter(SSZVector):
     LENGTH = PROD_CONFIG.PARAMETER_LEN
 
 
-Randomness = List[Fp]
-"""
-A type alias for the randomness `rho` (ρ) used during signing.
+class Randomness(SSZVector):
+    """
+    The randomness `rho` (ρ) used during signing.
 
-This value provides a variable input to the message hash, allowing the signer to
-repeatedly try hashing until a valid "codeword" is found. It must be included in
-the final signature for the verifier to reproduce the same hash.
-"""
+    This value provides a variable input to the message hash, allowing the signer to
+    repeatedly try hashing until a valid "codeword" is found. It must be included in
+    the final signature for the verifier to reproduce the same hash.
+
+    SSZ notation: `Vector[Fp, RAND_LEN_FE]`
+    """
+
+    ELEMENT_TYPE = Fp
+    LENGTH = PROD_CONFIG.RAND_LEN_FE
 
 
 def _serialize_digests(digests: HashDigestList) -> bytes:
@@ -155,7 +161,7 @@ class HashTreeOpening(StrictBaseModel):
     """SSZ-compliant list of sibling hashes, from bottom to top."""
 
 
-class HashTreeLayer(StrictBaseModel):
+class HashTreeLayer(Container):
     """
     Represents a single horizontal "slice" of the sparse Merkle tree.
 
@@ -169,6 +175,33 @@ class HashTreeLayer(StrictBaseModel):
     """SSZ-compliant list of hash digests stored for this layer."""
 
 
+LAYERS_LIMIT = PROD_CONFIG.LOG_LIFETIME + 1
+"""
+The maximum number of layers in a subtree.
+
+This is `LOG_LIFETIME + 1` to accommodate all layers from 0 (leaves) to LOG_LIFETIME (root),
+inclusive. For PROD_CONFIG with LOG_LIFETIME=32, this allows up to 33 layers.
+"""
+
+
+class HashTreeLayers(SSZList):
+    """
+    Variable-length list of Merkle tree layers.
+
+    In SSZ notation: `List[HashTreeLayer, LAYERS_LIMIT]`
+
+    This type represents the layers of a subtree, from the lowest layer up to the root.
+
+    The number of layers varies based on the subtree structure:
+    - Bottom trees: `LOG_LIFETIME/2` layers
+    - Top trees: `LOG_LIFETIME/2` layers
+    - Maximum: `LOG_LIFETIME + 1` layers
+    """
+
+    ELEMENT_TYPE = HashTreeLayer
+    LIMIT = LAYERS_LIMIT
+
+
 class PublicKey(StrictBaseModel):
     """
     The public-facing component of a key pair.
@@ -298,7 +331,7 @@ def __bytes__(self) -> bytes:
         """
         return (
             _serialize_digests(self.path.siblings)
-            + Fp.serialize_list(self.rho)
+            + Fp.serialize_list(cast(List[Fp], list(self.rho.data)))
             + _serialize_digests(self.hashes)
         )
 
@@ -399,7 +432,7 @@ def from_bytes(cls, data: bytes, config: XmssConfig) -> Signature:
 
         return cls(
             path=HashTreeOpening(siblings=siblings),
-            rho=rho,
+            rho=Randomness(data=rho),
             hashes=hashes,
         )
 

src/lean_spec/subspecs/xmss/merkle_tree.py

@@ -46,6 +46,7 @@
     HashDigestList,
     HashDigestVector,
     HashTreeLayer,
+    HashTreeLayers,
     HashTreeOpening,
     Parameter,
 )
@@ -166,7 +167,9 @@ def build(
 
         # Return the completed tree containing all computed layers.
         # A full tree is represented as a HashSubTree with lowest_layer=0
-        return HashSubTree(depth=Uint64(depth), lowest_layer=Uint64(0), layers=layers)
+        return HashSubTree(
+            depth=Uint64(depth), lowest_layer=Uint64(0), layers=HashTreeLayers(data=layers)
+        )
 
     def root(self, tree: HashSubTree) -> HashDigestVector:
         """
@@ -176,7 +179,7 @@ def root(self, tree: HashSubTree) -> HashDigestVector:
         and serves as the primary component of the master public key.
         """
         # The root is the single node in the final layer.
-        return cast(HashDigestVector, tree.layers[-1].nodes[0])
+        return cast(HashDigestVector, cast(HashTreeLayer, tree.layers.data[-1]).nodes[0])
 
     def path(self, tree: HashSubTree, position: Uint64) -> HashTreeOpening:
         """
@@ -204,10 +207,11 @@ def path(self, tree: HashSubTree, position: Uint64) -> HashTreeOpening:
             raise ValueError("Cannot generate path for empty tree.")
 
         # Check that the position is within the tree's range.
-        if position < tree.layers[0].start_index:
+        first_layer = cast(HashTreeLayer, tree.layers.data[0])
+        if position < first_layer.start_index:
             raise ValueError("Position (before start) is invalid.")
 
-        if position >= tree.layers[0].start_index + Uint64(len(tree.layers[0].nodes)):
+        if position >= first_layer.start_index + Uint64(len(first_layer.nodes)):
             raise ValueError("Position (after end) is invalid.")
 
         co_path: List[List[Fp]] = []
@@ -218,7 +222,7 @@ def path(self, tree: HashSubTree, position: Uint64) -> HashTreeOpening:
             # Determine the sibling's position by flipping the last bit (XOR with 1).
             sibling_position = current_position ^ Uint64(1)
             # Find the sibling's index within our sparsely stored `nodes` vector.
-            layer = tree.layers[level]
+            layer = cast(HashTreeLayer, tree.layers.data[level])
             sibling_index_in_vec = sibling_position - layer.start_index
             # Add the sibling's hash to the co-path.
             sibling_node = cast(HashDigestVector, layer.nodes.data[int(sibling_index_in_vec)])

src/lean_spec/subspecs/xmss/message_hash.py

@@ -191,7 +191,7 @@ def apply(
 
             # The input is: rho || P || epoch || message || iteration.
             combined_input = (
-                rho
+                cast(List[Fp], list(rho.data))
                 + cast(List[Fp], list(parameter.data))
                 + epoch_fe
                 + message_fe

src/lean_spec/subspecs/xmss/prf.py

@@ -24,7 +24,7 @@
     TEST_CONFIG,
     XmssConfig,
 )
-from .containers import PRFKey
+from .containers import PRFKey, Randomness
 
 PRF_DOMAIN_SEP: bytes = bytes(
     [
@@ -172,7 +172,7 @@ def apply(self, key: PRFKey, epoch: Uint64, chain_index: Uint64) -> List[Fp]:
 
     def get_randomness(
         self, key: PRFKey, epoch: Uint64, message: bytes, counter: Uint64
-    ) -> List[Fp]:
+    ) -> Randomness:
         """
         Derives pseudorandom field elements for use in deterministic signing.
 
@@ -200,7 +200,7 @@ def get_randomness(
             counter: The attempt number (used when retrying encoding).
 
         Returns:
-            A list of field elements to use as randomness for encoding (i.e., `rho`).
+            Randomness for encoding (i.e., `rho`).
         """
         config = self.config
 
@@ -224,16 +224,18 @@ def get_randomness(
         num_bytes_to_read = PRF_BYTES_PER_FE * config.RAND_LEN_FE
         prf_output_bytes = hashlib.shake_128(input_data).digest(num_bytes_to_read)
 
-        # Convert to field elements
-        return [
-            Fp(
-                value=int.from_bytes(
-                    prf_output_bytes[i * PRF_BYTES_PER_FE : (i + 1) * PRF_BYTES_PER_FE],
-                    "big",
+        # Convert to field elements and wrap in Randomness
+        return Randomness(
+            data=[
+                Fp(
+                    value=int.from_bytes(
+                        prf_output_bytes[i * PRF_BYTES_PER_FE : (i + 1) * PRF_BYTES_PER_FE],
+                        "big",
+                    )
                 )
-            )
-            for i in range(config.RAND_LEN_FE)
-        ]
+                for i in range(config.RAND_LEN_FE)
+            ]
+        )
 
 
 PROD_PRF = Prf(config=PROD_CONFIG)

src/lean_spec/subspecs/xmss/rand.py

@@ -40,7 +40,7 @@ def domain(self) -> List[Fp]:
 
     def rho(self) -> Randomness:
         """Generates randomness `rho` for message encoding."""
-        return self.field_elements(self.config.RAND_LEN_FE)
+        return Randomness(data=self.field_elements(self.config.RAND_LEN_FE))
 
 
 PROD_RAND = Rand(config=PROD_CONFIG)

src/lean_spec/subspecs/xmss/subtree.py

@@ -16,6 +16,7 @@
     HashDigestList,
     HashDigestVector,
     HashTreeLayer,
+    HashTreeLayers,
     HashTreeOpening,
     Parameter,
 )
@@ -72,10 +73,12 @@ class HashSubTree(StrictBaseModel):
     layers 16 through 32 (the root).
     """
 
-    layers: List[HashTreeLayer]
+    layers: HashTreeLayers
     """
     The layers of this subtree, from `lowest_layer` to the root.
 
+    SSZ notation: `List[HashTreeLayer, LAYERS_LIMIT]`
+
     - `layers[0]` corresponds to layer `lowest_layer` in the full tree
     - `layers[-1]` corresponds to the highest layer in this subtree
     - For bottom trees: the last layer contains a single root
@@ -179,7 +182,11 @@ def new(
             layers.append(current_layer)
 
         # Return the completed subtree.
-        return cls(depth=Uint64(depth), lowest_layer=Uint64(lowest_layer), layers=layers)
+        return cls(
+            depth=Uint64(depth),
+            lowest_layer=Uint64(lowest_layer),
+            layers=HashTreeLayers(data=layers),
+        )
 
     @classmethod
     def new_top_tree(
@@ -326,7 +333,7 @@ def new_bottom_tree(
         # the bottom_tree_index (if it's the left child of its parent in the top tree)
         # or bottom_tree_index (if it's the right child). Since we're at layer depth/2,
         # the position is simply bottom_tree_index.
-        middle_layer = full_tree.layers[depth // 2]
+        middle_layer = cast(HashTreeLayer, full_tree.layers.data[depth // 2])
 
         # The root is at position (start_index >> (depth // 2)) = bottom_tree_index
         # within the middle layer. We need to find it in the stored nodes.
@@ -336,7 +343,7 @@ def new_bottom_tree(
         root = list(root_data)
 
         # Truncate layers to keep only 0 through depth/2 - 1.
-        truncated_layers = full_tree.layers[: (depth // 2)]
+        truncated_layers = list(full_tree.layers.data[: (depth // 2)])
 
         # Add a final layer containing just the root.
         root_vector = HashDigestVector(data=root)
@@ -345,7 +352,11 @@ def new_bottom_tree(
         )
         truncated_layers.append(root_layer)
 
-        return cls(depth=Uint64(depth), lowest_layer=Uint64(0), layers=truncated_layers)
+        return cls(
+            depth=Uint64(depth),
+            lowest_layer=Uint64(0),
+            layers=HashTreeLayers(data=truncated_layers),
+        )
 
     def root(self) -> HashDigestVector:
         """
@@ -363,7 +374,7 @@ def root(self) -> HashDigestVector:
         if len(self.layers) == 0:
             raise ValueError("Cannot get root of empty subtree.")
 
-        highest_layer = self.layers[-1]
+        highest_layer = cast(HashTreeLayer, self.layers.data[-1])
         if len(highest_layer.nodes.data) == 0:
             raise ValueError("Highest layer of subtree is empty.")
 
@@ -394,7 +405,7 @@ def path(self, position: Uint64) -> HashTreeOpening:
         if len(self.layers) == 0:
             raise ValueError("Cannot generate path for empty subtree.")
 
-        lowest_layer = self.layers[0]
+        lowest_layer = cast(HashTreeLayer, self.layers.data[0])
         if position < lowest_layer.start_index:
             raise ValueError("Position is before the subtree's start index.")
 
@@ -407,8 +418,9 @@ def path(self, position: Uint64) -> HashTreeOpening:
 
         # Iterate through layers from lowest to highest, EXCLUDING the final root layer.
         # The root layer doesn't contribute a sibling to the authentication path.
-        # self.layers[:-1] gives all layers except the last (root) layer.
-        for layer in self.layers[:-1]:
+        # self.layers.data[:-1] gives all layers except the last (root) layer.
+        for layer_raw in self.layers.data[:-1]:
+            layer = cast(HashTreeLayer, layer_raw)
             # Determine the sibling's position by flipping the last bit.
             sibling_position = current_position ^ Uint64(1)
             sibling_index = sibling_position - layer.start_index
@@ -498,7 +510,7 @@ def combined_path(
     # Verify that the provided bottom_tree actually corresponds to this position.
     # The bottom tree's lowest layer starts at bottom_tree_index * leafs_per_bottom_tree.
     expected_start = bottom_tree_index * Uint64(leafs_per_bottom_tree)
-    actual_start = bottom_tree.layers[0].start_index
+    actual_start = cast(HashTreeLayer, bottom_tree.layers.data[0]).start_index
 
     if actual_start != expected_start:
         raise ValueError(