diff --git a/packages/testing/src/consensus_testing/test_types/store_checks.py b/packages/testing/src/consensus_testing/test_types/store_checks.py index 53c9bbf8..a07c52a0 100644 --- a/packages/testing/src/consensus_testing/test_types/store_checks.py +++ b/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 diff --git a/src/lean_spec/forks/__init__.py b/src/lean_spec/forks/__init__.py index 8f99bc7b..6af6e401 100644 --- a/src/lean_spec/forks/__init__.py +++ b/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,7 +40,6 @@ "Block", "BlockBody", "BlockHeader", - "BlockLookup", "BlockSignatures", "Config", "DEFAULT_REGISTRY", @@ -45,6 +47,7 @@ "ForkProtocol", "ForkRegistry", "LstarSpec", + "LstarStore", "SignedAggregatedAttestation", "SignedAttestation", "SignedBlock", diff --git a/src/lean_spec/forks/lstar/__init__.py b/src/lean_spec/forks/lstar/__init__.py index d1caedae..801ac3e2 100644 --- a/src/lean_spec/forks/lstar/__init__.py +++ b/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"] diff --git a/src/lean_spec/forks/lstar/containers/block/__init__.py b/src/lean_spec/forks/lstar/containers/block/__init__.py index 2f9ef394..3bce45ae 100644 --- a/src/lean_spec/forks/lstar/containers/block/__init__.py +++ b/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", diff --git a/src/lean_spec/forks/lstar/containers/block/block.py b/src/lean_spec/forks/lstar/containers/block/block.py index ada831f2..d2e4d617 100644 --- a/src/lean_spec/forks/lstar/containers/block/block.py +++ b/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.""" diff --git a/src/lean_spec/forks/lstar/spec.py b/src/lean_spec/forks/lstar/spec.py index 2e6ef55b..164d68bf 100644 --- a/src/lean_spec/forks/lstar/spec.py +++ b/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: diff --git a/src/lean_spec/forks/lstar/store.py b/src/lean_spec/forks/lstar/store.py index db9ca1de..4d96ba6b 100644 --- a/src/lean_spec/forks/lstar/store.py +++ b/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. diff --git a/src/lean_spec/subspecs/sync/service.py b/src/lean_spec/subspecs/sync/service.py index 49946818..d85a5cc5 100644 --- a/src/lean_spec/subspecs/sync/service.py +++ b/src/lean_spec/subspecs/sync/service.py @@ -42,7 +42,6 @@ from lean_spec.forks import ( Block, - BlockLookup, LstarSpec, SignedAggregatedAttestation, SignedAttestation, @@ -91,7 +90,7 @@ def head_slot(self) -> Slot: return store.blocks[store.head].slot -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