aggregation: cleanup and reorg (#245)

* aggregation: cleanup and reorg

* cleanup

* cleanup

* fix Bytes32

* fix linter

Author: tcoratger

packages/testing/src/consensus_testing/keys.py

@@ -39,14 +39,15 @@
 
 from lean_spec.config import LEAN_ENV
 from lean_spec.subspecs.containers import AttestationData
+from lean_spec.subspecs.containers.attestation import AggregationBits
 from lean_spec.subspecs.containers.block.types import (
     AggregatedAttestations,
     AttestationSignatures,
 )
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.xmss.aggregation import (
-    AttestationSignatureKey,
-    MultisigAggregatedSignature,
+    AggregatedSignatureProof,
+    SignatureKey,
 )
 from lean_spec.subspecs.xmss.containers import KeyPair, PublicKey, Signature
 from lean_spec.subspecs.xmss.interface import (
@@ -276,7 +277,7 @@ def sign_attestation_data(
     def build_attestation_signatures(
         self,
         aggregated_attestations: AggregatedAttestations,
-        signature_lookup: Mapping[AttestationSignatureKey, Signature] | None = None,
+        signature_lookup: Mapping[SignatureKey, Signature] | None = None,
     ) -> AttestationSignatures:
         """
         Build `AttestationSignatures` for already-aggregated attestations.
@@ -286,29 +287,34 @@ def build_attestation_signatures(
         """
         lookup = signature_lookup or {}
 
-        proof_blobs: list[MultisigAggregatedSignature] = []
+        proofs: list[AggregatedSignatureProof] = []
         for agg in aggregated_attestations:
             validator_ids = agg.aggregation_bits.to_validator_indices()
             message = agg.data.data_root_bytes()
             epoch = agg.data.slot
 
             public_keys: list[PublicKey] = [self.get_public_key(vid) for vid in validator_ids]
             signatures: list[Signature] = [
-                (lookup.get((vid, message)) or self.sign_attestation_data(vid, agg.data))
+                (
+                    lookup.get(SignatureKey(vid, message))
+                    or self.sign_attestation_data(vid, agg.data)
+                )
                 for vid in validator_ids
             ]
 
             # If the caller supplied raw signatures and any are invalid,
             # aggregation should fail with exception.
-            aggregated_signature = MultisigAggregatedSignature.aggregate_signatures(
+            participants = AggregationBits.from_validator_indices(validator_ids)
+            proof = AggregatedSignatureProof.aggregate(
+                participants=participants,
                 public_keys=public_keys,
                 signatures=signatures,
                 message=message,
                 epoch=epoch,
             )
-            proof_blobs.append(aggregated_signature)
+            proofs.append(proof)
 
-        return AttestationSignatures(data=proof_blobs)
+        return AttestationSignatures(data=proofs)
 
 
 def _generate_single_keypair(

packages/testing/src/consensus_testing/test_fixtures/fork_choice.py

@@ -29,7 +29,7 @@
 from lean_spec.subspecs.forkchoice import Store
 from lean_spec.subspecs.koalabear import Fp
 from lean_spec.subspecs.ssz import hash_tree_root
-from lean_spec.subspecs.xmss.aggregation import AttestationSignatureKey
+from lean_spec.subspecs.xmss.aggregation import SignatureKey
 from lean_spec.subspecs.xmss.containers import Signature
 from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
 from lean_spec.types import Bytes32, Uint64
@@ -404,14 +404,14 @@ def _build_attestations_from_spec(
         block_registry: dict[str, Block],
         parent_root: Bytes32,
         key_manager: XmssKeyManager,
-    ) -> tuple[list[Attestation], dict[AttestationSignatureKey, Signature]]:
+    ) -> tuple[list[Attestation], dict[SignatureKey, Signature]]:
         """Build attestations list from BlockSpec and their signatures."""
         if spec.attestations is None:
             return [], {}
 
         parent_state = store.states[parent_root]
         attestations = []
-        signature_lookup: dict[AttestationSignatureKey, Signature] = {}
+        signature_lookup: dict[SignatureKey, Signature] = {}
 
         for att_spec in spec.attestations:
             if isinstance(att_spec, SignedAttestationSpec):
@@ -423,9 +423,8 @@ def _build_attestations_from_spec(
 
             attestation = Attestation(validator_id=signed_att.validator_id, data=signed_att.message)
             attestations.append(attestation)
-            signature_lookup[(attestation.validator_id, attestation.data.data_root_bytes())] = (
-                signed_att.signature
-            )
+            sig_key = SignatureKey(attestation.validator_id, attestation.data.data_root_bytes())
+            signature_lookup[sig_key] = signed_att.signature
 
         return attestations, signature_lookup
 

packages/testing/src/consensus_testing/test_fixtures/state_transition.py

@@ -9,6 +9,7 @@
 from lean_spec.subspecs.containers.block.types import AggregatedAttestations
 from lean_spec.subspecs.containers.state.state import State
 from lean_spec.subspecs.ssz.hash import hash_tree_root
+from lean_spec.subspecs.xmss.aggregation import SignatureKey
 from lean_spec.types import Bytes32, Uint64
 
 from ..keys import get_shared_key_manager
@@ -262,7 +263,9 @@ def _build_block_from_spec(self, spec: BlockSpec, state: State) -> tuple[Block,
         if plain_attestations:
             key_manager = get_shared_key_manager(max_slot=spec.slot)
             gossip_signatures = {
-                (att.validator_id, att.data.data_root_bytes()): key_manager.sign_attestation_data(
+                SignatureKey(
+                    att.validator_id, att.data.data_root_bytes()
+                ): key_manager.sign_attestation_data(
                     att.validator_id,
                     att.data,
                 )

packages/testing/src/consensus_testing/test_fixtures/verify_signatures.py

@@ -21,6 +21,7 @@
 from lean_spec.subspecs.containers.state.state import State
 from lean_spec.subspecs.koalabear import Fp
 from lean_spec.subspecs.ssz import hash_tree_root
+from lean_spec.subspecs.xmss.aggregation import SignatureKey
 from lean_spec.subspecs.xmss.constants import TARGET_CONFIG
 from lean_spec.subspecs.xmss.containers import Signature
 from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
@@ -187,7 +188,7 @@ def _build_block_from_spec(
         # fixed-point collection when available_attestations/known_block_roots are used.
         # This might contain invalid signatures as we are not validating them here.
         gossip_signatures = {
-            (att.validator_id, att.data.data_root_bytes()): sig
+            SignatureKey(att.validator_id, att.data.data_root_bytes()): sig
             for att, sig in zip(attestations, attestation_signature_inputs, strict=True)
         }
 

src/lean_spec/subspecs/chain/config.py

@@ -7,7 +7,8 @@
 
 from typing_extensions import Final
 
-from lean_spec.types import StrictBaseModel, Uint64
+from lean_spec.types.base import StrictBaseModel
+from lean_spec.types.uint import Uint64
 
 # --- Time Parameters ---
 

src/lean_spec/subspecs/containers/__init__.py

@@ -10,10 +10,8 @@
 
 from .attestation import (
     AggregatedAttestation,
-    AggregationBits,
     Attestation,
     AttestationData,
-    AttestationsByValidator,
     SignedAttestation,
 )
 from .block import (
@@ -30,10 +28,8 @@
 
 __all__ = [
     "AggregatedAttestation",
-    "AggregationBits",
     "Attestation",
     "AttestationData",
-    "AttestationsByValidator",
     "Block",
     "BlockBody",
     "BlockHeader",

src/lean_spec/subspecs/containers/attestation/__init__.py

@@ -1,18 +1,17 @@
 """Attestation containers and related types for the Lean spec."""
 
+from .aggregation_bits import AggregationBits
 from .attestation import (
     AggregatedAttestation,
     Attestation,
     AttestationData,
     SignedAttestation,
 )
-from .types import AggregationBits, AttestationsByValidator
 
 __all__ = [
     "AggregatedAttestation",
     "AggregationBits",
     "Attestation",
     "AttestationData",
-    "AttestationsByValidator",
     "SignedAttestation",
 ]

…subspecs/containers/attestation/types.py …ntainers/attestation/aggregation_bits.pysrc/lean_spec/subspecs/containers/attestation/types.py renamed to src/lean_spec/subspecs/containers/attestation/aggregation_bits.py src/lean_spec/subspecs/containers/attestation/types.py renamed to src/lean_spec/subspecs/containers/attestation/aggregation_bits.py

@@ -1,23 +1,19 @@
-"""Attestation-related SSZ types for the Lean consensus specification."""
+"""Aggregation bits for tracking validator participation."""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-from lean_spec.types import Uint64
+from lean_spec.subspecs.chain.config import VALIDATOR_REGISTRY_LIMIT
+from lean_spec.types import Boolean, Uint64
 from lean_spec.types.bitfields import BaseBitlist
 
-from ...chain.config import VALIDATOR_REGISTRY_LIMIT
-
-if TYPE_CHECKING:
-    from .attestation import AttestationData
-
-AttestationsByValidator = dict[Uint64, "AttestationData"]
-"""Mapping from validator index to attestation data."""
-
 
 class AggregationBits(BaseBitlist):
-    """Bitlist representing validator participation in an attestation."""
+    """
+    Bitlist representing validator participation in an attestation or signature.
+
+    A general-purpose bitfield for tracking which validators have participated
+    in some collective action (attestation, signature aggregation, etc.).
+    """
 
     LIMIT = int(VALIDATOR_REGISTRY_LIMIT)
 
@@ -36,19 +32,23 @@ def from_validator_indices(cls, indices: list[Uint64]) -> AggregationBits:
             AssertionError: If no indices are provided.
             AssertionError: If any index is outside the supported LIMIT.
         """
-        ids = [int(i) for i in indices]
-        if not ids:
+        # Require at least one validator for a valid aggregation.
+        if not indices:
             raise AssertionError("Aggregated attestation must reference at least one validator")
 
-        max_id = max(ids)
-        if max_id >= cls.LIMIT:
-            raise AssertionError("Validator index out of range for aggregation bits")
+        # Convert to a set of native ints.
+        #
+        # This combines int conversion and deduplication in a single O(N) pass.
+        ids = {int(i) for i in indices}
 
-        bits = [False] * (max_id + 1)
-        for i in ids:
-            bits[i] = True
+        # Validate bounds: max index must be within registry limit.
+        if (max_id := max(ids)) >= cls.LIMIT:
+            raise AssertionError("Validator index out of range for aggregation bits")
 
-        return cls(data=bits)
+        # Build bit list:
+        # - True at positions present in indices,
+        # - False elsewhere.
+        return cls(data=[Boolean(i in ids) for i in range(max_id + 1)])
 
     def to_validator_indices(self) -> list[Uint64]:
         """
@@ -60,6 +60,7 @@ def to_validator_indices(self) -> list[Uint64]:
         Raises:
             AssertionError: If no bits are set.
         """
+        # Extract indices where bit is set; fail if none found.
         if not (indices := [Uint64(i) for i, bit in enumerate(self.data) if bool(bit)]):
             raise AssertionError("Aggregated attestation must reference at least one validator")
 

src/lean_spec/subspecs/containers/attestation/attestation.py

@@ -18,11 +18,11 @@
 
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.ssz import hash_tree_root
-from lean_spec.types import Container, Uint64
+from lean_spec.types import Bytes32, Container, Uint64
 
 from ...xmss.containers import Signature
 from ..checkpoint import Checkpoint
-from .types import AggregationBits
+from .aggregation_bits import AggregationBits
 
 
 class AttestationData(Container):
@@ -40,9 +40,9 @@ class AttestationData(Container):
     source: Checkpoint
     """The checkpoint representing the source block as observed by the validator."""
 
-    def data_root_bytes(self) -> bytes:
+    def data_root_bytes(self) -> Bytes32:
         """The root of the attestation data."""
-        return bytes(hash_tree_root(self))
+        return hash_tree_root(self)
 
 
 class Attestation(Container):

src/lean_spec/subspecs/containers/block/block.py

@@ -12,9 +12,7 @@
 from typing import TYPE_CHECKING
 
 from lean_spec.subspecs.containers.slot import Slot
-from lean_spec.subspecs.xmss.aggregation import (
-    MultisigError,
-)
+from lean_spec.subspecs.xmss.aggregation import AggregationError
 from lean_spec.subspecs.xmss.interface import TARGET_SIGNATURE_SCHEME, GeneralizedXmssScheme
 from lean_spec.types import Bytes32, Uint64
 from lean_spec.types.container import Container
@@ -189,12 +187,12 @@ def verify_signatures(
 
             public_keys = [validators[vid].get_pubkey() for vid in validator_ids]
             try:
-                aggregated_signature.verify_aggregated_payload(
+                aggregated_signature.verify(
                     public_keys=public_keys,
                     message=attestation_data_root,
                     epoch=aggregated_attestation.data.slot,
                 )
-            except MultisigError as exc:
+            except AggregationError as exc:
                 raise AssertionError(
                     f"Attestation aggregated signature verification failed: {exc}"
                 ) from exc