refactor(forks): make Store generic over StateT, BlockT (Stage 5 of #686)

packages/testing/src/consensus_testing/test_types/store_checks.py

@@ -3,7 +3,7 @@
 from typing import Literal
 
 from lean_spec.forks.lstar.containers import AttestationData
-from lean_spec.forks.lstar.containers.block.block import Block, BlockLookup
+from lean_spec.forks.lstar.containers.block.block import Block
 from lean_spec.forks.lstar.spec import LstarSpec
 from lean_spec.forks.lstar.store import Store
 from lean_spec.subspecs.ssz import hash_tree_root
@@ -12,7 +12,7 @@
 from .utils import resolve_block_root
 
 
-def _ancestor_set(blocks: BlockLookup, head: Bytes32) -> set[Bytes32]:
+def _ancestor_set(blocks: dict[Bytes32, Block], head: Bytes32) -> set[Bytes32]:
     """Walk parent links from head and collect every reachable block root."""
     seen: set[Bytes32] = set()
     root = head

src/lean_spec/forks/__init__.py

@@ -13,14 +13,17 @@
     SignedBlock,
     Validator,
 )
-from .lstar.containers.block import BlockLookup, BlockSignatures
+from .lstar.containers.block import BlockSignatures
 from .lstar.containers.block.types import AggregatedAttestations, AttestationSignatures
 from .lstar.containers.state import State, Validators
-from .lstar.spec import LstarSpec
-from .lstar.store import AttestationSignatureEntry, Store
+from .lstar.spec import LstarSpec, LstarStore
+from .lstar.store import AttestationSignatureEntry
 from .protocol import ForkProtocol, SpecStateType, SpecStoreType
 from .registry import ForkRegistry
 
+Store = LstarStore
+"""Public alias resolving to the concrete LstarStore until other forks land."""
+
 FORK_SEQUENCE: list[ForkProtocol] = [LstarSpec()]
 """Ordered oldest to newest. ForkRegistry enforces strictly increasing VERSION."""
 
@@ -37,14 +40,14 @@
     "Block",
     "BlockBody",
     "BlockHeader",
-    "BlockLookup",
     "BlockSignatures",
     "Config",
     "DEFAULT_REGISTRY",
     "FORK_SEQUENCE",
     "ForkProtocol",
     "ForkRegistry",
     "LstarSpec",
+    "LstarStore",
     "SignedAggregatedAttestation",
     "SignedAttestation",
     "SignedBlock",

src/lean_spec/forks/lstar/__init__.py

@@ -1,6 +1,7 @@
 """Lstar fork"""
 
 from .containers.state import State
-from .store import AttestationSignatureEntry, Store
+from .spec import LstarStore as Store
+from .store import AttestationSignatureEntry
 
 __all__ = ["AttestationSignatureEntry", "State", "Store"]

src/lean_spec/forks/lstar/containers/block/__init__.py

@@ -4,7 +4,6 @@
     Block,
     BlockBody,
     BlockHeader,
-    BlockLookup,
     BlockSignatures,
     SignedBlock,
 )
@@ -17,7 +16,6 @@
     "Block",
     "BlockBody",
     "BlockHeader",
-    "BlockLookup",
     "BlockSignatures",
     "SignedBlock",
     "AggregatedAttestations",

src/lean_spec/forks/lstar/containers/block/block.py

@@ -66,10 +66,6 @@ class Block(Container):
     """The block's payload."""
 
 
-BlockLookup = dict[Bytes32, Block]
-"""Mapping from block root to Block objects."""
-
-
 class BlockSignatures(Container):
     """Aggregated signature payload for a block."""
 

src/lean_spec/forks/lstar/spec.py

@@ -62,6 +62,9 @@
 from ..protocol import ForkProtocol, SpecBlockType, SpecStateType
 from .store import AttestationSignatureEntry, Store
 
+LstarStore = Store[State, Block]
+"""Concrete Store specialization owned by the lstar fork."""
+
 
 class LstarSpec(ForkProtocol):
     """Lstar fork."""
