fp: make koalabear fp a SSZType (leanEthereum#199)

Author: tcoratger

src/lean_spec/subspecs/koalabear/field.py

@@ -1,10 +1,8 @@
 """Core definition of the KoalaBear prime field Fp."""
 
-from typing import Self
+from typing import IO, Self
 
-from pydantic import Field, field_validator
-
-from lean_spec.types import StrictBaseModel
+from lean_spec.types import SSZType
 
 # =================================================================
 # Field Constants
@@ -72,16 +70,61 @@
 # =================================================================
 
 
-class Fp(StrictBaseModel):
-    """An element in the KoalaBear prime field F_p."""
+class Fp(SSZType):
+    """
+    An element in the KoalaBear prime field F_p.
+
+    This is an SSZ-serializable type.
+
+    Each field element is represented as a 4-byte little-endian unsigned integer.
+    """
+
+    def __init__(self, value: int) -> None:
+        """
+        Create a field element.
+
+        Args:
+            value: The value to wrap. Must be in the range [0, P).
+
+            Negative values will be normalized to the range [0, P).
+
+        Raises:
+            TypeError: If value is not an integer.
+        """
+        if not isinstance(value, int):
+            raise TypeError(f"Field value must be an integer, got {type(value).__name__}")
+
+        # Normalize to [0, P) - handles negative values correctly
+        self.value: int = value % P
+
+    @classmethod
+    def is_fixed_size(cls) -> bool:
+        """Fp elements are fixed-size (4 bytes)."""
+        return True
+
+    @classmethod
+    def get_byte_length(cls) -> int:
+        """Get the byte length of an Fp element."""
+        return P_BYTES
 
-    value: int = Field(ge=0, lt=P, description="Field element value in the range [0, P)")
+    def serialize(self, stream: IO[bytes]) -> int:
+        """Serialize the field element to a binary stream."""
+        data = self.value.to_bytes(P_BYTES, byteorder="little")
+        stream.write(data)
+        return len(data)
 
-    @field_validator("value", mode="before")
     @classmethod
-    def reduce_modulo_p(cls, v: int) -> int:
-        """Reduces an integer input modulo P before validation."""
-        return v % P
+    def deserialize(cls, stream: IO[bytes], scope: int) -> Self:
+        """Deserialize a field element from a binary stream."""
+        if scope != P_BYTES:
+            raise ValueError(f"Expected {P_BYTES} bytes for Fp, got {scope}")
+        data = stream.read(P_BYTES)
+        if len(data) != P_BYTES:
+            raise ValueError(f"Expected {P_BYTES} bytes for Fp, got {len(data)}")
+        value = int.from_bytes(data, byteorder="little")
+        if value >= P:
+            raise ValueError(f"Value {value} exceeds field modulus {P}")
+        return cls(value=value)
 
     def __add__(self, other: Self) -> Self:
         """Field addition."""
@@ -135,6 +178,20 @@ def two_adic_generator(cls, bits: int) -> Self:
             raise ValueError(f"bits must be between 0 and {TWO_ADICITY}")
         return cls(value=TWO_ADIC_GENERATORS[bits])
 
+    def __eq__(self, other: object) -> bool:
+        """Check equality of two field elements."""
+        if not isinstance(other, Fp):
+            return False
+        return self.value == other.value
+
+    def __hash__(self) -> int:
+        """Compute hash of the field element."""
+        return hash(self.value)
+
+    def __repr__(self) -> str:
+        """String representation."""
+        return f"Fp(value={self.value})"
+
     def __bytes__(self) -> bytes:
         """
         Serialize the field element using Python's bytes protocol.
@@ -150,7 +207,7 @@ def __bytes__(self) -> bytes:
             >>> len(data) == 4
             True
         """
-        return self.value.to_bytes(P_BYTES, byteorder="little")
+        return self.encode_bytes()
 
     @classmethod
     def from_bytes(cls, data: bytes) -> Self:
@@ -175,15 +232,7 @@ def from_bytes(cls, data: bytes) -> Self:
             >>> recovered == fp
             True
         """
-        if len(data) != P_BYTES:
-            raise ValueError(f"Expected {P_BYTES} bytes, got {len(data)}")
-
-        value = int.from_bytes(data, byteorder="little")
-
-        if value >= P:
-            raise ValueError(f"Value {value} (0x{value:08x}) exceeds field modulus {P} (0x{P:08x})")
-
-        return cls(value=value)
+        return cls.decode_bytes(data)
 
     @classmethod
     def serialize_list(cls, elements: list[Self]) -> bytes:

src/lean_spec/types/__init__.py

@@ -6,6 +6,7 @@
 from .byte_arrays import ZERO_HASH, Bytes32, Bytes52, Bytes3116
 from .collections import SSZList, SSZVector
 from .container import Container
+from .ssz_base import SSZType
 from .uint import Uint64
 from .validator import ValidatorIndex, is_proposer
 
@@ -22,6 +23,7 @@
     "is_proposer",
     "SSZList",
     "SSZVector",
+    "SSZType",
     "Boolean",
     "Container",
 ]

tests/lean_spec/subspecs/koalabear/test_field.py

@@ -38,8 +38,8 @@ def test_base_field_arithmetic() -> None:
     # Test equality against the same and different types
     assert a == Fp(value=5)
     assert a != b
-    assert a != 5  # type: ignore[comparison-overlap]
-    assert a != "5"  # type: ignore[comparison-overlap]
+    assert a != 5
+    assert a != "5"
 
     # Test error on inverting the zero element
     with pytest.raises(ZeroDivisionError, match="Cannot invert the zero element."):
@@ -91,10 +91,10 @@ def test_bytes_protocol() -> None:
         assert Fp.from_bytes(bytes(fp)) == fp
 
     # Test error handling for invalid data length
-    with pytest.raises(ValueError, match="Expected 4 bytes, got 3"):
+    with pytest.raises(ValueError, match="Expected 4 bytes for Fp, got 3"):
         Fp.from_bytes(b"\x01\x02\x03")
 
-    with pytest.raises(ValueError, match="Expected 4 bytes, got 5"):
+    with pytest.raises(ValueError, match="Expected 4 bytes for Fp, got 5"):
         Fp.from_bytes(b"\x01\x02\x03\x04\x05")
 
     # Test error handling for values exceeding the modulus
@@ -172,3 +172,125 @@ def test_serialize_list_roundtrip_property() -> None:
 
         assert recovered == elements
         assert len(data) == count * 4
+
+
+def test_ssz_type_properties() -> None:
+    """Test that Fp correctly implements SSZ type interface."""
+    # Test is_fixed_size
+    assert Fp.is_fixed_size() is True
+
+    # Test get_byte_length
+    assert Fp.get_byte_length() == 4
+
+
+def test_ssz_serialize() -> None:
+    """Test SSZ serialization using the serialize method."""
+    import io
+
+    fp = Fp(value=42)
+
+    # Test serialize to stream
+    stream = io.BytesIO()
+    bytes_written = fp.serialize(stream)
+    assert bytes_written == 4
+    assert stream.getvalue() == b"\x2a\x00\x00\x00"  # 42 in LE
+
+
+def test_ssz_deserialize() -> None:
+    """Test SSZ deserialization using the deserialize method."""
+    import io
+
+    # Test successful deserialization
+    data = b"\x2a\x00\x00\x00"  # 42 in LE
+    stream = io.BytesIO(data)
+    fp = Fp.deserialize(stream, 4)
+    assert fp == Fp(value=42)
+
+
+def test_ssz_deserialize_wrong_scope() -> None:
+    """Test deserialize error when scope doesn't match P_BYTES."""
+    import io
+
+    data = b"\x2a\x00\x00\x00"
+    stream = io.BytesIO(data)
+    with pytest.raises(ValueError, match="Expected 4 bytes for Fp, got 3"):
+        Fp.deserialize(stream, 3)
+
+
+def test_ssz_deserialize_short_data() -> None:
+    """Test deserialize error when stream has insufficient data."""
+    import io
+
+    stream = io.BytesIO(b"\x01\x02\x03")  # Only 3 bytes
+    with pytest.raises(ValueError, match="Expected 4 bytes for Fp, got 3"):
+        Fp.deserialize(stream, 4)
+
+
+def test_ssz_deserialize_exceeds_modulus() -> None:
+    """Test deserialize error when value exceeds field modulus."""
+    import io
+
+    # P = 2^31 - 2^24 + 1 = 2130706433
+    # Encode a value >= P (use P itself)
+    invalid_data = P.to_bytes(4, byteorder="little")
+    stream = io.BytesIO(invalid_data)
+    with pytest.raises(ValueError, match="exceeds field modulus"):
+        Fp.deserialize(stream, 4)
+
+
+def test_ssz_encode_decode_bytes() -> None:
+    """Test SSZ encode_bytes and decode_bytes methods."""
+    # Test encode_bytes
+    fp = Fp(value=100)
+    data = fp.encode_bytes()
+    assert len(data) == 4
+    assert data == b"\x64\x00\x00\x00"  # 100 in LE
+
+    # Test decode_bytes
+    fp2 = Fp.decode_bytes(data)
+    assert fp2 == fp
+
+    # Test roundtrip for various values
+    test_values = [0, 1, 42, 255, 256, 1000, 65535, 65536, 1000000, P - 1]
+    for value in test_values:
+        fp = Fp(value=value)
+        data = fp.encode_bytes()
+        recovered = Fp.decode_bytes(data)
+        assert recovered == fp, f"Failed for value={value}"
+
+
+def test_ssz_roundtrip() -> None:
+    """Comprehensive SSZ roundtrip test with many values."""
+    import random
+
+    random.seed(12345)
+
+    for _ in range(100):
+        # Test with random values
+        value = random.randint(0, P - 1)
+        fp = Fp(value=value)
+
+        # Test all serialization methods give same result
+        data1 = bytes(fp)
+        data2 = fp.encode_bytes()
+        assert data1 == data2
+
+        # Test all deserialization methods work
+        recovered1 = Fp.from_bytes(data1)
+        recovered2 = Fp.decode_bytes(data2)
+        assert recovered1 == fp
+        assert recovered2 == fp
+
+
+def test_ssz_deterministic() -> None:
+    """Test that SSZ serialization is deterministic."""
+    fp = Fp(value=999)
+
+    # Serialize multiple times
+    data1 = fp.encode_bytes()
+    data2 = fp.encode_bytes()
+    data3 = bytes(fp)
+
+    # All should be identical
+    assert data1 == data2
+    assert data1 == data3