types: cleanup (#244)

* types: cleanup

* fix tests

Author: tcoratger

packages/testing/src/consensus_testing/keys.py

@@ -44,8 +44,10 @@
     AttestationSignatures,
 )
 from lean_spec.subspecs.containers.slot import Slot
-from lean_spec.subspecs.containers.state.types import AttestationSignatureKey
-from lean_spec.subspecs.xmss.aggregation import MultisigAggregatedSignature
+from lean_spec.subspecs.xmss.aggregation import (
+    AttestationSignatureKey,
+    MultisigAggregatedSignature,
+)
 from lean_spec.subspecs.xmss.containers import KeyPair, PublicKey, Signature
 from lean_spec.subspecs.xmss.interface import (
     PROD_SIGNATURE_SCHEME,

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

@@ -26,10 +26,10 @@
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.containers.state import Validators
 from lean_spec.subspecs.containers.state.state import State
-from lean_spec.subspecs.containers.state.types import AttestationSignatureKey
 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.containers import Signature
 from lean_spec.subspecs.xmss.types import HashDigestList, HashTreeOpening, Randomness
 from lean_spec.types import Bytes32, Uint64

src/lean_spec/subspecs/containers/__init__.py

@@ -13,8 +13,7 @@
     AggregationBits,
     Attestation,
     AttestationData,
-    NaiveAggregatedSignature,
-    SignedAggregatedAttestation,
+    AttestationsByValidator,
     SignedAttestation,
 )
 from .block import (
@@ -31,19 +30,18 @@
 
 __all__ = [
     "AggregatedAttestation",
-    "NaiveAggregatedSignature",
     "AggregationBits",
-    "AttestationData",
     "Attestation",
-    "SignedAttestation",
-    "SignedAggregatedAttestation",
+    "AttestationData",
+    "AttestationsByValidator",
     "Block",
-    "BlockWithAttestation",
     "BlockBody",
     "BlockHeader",
+    "BlockWithAttestation",
     "Checkpoint",
     "Config",
+    "SignedAttestation",
     "SignedBlockWithAttestation",
-    "Validator",
     "State",
+    "Validator",
 ]

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

@@ -4,17 +4,15 @@
     AggregatedAttestation,
     Attestation,
     AttestationData,
-    SignedAggregatedAttestation,
     SignedAttestation,
 )
-from .types import AggregationBits, NaiveAggregatedSignature
+from .types import AggregationBits, AttestationsByValidator
 
 __all__ = [
-    "AttestationData",
-    "Attestation",
-    "SignedAttestation",
-    "SignedAggregatedAttestation",
     "AggregatedAttestation",
-    "NaiveAggregatedSignature",
     "AggregationBits",
+    "Attestation",
+    "AttestationData",
+    "AttestationsByValidator",
+    "SignedAttestation",
 ]

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

@@ -22,7 +22,7 @@
 
 from ...xmss.containers import Signature
 from ..checkpoint import Checkpoint
-from .types import AggregationBits, NaiveAggregatedSignature
+from .types import AggregationBits
 
 
 class AttestationData(Container):
@@ -107,22 +107,3 @@ def aggregate_by_data(
             )
             for data, validator_ids in data_to_validator_ids.items()
         ]
-
-
-class SignedAggregatedAttestation(Container):
-    """Aggregated attestation bundled with aggregated signatures."""
-
-    message: AggregatedAttestation
-    """Aggregated attestation data."""
-
-    signature: NaiveAggregatedSignature
-    """Aggregated attestation plus its combined signature.
-
-    Stores a naive list of validator signatures that mirrors the attestation
-    order.
-
-    TODO:
-    - signatures will be replaced by MegaBytes in next PR to include leanVM proof.
-    - this will be replaced by a SNARK in future devnets.
-    - this will be aggregated by aggregators in future devnets.
-    """

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

@@ -2,11 +2,18 @@
 
 from __future__ import annotations
 
-from lean_spec.types import SSZList, Uint64
+from typing import TYPE_CHECKING
+
+from lean_spec.types import Uint64
 from lean_spec.types.bitfields import BaseBitlist
 
 from ...chain.config import VALIDATOR_REGISTRY_LIMIT
-from ...xmss.containers import Signature
+
+if TYPE_CHECKING:
+    from .attestation import AttestationData
+
+AttestationsByValidator = dict[Uint64, "AttestationData"]
+"""Mapping from validator index to attestation data."""
 
 
 class AggregationBits(BaseBitlist):
@@ -57,10 +64,3 @@ def to_validator_indices(self) -> list[Uint64]:
             raise AssertionError("Aggregated attestation must reference at least one validator")
 
         return indices
-
-
-class NaiveAggregatedSignature(SSZList[Signature]):
-    """Naive list of validator signatures used for aggregation placeholders."""
-
-    ELEMENT_TYPE = Signature
-    LIMIT = int(VALIDATOR_REGISTRY_LIMIT)

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

@@ -11,12 +11,14 @@
 from .types import (
     AggregatedAttestations,
     AttestationSignatures,
+    BlockLookup,
 )
 
 __all__ = [
     "Block",
     "BlockBody",
     "BlockHeader",
+    "BlockLookup",
     "BlockSignatures",
     "BlockWithAttestation",
     "SignedBlockWithAttestation",

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

@@ -1,11 +1,21 @@
 """Block-specific SSZ types for the Lean Ethereum consensus specification."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from lean_spec.subspecs.xmss.aggregation import MultisigAggregatedSignature
-from lean_spec.types import SSZList
+from lean_spec.types import Bytes32, SSZList
 
 from ...chain.config import VALIDATOR_REGISTRY_LIMIT
 from ..attestation import AggregatedAttestation
 
+if TYPE_CHECKING:
+    from .block import Block
+
+BlockLookup = dict[Bytes32, "Block"]
+"""Mapping from block root to Block objects."""
+
 
 class AggregatedAttestations(SSZList[AggregatedAttestation]):
     """List of aggregated attestations included in a block."""

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

@@ -2,11 +2,6 @@
 
 from .state import State
 from .types import (
-    AggregatedSignaturePayload,
-    AggregatedSignaturePayloads,
-    AttestationsByValidator,
-    AttestationSignatureKey,
-    BlockLookup,
     HistoricalBlockHashes,
     JustificationRoots,
     JustificationValidators,
@@ -16,16 +11,11 @@
 )
 
 __all__ = [
-    "State",
-    "AggregatedSignaturePayload",
-    "AggregatedSignaturePayloads",
-    "AttestationSignatureKey",
-    "AttestationsByValidator",
-    "BlockLookup",
     "HistoricalBlockHashes",
     "JustificationRoots",
     "JustificationValidators",
     "JustifiedSlots",
+    "State",
     "StateLookup",
     "Validators",
 ]

src/lean_spec/subspecs/containers/state/state.py

@@ -3,7 +3,11 @@
 from typing import AbstractSet, Iterable
 
 from lean_spec.subspecs.ssz.hash import hash_tree_root
-from lean_spec.subspecs.xmss.aggregation import MultisigAggregatedSignature
+from lean_spec.subspecs.xmss.aggregation import (
+    AggregatedSignaturePayloads,
+    AttestationSignatureKey,
+    MultisigAggregatedSignature,
+)
 from lean_spec.subspecs.xmss.containers import PublicKey, Signature
 from lean_spec.types import (
     ZERO_HASH,
@@ -21,8 +25,6 @@
 from ..config import Config
 from ..slot import Slot
 from .types import (
-    AggregatedSignaturePayloads,
-    AttestationSignatureKey,
     HistoricalBlockHashes,
     JustificationRoots,
     JustificationValidators,