@@ -80,7 +83,7 @@ class LstarSpec(ForkProtocol):
     block_signatures_class: type[BlockSignatures] = BlockSignatures
     aggregated_attestations_class: type[AggregatedAttestations] = AggregatedAttestations
     attestation_signatures_class: type[AttestationSignatures] = AttestationSignatures
-    store_class: type[Store] = Store
+    store_class: type[Store[State, Block]] = LstarStore
 
     attestation_data_class: type[AttestationData] = AttestationData
     attestation_class: type[Attestation] = Attestation
@@ -878,7 +881,7 @@ def create_store(  # type: ignore[override]  # ty: ignore[invalid-method-overrid
         state: SpecStateType,
         anchor_block: SpecBlockType,
         validator_id: ValidatorIndex | None,
-    ) -> Store:
+    ) -> LstarStore:
         """Initialize a forkchoice store from an anchor state and block.
 
         The anchor block and state form the starting point for fork choice.
@@ -935,7 +938,7 @@ def create_store(  # type: ignore[override]  # ty: ignore[invalid-method-overrid
             validator_id=validator_id,
         )
 
-    def prune_stale_attestation_data(self, store: Store) -> Store:
+    def prune_stale_attestation_data(self, store: LstarStore) -> LstarStore:
         """Remove attestation data that can no longer influence fork choice.
 
         An attestation becomes stale when its target checkpoint falls at or before
@@ -972,7 +975,7 @@ def prune_stale_attestation_data(self, store: Store) -> Store:
             }
         )
 
-    def validate_attestation(self, store: Store, attestation_data: AttestationData) -> None:
+    def validate_attestation(self, store: LstarStore, attestation_data: AttestationData) -> None:
         """Validate incoming attestation before processing.
 
         Ensures the vote respects the basic laws of time and topology:
@@ -1026,11 +1029,11 @@ def validate_attestation(self, store: Store, attestation_data: AttestationData)
 
     def on_gossip_attestation(
         self,
-        store: Store,
+        store: LstarStore,
         signed_attestation: SignedAttestation,
         scheme: GeneralizedXmssScheme = TARGET_SIGNATURE_SCHEME,
         is_aggregator: bool = False,
-    ) -> Store:
+    ) -> LstarStore:
         """Process a signed attestation received via gossip network.
 
         This method:
@@ -1090,9 +1093,9 @@ def on_gossip_attestation(
 
     def on_gossip_aggregated_attestation(
         self,
-        store: Store,
+        store: LstarStore,
         signed_attestation: SignedAggregatedAttestation,
-    ) -> Store:
+    ) -> LstarStore:
         """Process a signed aggregated attestation received via aggregation topic.
 
         This method:
@@ -1155,10 +1158,10 @@ def on_gossip_aggregated_attestation(
 
     def on_block(
         self,
-        store: Store,
+        store: LstarStore,
         signed_block: SignedBlock,
         scheme: GeneralizedXmssScheme = TARGET_SIGNATURE_SCHEME,
-    ) -> Store:
+    ) -> LstarStore:
         """Process a new block and update the forkchoice state.
 
         This method integrates a block into the forkchoice store by:
@@ -1263,7 +1266,7 @@ def on_block(
 
     def extract_attestations_from_aggregated_payloads(
         self,
-        store: Store,
+        store: LstarStore,
         aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]],
     ) -> dict[ValidatorIndex, AttestationData]:
         """Extract attestations from aggregated payloads.
@@ -1281,7 +1284,7 @@ def extract_attestations_from_aggregated_payloads(
                         attestations[validator_id] = attestation_data
         return attestations
 
-    def compute_block_weights(self, store: Store) -> dict[Bytes32, int]:
+    def compute_block_weights(self, store: LstarStore) -> dict[Bytes32, int]:
         """Compute attestation-based weight for each block above the finalized slot.
 
         Walks backward from each validator's latest head vote, incrementing weight
