xmss: make a few types Uint64 to prepare for ssz (leanEthereum#200)

* xmss: make a few types Uint64 to prepare for ssz

* fix tests

* small fixes

Author: tcoratger

src/lean_spec/subspecs/xmss/containers.py

@@ -6,7 +6,7 @@
 
 from pydantic import Field
 
-from ...types import StrictBaseModel
+from ...types import StrictBaseModel, Uint64
 from ..koalabear import P_BYTES, Fp
 from .constants import PRF_KEY_LENGTH
 
@@ -105,7 +105,7 @@ class HashTreeLayer(StrictBaseModel):
     for the active range of leaves, not the entire conceptual layer.
     """
 
-    start_index: int
+    start_index: Uint64
     """The starting index of the first node in this layer."""
     nodes: List[HashDigest]
     """A list of the actual hash digests stored for this layer."""
@@ -120,7 +120,7 @@ class HashTree(StrictBaseModel):
     long lifetimes, prefer `HashSubTree` with the top-bottom tree approach.
     """
 
-    depth: int
+    depth: Uint64
     """The total depth of the tree (e.g., 32 for a 2^32 leaf space)."""
     layers: List[HashTreeLayer]
     """

src/lean_spec/subspecs/xmss/merkle_tree.py

@@ -81,7 +81,7 @@ def enforce_strict_types(self) -> "MerkleTree":
                 )
         return self
 
-    def _get_padded_layer(self, nodes: List[HashDigest], start_index: int) -> HashTreeLayer:
+    def _get_padded_layer(self, nodes: List[HashDigest], start_index: Uint64) -> HashTreeLayer:
         """
         Pads a layer of nodes with random hashes to simplify tree construction.
 
@@ -99,21 +99,21 @@ def _get_padded_layer(self, nodes: List[HashDigest], start_index: int) -> HashTr
             A new `HashTreeLayer` with the necessary padding applied.
         """
         nodes_with_padding: List[HashDigest] = []
-        end_index = start_index + len(nodes) - 1
+        end_index = start_index + Uint64(len(nodes)) - Uint64(1)
 
         # Prepend random padding if the layer starts at an odd index.
-        if start_index % 2 == 1:
+        if start_index % Uint64(2) == Uint64(1):
             nodes_with_padding.append(self.rand.domain())
 
         # The actual start index of the padded layer is always the even
         # number at or immediately before the original start_index.
-        actual_start_index = start_index - (start_index % 2)
+        actual_start_index = start_index - (start_index % Uint64(2))
 
         # Add the actual node content.
         nodes_with_padding.extend(nodes)
 
         # Append random padding if the layer ends at an even index.
-        if end_index % 2 == 0:
+        if end_index % Uint64(2) == Uint64(0):
             nodes_with_padding.append(self.rand.domain())
 
         return HashTreeLayer(start_index=actual_start_index, nodes=nodes_with_padding)
@@ -162,7 +162,7 @@ def build(
 
         # Start with the leaf hashes and apply the initial padding.
         layers: List[HashTreeLayer] = []
-        current_layer = self._get_padded_layer(leaf_hashes, int(start_index))
+        current_layer = self._get_padded_layer(leaf_hashes, start_index)
         layers.append(current_layer)
 
         # Iterate from the leaf layer (level 0) up to the root.
@@ -179,20 +179,20 @@ def build(
                 )
             ):
                 # Calculate the position of the parent node in the next level up.
-                parent_index = (current_layer.start_index // 2) + i
+                parent_index = (current_layer.start_index // Uint64(2)) + Uint64(i)
                 # Create the tweak for hashing these two children.
                 tweak = TreeTweak(level=level + 1, index=parent_index)
                 # Hash the left and right children to get their parent.
                 parent_node = self.hasher.apply(parameter, tweak, list(children))
                 parents.append(parent_node)
 
             # Pad the new list of parents to prepare for the next iteration.
-            new_start_index = current_layer.start_index // 2
+            new_start_index = current_layer.start_index // Uint64(2)
             current_layer = self._get_padded_layer(parents, new_start_index)
             layers.append(current_layer)
 
         # Return the completed tree containing all computed layers.
-        return HashTree(depth=depth, layers=layers)
+        return HashTree(depth=Uint64(depth), layers=layers)
 
     def root(self, tree: HashTree) -> HashDigest:
         """
@@ -230,26 +230,26 @@ def path(self, tree: HashTree, position: Uint64) -> HashTreeOpening:
             raise ValueError("Cannot generate path for empty tree.")
 
         # Check that the position is within the tree's range.
-        if int(position) < tree.layers[0].start_index:
+        if position < tree.layers[0].start_index:
             raise ValueError("Position (before start) is invalid.")
 
-        if int(position) >= tree.layers[0].start_index + len(tree.layers[0].nodes):
+        if position >= tree.layers[0].start_index + Uint64(len(tree.layers[0].nodes)):
             raise ValueError("Position (after end) is invalid.")
 
         co_path: List[HashDigest] = []
-        current_position = int(position)
+        current_position = position
 
         # Iterate from the leaf layer (level 0) up to the layer below the root.
-        for level in range(tree.depth):
+        for level in range(int(tree.depth)):
             # Determine the sibling's position by flipping the last bit (XOR with 1).
-            sibling_position = current_position ^ 1
+            sibling_position = current_position ^ Uint64(1)
             # Find the sibling's index within our sparsely stored `nodes` vector.
             layer = tree.layers[level]
             sibling_index_in_vec = sibling_position - layer.start_index
             # Add the sibling's hash to the co-path.
-            co_path.append(layer.nodes[sibling_index_in_vec])
+            co_path.append(layer.nodes[int(sibling_index_in_vec)])
             # Move up to the parent's position for the next iteration.
-            current_position //= 2
+            current_position = current_position // Uint64(2)
 
         return HashTreeOpening(siblings=co_path)
 

src/lean_spec/subspecs/xmss/subtree.py

@@ -19,7 +19,7 @@
     from .tweak_hash import TweakHasher
 
 
-def _get_padded_layer(rand: Rand, nodes: List[HashDigest], start_index: int) -> HashTreeLayer:
+def _get_padded_layer(rand: Rand, nodes: List[HashDigest], start_index: Uint64) -> HashTreeLayer:
     """
     Pads a layer of nodes with random hashes to simplify tree construction.
 
@@ -38,21 +38,21 @@ def _get_padded_layer(rand: Rand, nodes: List[HashDigest], start_index: int) ->
         A new `HashTreeLayer` with the necessary padding applied.
     """
     nodes_with_padding: List[HashDigest] = []
-    end_index = start_index + len(nodes) - 1
+    end_index = start_index + Uint64(len(nodes)) - Uint64(1)
 
     # Prepend random padding if the layer starts at an odd index.
-    if start_index % 2 == 1:
+    if start_index % Uint64(2) == Uint64(1):
         nodes_with_padding.append(rand.domain())
 
     # The actual start index of the padded layer is always the even
     # number at or immediately before the original start_index.
-    actual_start_index = start_index - (start_index % 2)
+    actual_start_index = start_index - (start_index % Uint64(2))
 
     # Add the actual node content.
     nodes_with_padding.extend(nodes)
 
     # Append random padding if the layer ends at an even index.
-    if end_index % 2 == 0:
+    if end_index % Uint64(2) == Uint64(0):
         nodes_with_padding.append(rand.domain())
 
     return HashTreeLayer(start_index=actual_start_index, nodes=nodes_with_padding)
@@ -84,15 +84,15 @@ class HashSubTree(StrictBaseModel):
     - Two consecutive bottom trees (sliding window)
     """
 
-    depth: int
+    depth: Uint64
     """
     The total depth of the full tree (e.g., 32 for a 2^32 leaf space).
 
     This represents the depth of the complete Merkle tree, not just this subtree.
     A subtree starting from layer `k` will have `depth - k` layers stored.
     """
 
-    lowest_layer: int
+    lowest_layer: Uint64
     """
     The lowest layer included in this subtree.
 
@@ -175,7 +175,7 @@ def new(
 
         # Start with the lowest layer nodes and apply initial padding.
         layers: List[HashTreeLayer] = []
-        current_layer = _get_padded_layer(rand, lowest_layer_nodes, start_index)
+        current_layer = _get_padded_layer(rand, lowest_layer_nodes, Uint64(start_index))
         layers.append(current_layer)
 
         # Build the tree layer by layer from lowest_layer up to the root.
@@ -192,20 +192,20 @@ def new(
                 )
             ):
                 # Calculate the position of the parent node in the next level up.
-                parent_index = (current_layer.start_index // 2) + i
+                parent_index = (current_layer.start_index // Uint64(2)) + Uint64(i)
                 # Create the tweak for hashing these two children.
                 tweak = TreeTweak(level=level + 1, index=parent_index)
                 # Hash the left and right children to get their parent.
                 parent_node = hasher.apply(parameter, tweak, list(children))
                 parents.append(parent_node)
 
             # Pad the new list of parents to prepare for the next iteration.
-            new_start_index = current_layer.start_index // 2
+            new_start_index = current_layer.start_index // Uint64(2)
             current_layer = _get_padded_layer(rand, parents, new_start_index)
             layers.append(current_layer)
 
         # Return the completed subtree.
-        return cls(depth=depth, lowest_layer=lowest_layer, layers=layers)
+        return cls(depth=Uint64(depth), lowest_layer=Uint64(lowest_layer), layers=layers)
 
     @classmethod
     def new_top_tree(
@@ -353,16 +353,16 @@ def new_bottom_tree(
 
         # 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.
-        root_position_in_layer = bottom_tree_index - middle_layer.start_index
-        root = middle_layer.nodes[root_position_in_layer]
+        root_position_in_layer = Uint64(bottom_tree_index) - middle_layer.start_index
+        root = middle_layer.nodes[int(root_position_in_layer)]
 
         # Truncate layers to keep only 0 through depth/2 - 1.
         truncated_layers = full_tree.layers[: (depth // 2)]
 
         # Add a final layer containing just the root.
-        truncated_layers.append(HashTreeLayer(start_index=bottom_tree_index, nodes=[root]))
+        truncated_layers.append(HashTreeLayer(start_index=Uint64(bottom_tree_index), nodes=[root]))
 
-        return cls(depth=depth, lowest_layer=0, layers=truncated_layers)
+        return cls(depth=Uint64(depth), lowest_layer=Uint64(0), layers=truncated_layers)
 
     def root(self) -> HashDigest:
         """
@@ -412,35 +412,35 @@ def path(self, position: Uint64) -> HashTreeOpening:
             raise ValueError("Cannot generate path for empty subtree.")
 
         lowest_layer = self.layers[0]
-        if int(position) < lowest_layer.start_index:
+        if position < lowest_layer.start_index:
             raise ValueError("Position is before the subtree's start index.")
 
-        if int(position) >= lowest_layer.start_index + len(lowest_layer.nodes):
+        if position >= lowest_layer.start_index + Uint64(len(lowest_layer.nodes)):
             raise ValueError("Position is beyond the subtree's range.")
 
         co_path: List[HashDigest] = []
-        current_position = int(position)
+        current_position = position
 
         # 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]:
             # Determine the sibling's position by flipping the last bit.
-            sibling_position = current_position ^ 1
+            sibling_position = current_position ^ Uint64(1)
             sibling_index = sibling_position - layer.start_index
 
             # Ensure the sibling exists in this layer
-            if sibling_index < 0 or sibling_index >= len(layer.nodes):
+            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)} nodes"
                 )
 
             # Add the sibling's hash to the co-path.
-            co_path.append(layer.nodes[sibling_index])
+            co_path.append(layer.nodes[int(sibling_index)])
 
             # Move to the parent's position for the next iteration.
-            current_position //= 2
+            current_position = current_position // Uint64(2)
 
         return HashTreeOpening(siblings=co_path)
 
@@ -496,23 +496,23 @@ def combined_path(
     depth = top_tree.depth
 
     # Validate even depth (required for top-bottom split).
-    if depth % 2 != 0:
+    if depth % Uint64(2) != Uint64(0):
         raise ValueError(
             f"Top-bottom tree traversal requires even depth, got {depth}. "
             f"Cannot split tree into equal top and bottom halves."
         )
 
     # Calculate parameters for bottom trees.
-    leafs_per_bottom_tree = 1 << (depth // 2)
+    leafs_per_bottom_tree = 1 << int(depth // Uint64(2))
 
     # Determine which bottom tree this position belongs to.
     #
     # Bottom tree index = floor(position / sqrt(LIFETIME))
-    bottom_tree_index = int(position) // leafs_per_bottom_tree
+    bottom_tree_index = position // Uint64(leafs_per_bottom_tree)
 
     # 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 * leafs_per_bottom_tree
+    expected_start = bottom_tree_index * Uint64(leafs_per_bottom_tree)
     actual_start = bottom_tree.layers[0].start_index
 
     if actual_start != expected_start:

tests/lean_spec/subspecs/xmss/test_utils.py

@@ -15,6 +15,7 @@
     expand_activation_time,
     int_to_base_p,
 )
+from lean_spec.types import Uint64
 
 
 @pytest.mark.parametrize(
@@ -129,8 +130,8 @@ def test_bottom_tree_from_prf_key() -> None:
     )
 
     # Verify structure
-    assert bottom_tree.depth == config.LOG_LIFETIME
-    assert bottom_tree.lowest_layer == 0
+    assert bottom_tree.depth == Uint64(config.LOG_LIFETIME)
+    assert bottom_tree.lowest_layer == Uint64(0)
     assert len(bottom_tree.layers) > 0
 
     # Verify the root layer has exactly one node