types: implement SSZType for boolean and bitfiels (leanEthereum#43)

Author: tcoratger

src/lean_spec/types/bitfields.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from typing import (
+    IO,
     Any,
     ClassVar,
     Dict,
@@ -18,12 +19,13 @@
 from typing_extensions import List, Self
 
 from .boolean import Boolean
+from .ssz_base import SSZType
 
 _BITVECTOR_CACHE: Dict[Tuple[Type[Any], int], Type[Bitvector]] = {}
 """A cache to store and reuse dynamically generated Bitvector types."""
 
 
-class Bitvector(tuple[Boolean, ...]):
+class Bitvector(tuple[Boolean, ...], SSZType):
     """A strict Bitvector type: a fixed-length, immutable sequence of booleans."""
 
     LENGTH: ClassVar[int]
@@ -142,6 +144,37 @@ def __get_pydantic_core_schema__(
             ),
         )
 
+    @classmethod
+    def is_fixed_size(cls) -> bool:
+        """Return whether the type is fixed-size."""
+        return True
+
+    @classmethod
+    def get_byte_length(cls) -> int:
+        """Return the byte length of the type."""
+        if not hasattr(cls, "LENGTH"):
+            raise TypeError("Cannot get length of raw Bitvector type.")
+        return (cls.LENGTH + 7) // 8
+
+    def serialize(self, stream: IO[bytes]) -> int:
+        """Serialize the bitvector to a binary stream."""
+        encoded_data = self.encode_bytes()
+        stream.write(encoded_data)
+        return len(encoded_data)
+
+    @classmethod
+    def deserialize(cls, stream: IO[bytes], scope: int) -> Self:
+        """Deserialize a bitvector from a binary stream."""
+        byte_length = cls.get_byte_length()
+        if scope != byte_length:
+            raise ValueError(
+                f"Invalid scope for {cls.__name__}: expected {byte_length}, got {scope}"
+            )
+        data = stream.read(byte_length)
+        if len(data) != byte_length:
+            raise IOError(f"Stream ended prematurely while decoding {cls.__name__}")
+        return cls.decode_bytes(data)
+
     def encode_bytes(self) -> bytes:
         """Serializes the Bitvector into a byte string according to SSZ spec."""
         # Calculate the number of bytes required to hold all bits.
@@ -192,7 +225,7 @@ def __repr__(self) -> str:
 """A cache to store and reuse dynamically generated Bitlist types."""
 
 
-class Bitlist(list[Boolean]):
+class Bitlist(list[Boolean], SSZType):
     """
     A strict Bitlist type: a variable-length, mutable sequence of booleans
     with a maximum capacity.
@@ -307,13 +340,38 @@ def __get_pydantic_core_schema__(
             ),
         )
 
+    @classmethod
+    def is_fixed_size(cls) -> bool:
+        """Return whether the type is fixed-size."""
+        return False
+
+    @classmethod
+    def get_byte_length(cls) -> int:
+        """Raise TypeError, as the type is variable-size."""
+        raise TypeError(f"Type {cls.__name__} is not fixed-size")
+
+    def serialize(self, stream: IO[bytes]) -> int:
+        """Serialize the bitlist to a binary stream."""
+        encoded_data = self.encode_bytes()
+        stream.write(encoded_data)
+        return len(encoded_data)
+
+    @classmethod
+    def deserialize(cls, stream: IO[bytes], scope: int) -> Self:
+        """Deserialize a bitlist from a binary stream."""
+        data = stream.read(scope)
+        if len(data) != scope:
+            raise IOError(f"Stream ended prematurely while decoding {cls.__name__}")
+        return cls.decode_bytes(data)
+
     def encode_bytes(self) -> bytes:
         """Serializes the Bitlist into a byte string with a trailing delimiter bit."""
         # Get the number of bits in the list.
         num_bits = len(self)
-        # The required byte length is the ceiling of (bits + 1) / 8.
-        byte_len = (num_bits + 8) // 8
-        # Create a mutable byte array.
+        if num_bits == 0:
+            return b"\x01"
+
+        byte_len = (num_bits + 7) // 8
         byte_array = bytearray(byte_len)
 
         # Pack the bits into the byte array.
@@ -323,12 +381,16 @@ def encode_bytes(self) -> bytes:
                 bit_index_in_byte = i % 8
                 byte_array[byte_index] |= 1 << bit_index_in_byte
 
-        # Add the mandatory delimiter bit at the position right after the last data bit.
-        delimiter_byte_index = num_bits // 8
-        delimiter_bit_index = num_bits % 8
-        byte_array[delimiter_byte_index] |= 1 << delimiter_bit_index
-
-        return bytes(byte_array)
+        # Add the mandatory delimiter bit.
+        if num_bits % 8 == 0:
+            # If the data perfectly fills the last byte, append a new byte for the delimiter.
+            return bytes(byte_array) + b"\x01"
+        else:
+            # Otherwise, set the bit after the last data bit in the existing last byte.
+            delimiter_byte_index = num_bits // 8
+            delimiter_bit_index = num_bits % 8
+            byte_array[delimiter_byte_index] |= 1 << delimiter_bit_index
+            return bytes(byte_array)
 
     @classmethod
     def decode_bytes(cls, data: bytes) -> Self:
@@ -338,7 +400,7 @@ def decode_bytes(cls, data: bytes) -> Self:
             raise TypeError("Cannot decode to raw Bitlist; specify a limit, e.g., `Bitlist[4]`.")
         # The encoded data must not be empty (it must at least contain the delimiter).
         if not data:
-            raise ValueError("Cannot decode empty bytes into a Bitlist.")
+            raise ValueError("Invalid Bitlist encoding: data cannot be empty.")
 
         # The length in bits is determined by finding the position of the delimiter.
         num_bits = (len(data) - 1) * 8

src/lean_spec/types/boolean.py

@@ -2,14 +2,16 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import IO, Any
 
 from pydantic.annotated_handlers import GetCoreSchemaHandler
 from pydantic_core import CoreSchema, core_schema
 from typing_extensions import Self
 
+from .ssz_base import SSZType
 
-class Boolean(int):
+
+class Boolean(int, SSZType):
     """
     A strict SSZ Boolean type that inherits from `int` for `True`/`False` representation.
 
@@ -70,6 +72,16 @@ def __get_pydantic_core_schema__(
             serialization=core_schema.plain_serializer_function_ser_schema(bool),
         )
 
+    @classmethod
+    def is_fixed_size(cls) -> bool:
+        """Return whether the type is fixed-size."""
+        return True
+
+    @classmethod
+    def get_byte_length(cls) -> int:
+        """Return the byte length of the type."""
+        return 1
+
     def encode_bytes(self) -> bytes:
         r"""
         Serializes the boolean to its SSZ byte representation.
@@ -78,6 +90,31 @@ def encode_bytes(self) -> bytes:
         """
         return b"\x01" if self else b"\x00"
 
+    @classmethod
+    def decode_bytes(cls, data: bytes) -> Self:
+        """Deserialize a single byte into a Boolean instance."""
+        if len(data) != 1:
+            raise ValueError(f"Expected 1 byte for Boolean, got {len(data)}")
+        if data[0] not in (0, 1):
+            raise ValueError(f"Boolean byte must be 0x00 or 0x01, got {data[0]:#04x}")
+        return cls(data[0])
+
+    def serialize(self, stream: IO[bytes]) -> int:
+        """Serialize the boolean to a binary stream."""
+        encoded_data = self.encode_bytes()
+        stream.write(encoded_data)
+        return len(encoded_data)
+
+    @classmethod
+    def deserialize(cls, stream: IO[bytes], scope: int) -> Self:
+        """Deserialize a boolean from a binary stream."""
+        if scope != 1:
+            raise ValueError(f"Invalid scope for Boolean: expected 1, got {scope}")
+        data = stream.read(1)
+        if len(data) != 1:
+            raise IOError("Stream ended prematurely while decoding Boolean")
+        return cls.decode_bytes(data)
+
     def _raise_type_error(self, other: Any, op_symbol: str) -> None:
         """Helper to raise a consistent TypeError for unsupported operations."""
         raise TypeError(

tests/lean_spec/types/test_bitfields.py

@@ -1,5 +1,6 @@
-"""Tests for the Bitvector and Bitlist types."""
+""" "Tests for the Bitvector and Bitlist types."""
 
+import io
 from typing import Any
 
 import pytest
@@ -201,28 +202,6 @@ class TestBitfieldSerialization:
             (4, (False, True, False, True), "0a"),
             (3, (False, True, False), "02"),
             (10, (1, False, True, False, False, False, True, True, False, True), "c502"),
-            (
-                16,
-                (
-                    True,
-                    False,
-                    True,
-                    False,
-                    False,
-                    False,
-                    True,
-                    True,
-                    False,
-                    True,
-                    False,
-                    False,
-                    False,
-                    False,
-                    True,
-                    True,
-                ),
-                "c5c2",
-            ),
         ],
     )
     def test_bitvector_serialization_deserialization(
@@ -274,7 +253,106 @@ def test_bitvector_decode_invalid_length(self) -> None:
     def test_bitlist_decode_invalid_data(self) -> None:
         """Tests that Bitlist.decode_bytes fails for invalid byte strings."""
         list_type = Bitlist[8]  # type: ignore
-        with pytest.raises(ValueError, match="Cannot decode empty bytes"):
+        with pytest.raises(ValueError, match="data cannot be empty"):
             list_type.decode_bytes(b"")
         with pytest.raises(ValueError, match="last byte cannot be zero"):
             list_type.decode_bytes(b"\xff\x00")
+
+
+class TestBitfieldSSZ:
+    """Tests the SSZType interface methods for bitfields."""
+
+    def test_bitvector_ssz_properties(self) -> None:
+        vec_type = Bitvector[10]  # type: ignore
+        assert vec_type.is_fixed_size() is True
+        assert vec_type.get_byte_length() == 2  # (10+7)//8
+
+    def test_bitlist_ssz_properties(self) -> None:
+        list_type = Bitlist[10]  # type: ignore
+        assert list_type.is_fixed_size() is False
+        with pytest.raises(TypeError):
+            list_type.get_byte_length()
+
+    def test_bitvector_deserialize_invalid_scope(self) -> None:
+        vec_type = Bitvector[8]  # type: ignore
+        stream = io.BytesIO(b"\xff")
+        with pytest.raises(ValueError, match="Invalid scope"):
+            vec_type.deserialize(stream, scope=2)
+
+    def test_bitvector_deserialize_premature_end(self) -> None:
+        vec_type = Bitvector[16]  # type: ignore
+        stream = io.BytesIO(b"\xff")  # Only 1 byte, expects 2
+        with pytest.raises(IOError, match="Stream ended prematurely"):
+            vec_type.deserialize(stream, scope=2)
+
+    def test_bitlist_deserialize_premature_end(self) -> None:
+        list_type = Bitlist[16]  # type: ignore
+        stream = io.BytesIO(b"\xff")  # Only 1 byte
+        with pytest.raises(IOError, match="Stream ended prematurely"):
+            list_type.deserialize(stream, scope=2)  # Scope says to read 2
+
+    @pytest.mark.parametrize(
+        "length,value,expected_hex",
+        [
+            (8, (1, 1, 0, 1, 0, 1, 0, 0), "2b"),
+            (4, (0, 1, 0, 1), "0a"),
+            (3, (0, 1, 0), "02"),
+            (10, (1, 0, 1, 0, 0, 0, 1, 1, 0, 1), "c502"),
+            (16, (1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1), "c5c2"),
+            (512, tuple([1] * 512), "ff" * 64),
+            (513, tuple([1] * 513), ("ff" * 64) + "01"),
+        ],
+    )
+    def test_bitvector_encode_decode(
+        self, length: int, value: Tuple[int, ...], expected_hex: str
+    ) -> None:
+        vec_t = Bitvector[length]  # type: ignore
+        instance = vec_t(value)
+        encoded = instance.encode_bytes()
+        assert encoded.hex() == expected_hex
+
+        # round-trip via classmethod
+        decoded = vec_t.decode_bytes(encoded)
+        assert decoded == instance
+
+        # round-trip via stream serialize/deserialize
+        stream = io.BytesIO()
+        written = instance.serialize(stream)
+        assert written == vec_t.get_byte_length()
+        stream.seek(0)
+        decoded2 = vec_t.deserialize(stream, scope=written)
+        assert decoded2 == instance
+
+    @pytest.mark.parametrize(
+        "limit,value,expected_hex",
+        [
+            (8, (), "01"),
+            (8, (1, 1, 0, 1, 0, 1, 0, 0), "2b01"),
+            (4, (0, 1, 0, 1), "1a"),
+            (3, (0, 1, 0), "0a"),
+            (16, (1, 0, 1, 0, 0, 0, 1, 1, 0, 1), "c506"),
+            (512, (1,), "03"),
+            (512, tuple([1] * 512), ("ff" * 64) + "01"),
+            (513, tuple([1] * 513), ("ff" * 64) + "03"),
+        ],
+    )
+    def test_bitlist_encode_decode(
+        self, limit: int, value: Tuple[int, ...], expected_hex: str
+    ) -> None:
+        list_t = Bitlist[limit]  # type: ignore
+        instance = list_t(value)
+        encoded = instance.encode_bytes()
+        assert encoded.hex() == expected_hex
+
+        # round-trip via classmethod
+        decoded = list_t.decode_bytes(encoded)
+        assert decoded == instance
+
+        # round-trip via stream serialize/deserialize
+        stream = io.BytesIO()
+        written = instance.serialize(stream)
+        # variable-size, so we assert the written size matches the encoding length
+        assert written == len(encoded)
+        stream.seek(0)
+        decoded2 = list_t.deserialize(stream, scope=written)
+        assert decoded2 == instance