@@ -1306,7 +1309,7 @@ def compute_block_weights(self, store: Store) -> dict[Bytes32, int]:
 
     def _compute_lmd_ghost_head(
         self,
-        store: Store,
+        store: LstarStore,
         start_root: Bytes32,
         attestations: dict[ValidatorIndex, AttestationData],
         min_score: int = 0,
@@ -1388,7 +1391,7 @@ def _compute_lmd_ghost_head(
 
         return head
 
-    def update_head(self, store: Store) -> Store:
+    def update_head(self, store: LstarStore) -> LstarStore:
         """Compute updated store with new canonical head.
 
         Selects the canonical chain head using:
@@ -1418,7 +1421,7 @@ def update_head(self, store: Store) -> Store:
             }
         )
 
-    def accept_new_attestations(self, store: Store) -> Store:
+    def accept_new_attestations(self, store: LstarStore) -> LstarStore:
         """Process pending aggregated payloads and update forkchoice head.
 
         Moves aggregated payloads from latest_new_aggregated_payloads to
@@ -1456,7 +1459,7 @@ def accept_new_attestations(self, store: Store) -> Store:
         # Update head with newly accepted aggregated payloads
         return self.update_head(store)
 
-    def update_safe_target(self, store: Store) -> Store:
+    def update_safe_target(self, store: LstarStore) -> LstarStore:
         """Compute the deepest block that has 2/3+ supermajority attestation weight.
 
         The safe target is the furthest-from-genesis block where enough validators
@@ -1523,7 +1526,7 @@ def update_safe_target(self, store: Store) -> Store:
         # The head and attestation pools remain unchanged.
         return store.model_copy(update={"safe_target": safe_target})
 
-    def aggregate(self, store: Store) -> tuple[Store, list[SignedAggregatedAttestation]]:
+    def aggregate(self, store: LstarStore) -> tuple[LstarStore, list[SignedAggregatedAttestation]]:
         """Turn raw validator votes into compact aggregated attestations.
 
         Validators cast individual signatures over gossip. Before those
@@ -1660,10 +1663,10 @@ def aggregate(self, store: Store) -> tuple[Store, list[SignedAggregatedAttestati
 
     def tick_interval(
         self,
-        store: Store,
+        store: LstarStore,
         has_proposal: bool,
         is_aggregator: bool = False,
-    ) -> tuple[Store, list[SignedAggregatedAttestation]]:
+    ) -> tuple[LstarStore, list[SignedAggregatedAttestation]]:
         """Advance store time by one interval and perform interval-specific actions.
 
         Different actions are performed based on interval within slot:
@@ -1691,11 +1694,11 @@ def tick_interval(
 
     def on_tick(
         self,
-        store: Store,
+        store: LstarStore,
         target_interval: Interval,
         has_proposal: bool,
         is_aggregator: bool = False,
-    ) -> tuple[Store, list[SignedAggregatedAttestation]]:
+    ) -> tuple[LstarStore, list[SignedAggregatedAttestation]]:
         """Advance forkchoice store time to given interval count.
 
         Ticks store forward interval by interval, performing appropriate
@@ -1716,7 +1719,7 @@ def on_tick(
 
         return store, all_new_aggregates
 
-    def get_proposal_head(self, store: Store, slot: Slot) -> tuple[Store, Bytes32]:
+    def get_proposal_head(self, store: LstarStore, slot: Slot) -> tuple[LstarStore, Bytes32]:
         """Get the head for block proposal at given slot.
 
         Ensures store is up-to-date and processes any pending attestations
@@ -1732,7 +1735,7 @@ def get_proposal_head(self, store: Store, slot: Slot) -> tuple[Store, Bytes32]:
 
         return store, store.head
 
-    def get_attestation_target(self, store: Store) -> Checkpoint:
+    def get_attestation_target(self, store: LstarStore) -> Checkpoint:
         """Calculate target checkpoint for validator attestations.
 
         Determines appropriate attestation target based on head, safe target,
@@ -1770,7 +1773,7 @@ def get_attestation_target(self, store: Store) -> Checkpoint:
 
         return Checkpoint(root=target_block_root, slot=target_block.slot)
 
-    def produce_attestation_data(self, store: Store, slot: Slot) -> AttestationData:
+    def produce_attestation_data(self, store: LstarStore, slot: Slot) -> AttestationData:
         """Produce attestation data for the given slot.
 
         This method constructs an AttestationData object according to the lean protocol
@@ -1796,10 +1799,10 @@ def produce_attestation_data(self, store: Store, slot: Slot) -> AttestationData:
 
     def produce_block_with_signatures(
         self,
-        store: Store,
+        store: LstarStore,
         slot: Slot,
         validator_index: ValidatorIndex,
-    ) -> tuple[Store, Block, list[AggregatedSignatureProof]]:
+    ) -> tuple[LstarStore, Block, list[AggregatedSignatureProof]]:
         """Produce a block and its aggregated signature proofs for the target slot.
 
         Block production proceeds in four stages:

src/lean_spec/forks/lstar/store.py

@@ -4,22 +4,28 @@
 The Store tracks all information required for the LMD GHOST forkchoice algorithm.
 """
 
-__all__ = ["AttestationSignatureEntry", "Store"]
+__all__ = ["AttestationSignatureEntry", "BlockT", "StateT", "Store"]
 
-from typing import NamedTuple
+from typing import Generic, NamedTuple, TypeVar
+
+from pydantic import Field
 
 from lean_spec.forks.lstar.containers import (
     AttestationData,
     Config,
 )
-from lean_spec.forks.lstar.containers.block import BlockLookup
 from lean_spec.subspecs.chain.clock import Interval
 from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof
 from lean_spec.subspecs.xmss.containers import Signature
 from lean_spec.types import Bytes32, Checkpoint, ValidatorIndex
 from lean_spec.types.base import StrictBaseModel
+from lean_spec.types.container import Container
+
+StateT = TypeVar("StateT", bound=Container)
+"""Per-fork post-state type tracked alongside each known block."""
 
-from .containers.state import State
+BlockT = TypeVar("BlockT", bound=Container)
+"""Per-fork block type stored in the forkchoice view."""
 
 
 class AttestationSignatureEntry(NamedTuple):
@@ -34,7 +40,7 @@ class AttestationSignatureEntry(NamedTuple):
     signature: Signature
 
 
-class Store(StrictBaseModel):
+class Store(StrictBaseModel, Generic[StateT, BlockT]):
     """
     Forkchoice store tracking chain state and validator attestations.
 
@@ -91,7 +97,7 @@ class Store(StrictBaseModel):
     Fork choice will never revert finalized history.
     """
 
-    blocks: BlockLookup = {}
+    blocks: dict[Bytes32, BlockT] = Field(default_factory=dict)
     """
     Mapping from block root to Block objects.
 
@@ -100,7 +106,7 @@ class Store(StrictBaseModel):
     Every block that might participate in fork choice must appear here.
     """
 
-    states: dict[Bytes32, State] = {}
+    states: dict[Bytes32, StateT] = Field(default_factory=dict)
     """
     Mapping from block root to State objects.
 
@@ -113,14 +119,18 @@ class Store(StrictBaseModel):
     validator_id: ValidatorIndex | None
     """Index of the validator running this store instance."""
 
-    attestation_signatures: dict[AttestationData, set[AttestationSignatureEntry]] = {}
+    attestation_signatures: dict[AttestationData, set[AttestationSignatureEntry]] = Field(
+        default_factory=dict
+    )
     """
     Per-validator XMSS signatures learned from committee attesters.
 
     Keyed by AttestationData.
     """
 
-    latest_new_aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = {}
+    latest_new_aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = Field(
+        default_factory=dict
+    )
     """
     Aggregated signature proofs pending processing.
 
@@ -129,7 +139,9 @@ class Store(StrictBaseModel):
     Populated from blocks or gossip aggregated attestations.
     """
 
-    latest_known_aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = {}
+    latest_known_aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = Field(
+        default_factory=dict
+    )
     """
     Aggregated signature proofs that have been processed.