subspec: add ssz subspec (leanEthereum#39)

* types: add Byte type to encapsulate Uint8

* subspec: add ssz subspec

Author: tcoratger

src/lean_spec/subspecs/ssz/constants.py

@@ -1,7 +1,9 @@
 """Constants defined in the SSZ specification."""
 
+from lean_spec.types.hash import Bytes32
+
 BYTES_PER_CHUNK: int = 32
 """The number of bytes in a Merkle tree chunk."""
 
-ZERO_HASH: bytes = b"\x00" * BYTES_PER_CHUNK
+ZERO_HASH: Bytes32 = b"\x00" * BYTES_PER_CHUNK
 """A zero hash, used for padding in the Merkle tree."""

src/lean_spec/subspecs/ssz/gindex.py

@@ -2,7 +2,7 @@
 
 from typing import List
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class GeneralizedIndex(BaseModel):
@@ -12,8 +12,14 @@ class GeneralizedIndex(BaseModel):
     Helper methods are provided for tree navigation.
     """
 
+    model_config = ConfigDict(strict=True)
+
     value: int = Field(..., gt=0, description="The index value, must be a positive integer.")
 
+    def __hash__(self) -> int:
+        """Hashes the index value."""
+        return hash(self.value)
+
     @property
     def depth(self) -> int:
         """The depth of the node in the tree."""
@@ -35,12 +41,26 @@ def parent(self) -> "GeneralizedIndex":
             raise ValueError("Root node has no parent.")
         return type(self)(value=self.value // 2)
 
+    def child(self, right_side: bool) -> "GeneralizedIndex":
+        """
+        Returns the index of a child node.
+
+        - `right_side=False` for the left child (2k),
+        - `right_side=True` for the right child (2k+1).
+        """
+        return type(self)(value=self.value * 2 + int(right_side))
+
     def get_branch_indices(self) -> List["GeneralizedIndex"]:
         """Gets the indices of the sibling nodes along the path to the root."""
-        indices = [self.sibling]
-        while indices[-1].value > 1:
-            indices.append(indices[-1].parent.sibling)
-        return indices[:-1]
+        indices: List["GeneralizedIndex"] = []
+        current_index = self.value
+        while current_index > 1:
+            # Sibling is the current index XOR 1
+            sibling_index = current_index ^ 1
+            indices.append(type(self)(value=sibling_index))
+            # Move up to the parent
+            current_index //= 2
+        return indices
 
     def get_path_indices(self) -> List["GeneralizedIndex"]:
         """Gets the indices of the nodes along the path to the root."""

src/lean_spec/subspecs/ssz/merkle/__init__.py

@@ -0,0 +1 @@
+"""SSZ Merkle related functionality."""

src/lean_spec/subspecs/ssz/merkle/proof.py

@@ -0,0 +1,137 @@
+"""Merkle proofs for SSZ."""
+
+from __future__ import annotations
+
+from typing import Dict, List, Sequence, Set
+
+from pydantic import Field, model_validator
+
+from lean_spec.types.base import StrictBaseModel
+from lean_spec.types.hash import Bytes32
+
+from ..constants import ZERO_HASH
+from ..gindex import GeneralizedIndex
+from ..utils import hash_nodes
+
+Root = bytes
+"""The type of a Merkle tree root."""
+Proof = Sequence[Bytes32]
+"""The type of a Merkle proof."""
+ProofHashes = Sequence[Bytes32]
+"""The type of a Merkle proof's helper nodes."""
+
+
+class MerkleProof(StrictBaseModel):
+    """
+    Represents a Merkle multiproof, encapsulating its data and verification logic.
+
+    This object is immutable; once created, its contents cannot be changed.
+    """
+
+    leaves: Sequence[Bytes32] = Field(..., description="The leaf data being proven.")
+
+    indices: Sequence[GeneralizedIndex] = Field(
+        ..., description="The generalized indices of the leaves."
+    )
+
+    proof_hashes: ProofHashes = Field(..., description="The helper nodes required for the proof.")
+
+    @model_validator(mode="after")
+    def check_leaves_and_indices_length(self) -> MerkleProof:
+        """Ensures the number of leaves matches the number of indices."""
+        if len(self.leaves) != len(self.indices):
+            raise ValueError("The number of leaves must match the number of indices.")
+        return self
+
+    @classmethod
+    def from_single_leaf(
+        cls, leaf: Bytes32, proof_hashes: ProofHashes, index: GeneralizedIndex
+    ) -> MerkleProof:
+        """Creates a MerkleProof object from a traditional single-item proof."""
+        return cls(leaves=[leaf], proof_hashes=proof_hashes, indices=[index])
+
+    def _get_helper_indices(self) -> List[GeneralizedIndex]:
+        """
+        Calculates the generalized indices of all "helper" nodes needed to prove the leaves.
+
+        This is an internal helper method.
+        """
+        all_helper_indices: Set[GeneralizedIndex] = set()
+        all_path_indices: Set[GeneralizedIndex] = set()
+
+        for index in self.indices:
+            all_helper_indices.update(index.get_branch_indices())
+            all_path_indices.update(index.get_path_indices())
+
+        return sorted(all_helper_indices - all_path_indices, key=lambda g: g.value, reverse=True)
+
+    def calculate_root(self) -> Root:
+        """
+        Calculates the Merkle root from the proof's leaves and helper nodes.
+
+        Handles both single and multi-leaf proofs seamlessly.
+        """
+        # For a single leaf proof, use the more direct calculation.
+        if len(self.indices) == 1:
+            index = self.indices[0]
+            leaf = self.leaves[0]
+            if len(self.proof_hashes) != index.depth:
+                raise ValueError("Proof length must match the depth of the index.")
+
+            root = leaf
+            for i, branch_node in enumerate(self.proof_hashes):
+                if index.get_bit(i):
+                    root = hash_nodes(branch_node, root)
+                else:
+                    root = hash_nodes(root, branch_node)
+            return root
+
+        # For multi-leaf proofs, perform tree reconstruction.
+        helper_indices = self._get_helper_indices()
+        if len(self.proof_hashes) != len(helper_indices):
+            raise ValueError("Proof length does not match the required number of helper nodes.")
+
+        # 1. Start with the known nodes (leaves and proof hashes).
+        tree: Dict[int, Bytes32] = {
+            **{index.value: node for index, node in zip(self.indices, self.leaves, strict=False)},
+            **{
+                index.value: node
+                for index, node in zip(helper_indices, self.proof_hashes, strict=False)
+            },
+        }
+
+        # 2. Process nodes from deepest to shallowest.
+        #    The list of keys will grow as we create new parent nodes.
+        keys = sorted(tree.keys(), reverse=True)
+        pos = 0
+        while pos < len(keys):
+            key = keys[pos]
+            sibling_key = key ^ 1
+
+            # 3. If a node's sibling is also in the tree, we can create their parent.
+            if sibling_key in tree:
+                parent_key = key // 2
+
+                # Ensure we don't re-calculate a parent we already have.
+                if parent_key not in tree:
+                    # The order of hashing depends on which key is smaller.
+                    if key < sibling_key:
+                        tree[parent_key] = hash_nodes(tree[key], tree[sibling_key])
+                    else:
+                        tree[parent_key] = hash_nodes(tree[sibling_key], tree[key])
+                    keys.append(parent_key)
+            pos += 1
+
+        # 4. After processing all nodes, the root must be at index 1.
+        if 1 not in tree:
+            # This can happen if the proof is incomplete or for an empty leaf set.
+            return ZERO_HASH
+
+        return tree[1]
+
+    def verify(self, root: Root) -> bool:
+        """Verifies the Merkle proof against a known root."""
+        try:
+            return self.calculate_root() == root
+        except ValueError:
+            return False

src/lean_spec/subspecs/ssz/merkle/tree.py

@@ -0,0 +1,43 @@
+"""Merkle tree building logic."""
+
+from typing import List, Sequence
+
+from lean_spec.subspecs.ssz.constants import ZERO_HASH
+from lean_spec.subspecs.ssz.utils import get_power_of_two_ceil, hash_nodes
+from lean_spec.types.hash import Bytes32
+
+
+def build_merkle_tree(leaves: Sequence[Bytes32]) -> List[Bytes32]:
+    r"""
+    Builds a full Merkle tree and returns it as a flat list.
+
+    The tree is represented as a list where the node at a generalized
+    index `i` is located at `tree[i]`. The 0-index is a placeholder.
+    """
+    # Handle the edge case of no leaves.
+    if not leaves:
+        # Per the spec, a tree of an empty list is a single ZERO_HASH.
+        #
+        # The flat list format includes a placeholder at index 0.
+        return [ZERO_HASH] * 2
+
+    # Calculate the required size of the bottom layer (must be a power of two).
+    bottom_layer_size = get_power_of_two_ceil(len(leaves))
+
+    # Create the complete, padded leaf layer.
+    padded_leaves = list(leaves) + [ZERO_HASH] * (bottom_layer_size - len(leaves))
+
+    # Initialize the tree with placeholders for parent nodes and the padded leaves.
+    #
+    # The first half of the list will store the calculated parent nodes.
+    tree = [ZERO_HASH] * bottom_layer_size + padded_leaves
+
+    # Iterate backwards from the last parent node up to the root.
+    #
+    # This calculates the tree from the bottom up.
+    for i in range(bottom_layer_size - 1, 0, -1):
+        # A parent at index `i` is the hash of its two children at `2*i` and `2*i+1`.
+        tree[i] = hash_nodes(tree[i * 2], tree[i * 2 + 1])
+
+    # Return the fully constructed tree.
+    return tree

src/lean_spec/subspecs/ssz/utils.py

@@ -0,0 +1,21 @@
+"""Generic helper functions for SSZ and Merkle proofs."""
+
+import hashlib
+
+from lean_spec.types.hash import Bytes32
+
+
+def get_power_of_two_ceil(x: int) -> int:
+    """
+    Calculates the smallest power of two greater than or equal to x.
+
+    Examples: 0->1, 1->1, 2->2, 3->4, 4->4, 5->8.
+    """
+    if x <= 1:
+        return 1
+    return 1 << (x - 1).bit_length()
+
+
+def hash_nodes(node_a: Bytes32, node_b: Bytes32) -> Bytes32:
+    """Hashes two 32-byte nodes together using SHA-256."""
+    return hashlib.sha256(node_a + node_b).digest()