From 68470b59bf93bd4a7c860801fab204cc35920fd4 Mon Sep 17 00:00:00 2001 From: anshalshukla Date: Fri, 6 Mar 2026 20:47:56 +0530 Subject: [PATCH] simplify store mappings --- .../testing/src/consensus_testing/keys.py | 16 +- .../test_fixtures/fork_choice.py | 54 ++-- .../subspecs/containers/state/state.py | 175 ++++-------- src/lean_spec/subspecs/forkchoice/__init__.py | 3 +- src/lean_spec/subspecs/forkchoice/store.py | 259 +++++++----------- src/lean_spec/subspecs/validator/service.py | 16 +- src/lean_spec/subspecs/xmss/aggregation.py | 27 -- tests/lean_spec/helpers/builders.py | 12 +- .../containers/test_state_aggregation.py | 181 ++++++------ .../forkchoice/test_compute_block_weights.py | 12 +- .../forkchoice/test_store_attestations.py | 184 +++++-------- .../subspecs/forkchoice/test_store_pruning.py | 177 ++++-------- .../subspecs/forkchoice/test_validator.py | 46 ++-- .../subspecs/validator/test_service.py | 7 +- 14 files changed, 445 insertions(+), 724 deletions(-) diff --git a/packages/testing/src/consensus_testing/keys.py b/packages/testing/src/consensus_testing/keys.py index 323ba870..e55c6031 100755 --- a/packages/testing/src/consensus_testing/keys.py +++ b/packages/testing/src/consensus_testing/keys.py @@ -46,10 +46,7 @@ ) from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.koalabear import Fp -from lean_spec.subspecs.xmss.aggregation import ( - AggregatedSignatureProof, - SignatureKey, -) +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.subspecs.xmss.constants import TARGET_CONFIG from lean_spec.subspecs.xmss.containers import KeyPair, PublicKey, Signature from lean_spec.subspecs.xmss.interface import ( @@ -333,7 +330,8 @@ def sign_attestation_data( def build_attestation_signatures( self, aggregated_attestations: AggregatedAttestations, - signature_lookup: Mapping[SignatureKey, Signature] | None = None, + signature_lookup: Mapping[AttestationData, Mapping[ValidatorIndex, Signature]] + | None = None, ) -> AttestationSignatures: """ Build attestation signatures for already-aggregated attestations. @@ -349,12 +347,12 @@ def build_attestation_signatures( message = agg.data.data_root_bytes() slot = agg.data.slot + # Look up pre-computed signatures by attestation data and validator ID. + sigs_for_data = lookup.get(agg.data, {}) + public_keys: list[PublicKey] = [self.get_public_key(vid) for vid in validator_ids] signatures: list[Signature] = [ - ( - lookup.get(SignatureKey(vid, message)) - or self.sign_attestation_data(vid, agg.data) - ) + sigs_for_data.get(vid) or self.sign_attestation_data(vid, agg.data) for vid in validator_ids ] diff --git a/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py b/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py index 30693e24..bf6e7aef 100644 --- a/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py +++ b/packages/testing/src/consensus_testing/test_fixtures/fork_choice.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import ClassVar, Self +from typing import ClassVar, Mapping, Self from pydantic import model_validator @@ -38,7 +38,6 @@ from lean_spec.subspecs.containers.validator import ValidatorIndex from lean_spec.subspecs.forkchoice import Store from lean_spec.subspecs.ssz import hash_tree_root -from lean_spec.subspecs.xmss.aggregation import SignatureKey from lean_spec.subspecs.xmss.containers import Signature from lean_spec.types import Bytes32, Uint64 @@ -381,20 +380,23 @@ def _build_block_from_spec( # # Attestations vote for blocks and influence fork choice weight. # The spec may include attestations to include in this block. - attestations, attestation_signatures, valid_signature_keys = ( - self._build_attestations_from_spec( - spec, store, block_registry, parent_root, key_manager - ) + ( + attestations, + attestation_signatures, + valid_attestations, + ) = self._build_attestations_from_spec( + spec, store, block_registry, parent_root, key_manager ) - # Merge per-attestation signatures into the Store's gossip signature cache. - # Required so the Store can aggregate committee signatures later when building payloads. + # Merge valid attestation signatures into the Store's gossip cache. + # Only attestations with valid (non-dummy) signatures are merged. working_store = store - for attestation in attestations: - sig_key = SignatureKey(attestation.validator_id, attestation.data.data_root_bytes()) - if sig_key not in valid_signature_keys: - continue - if (signature := attestation_signatures.get(sig_key)) is None: + for attestation in valid_attestations: + sigs_for_data = attestation_signatures.get(attestation.data) + if ( + sigs_for_data is None + or (signature := sigs_for_data.get(attestation.validator_id)) is None + ): continue working_store = working_store.on_gossip_attestation( SignedAttestation( @@ -564,7 +566,11 @@ def _build_attestations_from_spec( block_registry: dict[str, Block], parent_root: Bytes32, key_manager: XmssKeyManager, - ) -> tuple[list[Attestation], dict[SignatureKey, Signature], set[SignatureKey]]: + ) -> tuple[ + list[Attestation], + Mapping[AttestationData, Mapping[ValidatorIndex, Signature]], + set[Attestation], + ]: """ Build attestations and signatures from block specification. @@ -580,7 +586,7 @@ def _build_attestations_from_spec( key_manager: Key manager for signing. Returns: - Tuple of (attestations list, signature lookup dict, valid signature keys). + Tuple of (attestations list, signature lookup by data, valid attestations). """ # No attestations specified means empty block body. if spec.attestations is None: @@ -588,8 +594,8 @@ def _build_attestations_from_spec( parent_state = store.states[parent_root] attestations = [] - signature_lookup: dict[SignatureKey, Signature] = {} - valid_signature_keys: set[SignatureKey] = set() + signature_lookup: Mapping[AttestationData, Mapping[ValidatorIndex, Signature]] = {} + valid_attestations: set[Attestation] = set() for aggregated_spec in spec.attestations: # Build attestation data once. @@ -614,19 +620,19 @@ def _build_attestations_from_spec( validator_id, attestation_data, ) + valid_attestations.add(attestation) else: # Dummy signature for testing invalid signature handling. # The Store should reject attestations with bad signatures. signature = create_dummy_signature() - # Index signature by validator and data root. - # This enables lookup during signature aggregation. - sig_key = SignatureKey(validator_id, attestation_data.data_root_bytes()) - signature_lookup[sig_key] = signature - if aggregated_spec.valid_signature: - valid_signature_keys.add(sig_key) + # Index signature by attestation data and validator ID. + signature_lookup.setdefault(attestation_data, {}).setdefault( + validator_id, + signature, + ) - return attestations, signature_lookup, valid_signature_keys + return attestations, signature_lookup, valid_attestations def _build_attestation_data_from_spec( self, diff --git a/src/lean_spec/subspecs/containers/state/state.py b/src/lean_spec/subspecs/containers/state/state.py index 4ff56805..57095136 100644 --- a/src/lean_spec/subspecs/containers/state/state.py +++ b/src/lean_spec/subspecs/containers/state/state.py @@ -9,10 +9,7 @@ from lean_spec.subspecs.chain.clock import Interval from lean_spec.subspecs.chain.config import INTERVALS_PER_SLOT from lean_spec.subspecs.ssz.hash import hash_tree_root -from lean_spec.subspecs.xmss.aggregation import ( - AggregatedSignatureProof, - SignatureKey, -) +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.subspecs.xmss.containers import PublicKey, Signature from lean_spec.types import ( ZERO_HASH, @@ -22,7 +19,7 @@ Uint64, ) -from ..attestation import AggregatedAttestation, AggregationBits, Attestation +from ..attestation import AggregatedAttestation, AggregationBits, Attestation, AttestationData from ..block import Block, BlockBody, BlockHeader from ..block.types import AggregatedAttestations from ..checkpoint import Checkpoint @@ -38,7 +35,7 @@ ) if TYPE_CHECKING: - from lean_spec.subspecs.forkchoice import Store + from lean_spec.subspecs.forkchoice import GossipSignatureEntry, Store class State(Container): @@ -741,7 +738,7 @@ def build_block( attestations: list[Attestation] | None = None, available_attestations: Iterable[Attestation] | None = None, known_block_roots: AbstractSet[Bytes32] | None = None, - aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] | None = None, + aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] | None = None, ) -> tuple[Block, State, list[AggregatedAttestation], list[AggregatedSignatureProof]]: """ Build a valid block on top of this state. @@ -754,9 +751,6 @@ def build_block( processing attestations may update the justified checkpoint, which may make additional attestations valid. - Signatures are looked up from the provided signature maps using - (validator_id, attestation_data_root) as the key. - Args: slot: Target slot for the block. proposer_index: Validator index of the proposer. @@ -764,8 +758,7 @@ def build_block( attestations: Initial attestations to include. available_attestations: Pool of attestations to collect from. known_block_roots: Set of known block roots for attestation validation. - gossip_signatures: Per-validator XMSS signatures learned from gossip. - aggregated_payloads: Aggregated signature payloads learned from blocks. + aggregated_payloads: Aggregated signature payloads keyed by attestation data. Returns: Tuple of (Block, post-State, collected attestations, signatures). @@ -804,9 +797,6 @@ def build_block( for attestation in available_attestations: data = attestation.data - validator_id = attestation.validator_id - data_root = data.data_root_bytes() - sig_key = SignatureKey(validator_id, data_root) # Skip if target block is unknown. if data.head.root not in known_block_roots: @@ -821,12 +811,20 @@ def build_block( continue # We can only include an attestation if we have some way to later provide - # an aggregated proof for its group: - # - at least one aggregated proof learned from a block that references - # this validator+data. - has_block_proof = bool(aggregated_payloads and sig_key in aggregated_payloads) + # an aggregated proof for this attestation. + # - at least one proof for the attestation data with this validator's + # participation bit set + if not aggregated_payloads or data not in aggregated_payloads: + continue + + vid = attestation.validator_id + has_proof_for_validator = any( + int(vid) < len(proof.participants.data) + and bool(proof.participants.data[int(vid)]) + for proof in aggregated_payloads[data] + ) - if has_block_proof: + if has_proof_for_validator: new_attestations.append(attestation) # Fixed point reached: no new attestations found. @@ -864,7 +862,7 @@ def build_block( def aggregate_gossip_signatures( self, attestations: Collection[Attestation], - gossip_signatures: dict[SignatureKey, "Signature"] | None = None, + gossip_signatures: dict[AttestationData, set[GossipSignatureEntry]] | None = None, ) -> list[tuple[AggregatedAttestation, AggregatedSignatureProof]]: """ Collect aggregated signatures from gossip network and aggregate them. @@ -877,8 +875,9 @@ def aggregate_gossip_signatures( ---------- attestations : Collection[Attestation] Individual attestations to aggregate and sign. - gossip_signatures : dict[SignatureKey, Signature] | None - Per-validator XMSS signatures learned from the gossip network. + gossip_signatures : dict[AttestationData, set[GossipSignatureEntry]] | None + Per-validator XMSS signatures learned from the gossip network, + keyed by the attestation data they signed. Returns: ------- @@ -912,18 +911,14 @@ def aggregate_gossip_signatures( gossip_keys: list[PublicKey] = [] gossip_ids: list[ValidatorIndex] = [] - # Attempt to collect each validator's signature from gossip. - # - # Signatures are keyed by (validator ID, data root). - # - If a signature exists, we add it to our collection. - if gossip_signatures: - for vid in validator_ids: - key = SignatureKey(vid, data_root) - if (sig := gossip_signatures.get(key)) is not None: - # Found a signature: collect it along with the public key. - gossip_sigs.append(sig) - gossip_keys.append(self.validators[vid].get_pubkey()) - gossip_ids.append(vid) + # Look up signatures by attestation data directly. + # Sort by validator ID for deterministic aggregation order. + if gossip_signatures and (entries := gossip_signatures.get(data)): + for entry in sorted(entries, key=lambda e: e.validator_id): + if entry.validator_id in validator_ids: + gossip_sigs.append(entry.signature) + gossip_keys.append(self.validators[entry.validator_id].get_pubkey()) + gossip_ids.append(entry.validator_id) # If we collected any gossip signatures, aggregate them into a proof. # @@ -948,27 +943,24 @@ def aggregate_gossip_signatures( def select_aggregated_proofs( self, attestations: list[Attestation], - aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] | None = None, + aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] | None = None, ) -> tuple[list[AggregatedAttestation], list[AggregatedSignatureProof]]: """ Select aggregated proofs for a set of attestations. This method selects aggregated proofs from aggregated_payloads, - prioritizing proofs from the most recent blocks. + using a greedy set-cover approach to minimize the number of proofs. Strategy: - 1. For each attestation group, aggregate as many signatures as possible - from the most recent block's proofs. - 2. If remaining validators exist after step 1, include proofs from - previous blocks that cover them. + For each attestation group, greedily pick proofs that cover the most + remaining validators until all are covered or no more proofs exist. Parameters: ---------- attestations : list[Attestation] Individual attestations to aggregate and sign. - aggregated_payloads : dict[SignatureKey, list[AggregatedSignatureProof]] | None - Aggregated proofs learned from previously-seen blocks. - The list for each key should be ordered with most recent proofs first. + aggregated_payloads : dict[AttestationData, set[AggregatedSignatureProof]] | None + Aggregated proofs keyed by attestation data. Returns: ------- @@ -980,87 +972,36 @@ def select_aggregated_proofs( # Group individual attestations by data for aggregated in AggregatedAttestation.aggregate_by_data(attestations): data = aggregated.data - data_root = data.data_root_bytes() - validator_ids = ( - aggregated.aggregation_bits.to_validator_indices() - ) # validators contributed to this attestation + validator_ids = aggregated.aggregation_bits.to_validator_indices() - # Validators that are missing in the current aggregation are put into remaining. + # Validators that still need proof coverage. remaining: set[ValidatorIndex] = set(validator_ids) - # Fallback to existing proofs - # - # Some validators may not have broadcast their signatures over gossip, - # but we might have seen proofs for them in previously-received blocks. - # - # Example scenario: - # - # - We need signatures from validators {0, 1, 2, 3, 4}. - # - Gossip gave us signatures for {0, 1}. - # - Remaining: {2, 3, 4}. - # - From old blocks, we have: - # • Proof A covering {2, 3} - # • Proof B covering {3, 4} - # • Proof C covering {4} - # - # We want to cover {2, 3, 4} with as few proofs as possible. - # A greedy approach: always pick the proof with the largest overlap. - # - # - Iteration 1: Proof A covers {2, 3} (2 validators). Pick it. - # Remaining: {4}. - # - Iteration 2: Proof B covers {4} (1 validator). Pick it. - # Remaining: {} → done. - # - # Result: 2 proofs instead of 3. - - while remaining and aggregated_payloads: - # Step 1: Find candidate proofs for a remaining validator. - # - # Proofs are indexed by (validator ID, data root). We pick any - # validator still in the remaining set and look up proofs that - # include them. - target_id = next(iter(remaining)) - candidates = aggregated_payloads.get(SignatureKey(target_id, data_root), []) - - # No proofs found for this validator. - # Remove it and try another validator. - if not candidates: - remaining.discard(target_id) - continue + # Look up all proofs for this attestation data directly. + candidates = sorted( + (aggregated_payloads.get(data, set()) if aggregated_payloads else set()), + key=lambda p: -len(p.participants.to_validator_indices()), + ) - # Step 2: Pick the proof covering the most remaining validators. - # - # At each step, we select the single proof that eliminates the highest - # number of *currently missing* validators from our list. - # - # The 'score' of a candidate proof is defined as the size of the - # intersection between: - # A. The validators inside the proof (`p.participants`) - # B. The validators we still need (`remaining`) - # - # Example: - # Remaining needed : {Alice, Bob, Charlie} - # Proof 1 covers : {Alice, Dave} -> Score: 1 (Only Alice counts) - # Proof 2 covers : {Bob, Charlie, Eve} -> Score: 2 (Bob & Charlie count) - # -> Result: We pick Proof 2 because it has the highest score. - best, covered = max( - ((p, set(p.participants.to_validator_indices())) for p in candidates), - # Calculate the intersection size (A ∩ B) for every candidate. - key=lambda pair: len(pair[1] & remaining), - ) + if not candidates: + continue - # Guard: If the best proof has zero overlap with remaining, stop. - if covered.isdisjoint(remaining): + # Greedy set-cover: candidates are pre-sorted by participant count + # (most validators first). Iterate in order and pick any proof that + # overlaps with remaining validators. + # + # TODO: We don't support recursive aggregation yet. + # In the future, we should be able to aggregate the proofs into a single proof. + for proof in candidates: + if not remaining: break - - # Step 3: Record the proof and remove covered validators. - # - # TODO: We don't support recursive aggregation yet. - # In the future, we should be able to aggregate the proofs into a single proof. + covered = set(proof.participants.to_validator_indices()) + if covered.isdisjoint(remaining): + continue results.append( ( - AggregatedAttestation(aggregation_bits=best.participants, data=data), - best, + AggregatedAttestation(aggregation_bits=proof.participants, data=data), + proof, ) ) remaining -= covered diff --git a/src/lean_spec/subspecs/forkchoice/__init__.py b/src/lean_spec/subspecs/forkchoice/__init__.py index 5d211007..8a916517 100644 --- a/src/lean_spec/subspecs/forkchoice/__init__.py +++ b/src/lean_spec/subspecs/forkchoice/__init__.py @@ -5,8 +5,9 @@ providing the core functionality for determining the canonical chain head. """ -from .store import Store +from .store import GossipSignatureEntry, Store __all__ = [ + "GossipSignatureEntry", "Store", ] diff --git a/src/lean_spec/subspecs/forkchoice/store.py b/src/lean_spec/subspecs/forkchoice/store.py index bab9ff46..2032903c 100644 --- a/src/lean_spec/subspecs/forkchoice/store.py +++ b/src/lean_spec/subspecs/forkchoice/store.py @@ -4,9 +4,10 @@ The Store tracks all information required for the LMD GHOST forkchoice algorithm. """ -__all__ = ["Store"] +__all__ = ["GossipSignatureEntry", "Store"] from collections import defaultdict +from typing import NamedTuple from lean_spec.subspecs.chain.clock import Interval from lean_spec.subspecs.chain.config import ( @@ -32,7 +33,6 @@ from lean_spec.subspecs.xmss.aggregation import ( AggregatedSignatureProof, AggregationError, - SignatureKey, ) from lean_spec.subspecs.xmss.containers import Signature from lean_spec.subspecs.xmss.interface import TARGET_SIGNATURE_SCHEME, GeneralizedXmssScheme @@ -44,6 +44,18 @@ from lean_spec.types.container import Container +class GossipSignatureEntry(NamedTuple): + """ + Single validator's XMSS signature for an attestation, as learned from gossip. + + Used as an element in the gossip_signatures map: one entry per validator + that attested to the same AttestationData. + """ + + validator_id: ValidatorIndex + signature: Signature + + class Store(Container): """ Forkchoice store tracking chain state and validator attestations. @@ -123,22 +135,14 @@ class Store(Container): validator_id: ValidatorIndex | None """Index of the validator running this store instance.""" - gossip_signatures: dict[SignatureKey, Signature] = {} + gossip_signatures: dict[AttestationData, set[GossipSignatureEntry]] = {} """ Per-validator XMSS signatures learned from committee attesters. - Keyed by SignatureKey(validator_id, attestation_data_root). + Keyed by AttestationData. """ - attestation_data_by_root: dict[Bytes32, AttestationData] = {} - """ - Mapping from attestation data root to full AttestationData. - - This allows reconstructing attestations from aggregated payloads. - Entries are pruned once their target checkpoint falls at or before finalization. - """ - - latest_new_aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] = {} + latest_new_aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = {} """ Aggregated signature proofs pending processing. @@ -147,7 +151,7 @@ class Store(Container): Populated from blocks or gossip aggregated attestations. """ - latest_known_aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] = {} + latest_known_aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = {} """ Aggregated signature proofs that have been processed. @@ -175,54 +179,31 @@ def prune_stale_attestation_data(self) -> "Store": """ finalized_slot = self.latest_finalized.slot - # Identify which attestations are stale. - # - # An attestation is stale if its target slot is at or before finalization. - # We collect all such roots in one pass to avoid repeated lookups. - stale_roots: set[Bytes32] = { - data_root - for data_root, data in self.attestation_data_by_root.items() - if data.target.slot <= finalized_slot - } - - # Early return if nothing to prune. - # - # Return the same instance to avoid unnecessary object creation. - if not stale_roots: - return self - # Filter out stale entries from all attestation-related mappings. # - # Each mapping is keyed by attestation data root, so we check membership - # against the stale set we computed above. - - new_attestation_data = { - root: data - for root, data in self.attestation_data_by_root.items() - if root not in stale_roots - } + # Each mapping is keyed by attestation data, so we check membership by slot + # against the finalized slot. new_gossip_sigs = { - key: sig - for key, sig in self.gossip_signatures.items() - if key.data_root not in stale_roots + attestation_data: sigs + for attestation_data, sigs in self.gossip_signatures.items() + if attestation_data.target.slot > finalized_slot } new_aggregated_new = { - key: proofs - for key, proofs in self.latest_new_aggregated_payloads.items() - if key.data_root not in stale_roots + attestation_data: proofs + for attestation_data, proofs in self.latest_new_aggregated_payloads.items() + if attestation_data.target.slot > finalized_slot } new_aggregated_known = { - key: proofs - for key, proofs in self.latest_known_aggregated_payloads.items() - if key.data_root not in stale_roots + attestation_data: proofs + for attestation_data, proofs in self.latest_known_aggregated_payloads.items() + if attestation_data.target.slot > finalized_slot } return self.model_copy( update={ - "attestation_data_by_root": new_attestation_data, "gossip_signatures": new_gossip_sigs, "latest_new_aggregated_payloads": new_aggregated_new, "latest_known_aggregated_payloads": new_aggregated_known, @@ -329,27 +310,23 @@ def on_gossip_attestation( public_key, attestation_data.slot, attestation_data.data_root_bytes(), scheme ), "Signature verification failed" - # Store signature and attestation data for later aggregation - new_committee_sigs = dict(self.gossip_signatures) - new_attestation_data_by_root = dict(self.attestation_data_by_root) - data_root = attestation_data.data_root_bytes() + # Store signature and attestation data for later aggregation. + # Copy the inner sets so we can add to them without mutating the previous store. + new_committee_sigs = {k: set(v) for k, v in self.gossip_signatures.items()} if is_aggregator: assert self.validator_id is not None, "Current validator ID must be set for aggregation" current_subnet = self.validator_id.compute_subnet_id(ATTESTATION_COMMITTEE_COUNT) attester_subnet = validator_id.compute_subnet_id(ATTESTATION_COMMITTEE_COUNT) if current_subnet == attester_subnet: - sig_key = SignatureKey(validator_id, data_root) - new_committee_sigs[sig_key] = signature - - # Store attestation data for later extraction - new_attestation_data_by_root[data_root] = attestation_data + new_committee_sigs.setdefault(attestation_data, set()).add( + GossipSignatureEntry(validator_id, signature) + ) # Return store with updated signature map and attestation data return self.model_copy( update={ "gossip_signatures": new_committee_sigs, - "attestation_data_by_root": new_attestation_data_by_root, } ) @@ -408,28 +385,16 @@ def on_gossip_aggregated_attestation( f"Committee aggregation signature verification failed: {exc}" ) from exc - # Shallow-copy the dict and its list values to maintain immutability + # Shallow-copy the dict and its inner sets to preserve immutability. new_aggregated_payloads = { - k: list(v) for k, v in self.latest_new_aggregated_payloads.items() + k: set(v) for k, v in self.latest_new_aggregated_payloads.items() } - data_root = data.data_root_bytes() - - # Store attestation data by root for later retrieval - new_attestation_data_by_root = dict(self.attestation_data_by_root) - new_attestation_data_by_root[data_root] = data - - for vid in validator_ids: - # Update Proof Map - # - # Store the proof so future block builders can reuse this aggregation - key = SignatureKey(vid, data_root) - new_aggregated_payloads.setdefault(key, []).append(proof) + new_aggregated_payloads.setdefault(data, set()).add(proof) # Return store with updated aggregated payloads and attestation data return self.model_copy( update={ "latest_new_aggregated_payloads": new_aggregated_payloads, - "attestation_data_by_root": new_attestation_data_by_root, } ) @@ -531,38 +496,19 @@ def on_block( ) # Copy the aggregated proof map for updates - # Shallow-copy the dict and its list values to maintain immutability + # Shallow-copy the dict and its inner sets to preserve immutability. # Block attestations go directly to "known" payloads (like is_from_block=True) - block_proofs: dict[SignatureKey, list[AggregatedSignatureProof]] = { - k: list(v) for k, v in store.latest_known_aggregated_payloads.items() + block_proofs: dict[AttestationData, set[AggregatedSignatureProof]] = { + k: set(v) for k, v in store.latest_known_aggregated_payloads.items() } - # Store attestation data by root for later retrieval - new_attestation_data_by_root = dict(store.attestation_data_by_root) - for att, proof in zip(aggregated_attestations, attestation_signatures, strict=True): - validator_ids = att.aggregation_bits.to_validator_indices() - data_root = att.data.data_root_bytes() - - # Store the attestation data - new_attestation_data_by_root[data_root] = att.data - - for vid in validator_ids: - # Update Proof Map - # - # Store the proof so future block builders can reuse this aggregation - key = SignatureKey(vid, data_root) - block_proofs.setdefault(key, []).append(proof) - - # Store proposer attestation data as well - proposer_data_root = proposer_attestation.data.data_root_bytes() - new_attestation_data_by_root[proposer_data_root] = proposer_attestation.data + block_proofs.setdefault(att.data, set()).add(proof) # Update store with new aggregated proofs and attestation data store = store.model_copy( update={ "latest_known_aggregated_payloads": block_proofs, - "attestation_data_by_root": new_attestation_data_by_root, } ) @@ -577,7 +523,7 @@ def on_block( # The proposer casts their attestation in interval 1, after block # proposal. Store the signature so it can be aggregated later. - new_gossip_sigs = dict(store.gossip_signatures) + new_gossip_sigs = {k: set(v) for k, v in store.gossip_signatures.items()} # Store proposer signature for future lookup if it belongs to the same committee # as the current validator. @@ -589,12 +535,9 @@ def on_block( ATTESTATION_COMMITTEE_COUNT ) if proposer_subnet_id == current_validator_subnet_id: - proposer_sig_key = SignatureKey( - proposer_attestation.validator_id, - proposer_attestation.data.data_root_bytes(), - ) - new_gossip_sigs[proposer_sig_key] = ( - signed_block_with_attestation.signature.proposer_signature + sig = signed_block_with_attestation.signature.proposer_signature + new_gossip_sigs.setdefault(proposer_attestation.data, set()).add( + GossipSignatureEntry(proposer_attestation.validator_id, sig) ) # Update store with proposer signature @@ -607,7 +550,7 @@ def on_block( return store def extract_attestations_from_aggregated_payloads( - self, aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] + self, aggregated_payloads: dict[AttestationData, set[AggregatedSignatureProof]] ) -> dict[ValidatorIndex, AttestationData]: """ Extract attestations from aggregated payloads. @@ -616,33 +559,19 @@ def extract_attestations_from_aggregated_payloads( for each validator that participated in the aggregation. Args: - aggregated_payloads: Mapping from SignatureKey to list of aggregated proofs. + aggregated_payloads: Mapping from AttestationData to set of aggregated proofs. Returns: Mapping from ValidatorIndex to AttestationData for each validator. """ attestations: dict[ValidatorIndex, AttestationData] = {} - for sig_key, proofs in aggregated_payloads.items(): - # Get the attestation data from the data root in the signature key - data_root = sig_key.data_root - attestation_data = self.attestation_data_by_root.get(data_root) - - if attestation_data is None: - # Skip if we don't have the attestation data - continue - - # Extract all validator IDs from all proofs for this signature key + for attestation_data, proofs in aggregated_payloads.items(): for proof in proofs: - validator_ids = proof.participants.to_validator_indices() - for vid in validator_ids: - # Store the attestation data for this validator - # If multiple attestations exist for same validator, - # keep the latest (highest slot) - existing = attestations.get(vid) + for validator_id in proof.participants.to_validator_indices(): + existing = attestations.get(validator_id) if existing is None or existing.slot < attestation_data.slot: - attestations[vid] = attestation_data - + attestations[validator_id] = attestation_data return attestations def compute_block_weights(self) -> dict[Bytes32, int]: @@ -828,13 +757,12 @@ def accept_new_attestations(self) -> "Store": New Store with migrated aggregated payloads and updated head. """ # Merge new aggregated payloads into known aggregated payloads - merged_aggregated_payloads = dict(self.latest_known_aggregated_payloads) - for sig_key, proofs in self.latest_new_aggregated_payloads.items(): - if sig_key in merged_aggregated_payloads: - # Merge proof lists for the same signature key - merged_aggregated_payloads[sig_key] = merged_aggregated_payloads[sig_key] + proofs - else: - merged_aggregated_payloads[sig_key] = proofs + merged_aggregated_payloads = { + attestation_data: set(proofs) + for attestation_data, proofs in self.latest_known_aggregated_payloads.items() + } + for attestation_data, proofs in self.latest_new_aggregated_payloads.items(): + merged_aggregated_payloads.setdefault(attestation_data, set()).update(proofs) # Create store with migrated aggregated payloads store = self.model_copy( @@ -913,17 +841,18 @@ def update_safe_target(self) -> "Store": # # The technique: start with a shallow copy of "known", then overlay # every entry from "new" on top. When both pools contain proofs for - # the same signature key, concatenate the proof lists. - all_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] = dict( - self.latest_known_aggregated_payloads - ) - for sig_key, proofs in self.latest_new_aggregated_payloads.items(): - if sig_key in all_payloads: - # Both pools have proofs for this key. Combine them. - all_payloads[sig_key] = all_payloads[sig_key] + proofs + # the same attestation data, merge the proof sets. + all_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = { + attestation_data: set(proofs) + for attestation_data, proofs in self.latest_known_aggregated_payloads.items() + } + for attestation_data, proofs in self.latest_new_aggregated_payloads.items(): + if attestation_data in all_payloads: + # Both pools have proofs for this attestation. Combine them. + all_payloads[attestation_data].update(proofs) else: - # Only "new" has proofs for this key. Add them directly. - all_payloads[sig_key] = proofs + # Only "new" has proofs for this attestation. Add them directly. + all_payloads[attestation_data] = set(proofs) # Convert the merged aggregated payloads into per-validator votes. # @@ -956,27 +885,24 @@ def aggregate_committee_signatures(self) -> tuple["Store", list[SignedAggregated Aggregate committee signatures for attestations in committee_signatures. This method aggregates signatures from the gossip_signatures map. - Attestations are reconstructed from gossip_signatures using attestation_data_by_root. Returns: Tuple of (new Store with updated payloads, list of new SignedAggregatedAttestation). """ - new_aggregated_payloads = dict(self.latest_new_aggregated_payloads) - - # Extract attestations from gossip_signatures - # Each SignatureKey contains (validator_id, data_root) - # We look up the full AttestationData from attestation_data_by_root - attestation_list: list[Attestation] = [] - for sig_key in self.gossip_signatures: - data_root = sig_key.data_root - attestation_data = self.attestation_data_by_root.get(data_root) - if attestation_data is not None: - attestation_list.append( - Attestation(validator_id=sig_key.validator_id, data=attestation_data) - ) + new_aggregated_payloads = { + attestation_data: set(proofs) + for attestation_data, proofs in self.latest_new_aggregated_payloads.items() + } committee_signatures = self.gossip_signatures + # Extract attestations from gossip_signatures + attestation_list: list[Attestation] = [ + Attestation(validator_id=entry.validator_id, data=attestation_data) + for attestation_data, signatures in self.gossip_signatures.items() + for entry in signatures + ] + head_state = self.states[self.head] # Perform aggregation aggregated_results = head_state.aggregate_gossip_signatures( @@ -990,17 +916,22 @@ def aggregate_committee_signatures(self) -> tuple["Store", list[SignedAggregated ] # Compute new aggregated payloads - new_gossip_sigs = dict(self.gossip_signatures) + new_gossip_sigs = { + attestation_data: set(signatures) + for attestation_data, signatures in self.gossip_signatures.items() + } for aggregated_attestation, aggregated_signature in aggregated_results: - data_root = aggregated_attestation.data.data_root_bytes() - validator_ids = aggregated_signature.participants.to_validator_indices() - for vid in validator_ids: - sig_key = SignatureKey(vid, data_root) - new_aggregated_payloads.setdefault(sig_key, []).append(aggregated_signature) - - # Prune successfully aggregated signature from gossip map - if sig_key in new_gossip_sigs: - del new_gossip_sigs[sig_key] + attestation_data = aggregated_attestation.data + new_aggregated_payloads.setdefault(attestation_data, set()).add(aggregated_signature) + + validator_ids = set(aggregated_signature.participants.to_validator_indices()) + existing_entries = new_gossip_sigs.get(attestation_data) + if existing_entries: + remaining = {e for e in existing_entries if e.validator_id not in validator_ids} + if remaining: + new_gossip_sigs[attestation_data] = remaining + else: + del new_gossip_sigs[attestation_data] return self.model_copy( update={ diff --git a/src/lean_spec/subspecs/validator/service.py b/src/lean_spec/subspecs/validator/service.py index 2bc7e7d6..464ef080 100644 --- a/src/lean_spec/subspecs/validator/service.py +++ b/src/lean_spec/subspecs/validator/service.py @@ -53,8 +53,9 @@ BlockWithAttestation, ) from lean_spec.subspecs.containers.slot import Slot +from lean_spec.subspecs.forkchoice import GossipSignatureEntry from lean_spec.subspecs.xmss import TARGET_SIGNATURE_SCHEME, GeneralizedXmssScheme -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types import Uint64 from .registry import ValidatorEntry, ValidatorRegistry @@ -528,20 +529,15 @@ def _store_proposer_attestation_signature( proposer_attestation = signed_block.message.proposer_attestation proposer_signature = signed_block.signature.proposer_signature - data_root = proposer_attestation.data.data_root_bytes() - sig_key = SignatureKey(validator_index, data_root) - new_gossip_sigs = dict(store.gossip_signatures) - new_gossip_sigs[sig_key] = proposer_signature - - # Also store the attestation data for later extraction during aggregation. - new_attestation_data_by_root = dict(store.attestation_data_by_root) - new_attestation_data_by_root[data_root] = proposer_attestation.data + new_gossip_sigs = {k: set(v) for k, v in store.gossip_signatures.items()} + new_gossip_sigs.setdefault(proposer_attestation.data, set()).add( + GossipSignatureEntry(validator_index, proposer_signature) + ) self.sync_service.store = store.model_copy( update={ "gossip_signatures": new_gossip_sigs, - "attestation_data_by_root": new_attestation_data_by_root, } ) diff --git a/src/lean_spec/subspecs/xmss/aggregation.py b/src/lean_spec/subspecs/xmss/aggregation.py index d16a7800..a9a54810 100644 --- a/src/lean_spec/subspecs/xmss/aggregation.py +++ b/src/lean_spec/subspecs/xmss/aggregation.py @@ -3,7 +3,6 @@ from __future__ import annotations from collections.abc import Sequence -from dataclasses import dataclass from typing import Self from lean_multisig_py import ( @@ -16,37 +15,11 @@ from lean_spec.config import LEAN_ENV, LeanEnvMode from lean_spec.subspecs.containers.attestation import AggregationBits from lean_spec.subspecs.containers.slot import Slot -from lean_spec.subspecs.containers.validator import ValidatorIndex from lean_spec.types import ByteListMiB, Bytes32, Container from .containers import PublicKey, Signature -@dataclass(frozen=True, slots=True) -class SignatureKey: - """ - Key for looking up individual validator signatures. - - Used to index signature caches by (validator, message) pairs. - """ - - _validator_id: ValidatorIndex - """The validator who produced the signature.""" - - data_root: Bytes32 - """The hash of the signed data (e.g., attestation data root).""" - - def __init__(self, validator_id: int | ValidatorIndex, data_root: Bytes32) -> None: - """Create a SignatureKey with the given validator_id and data_root.""" - object.__setattr__(self, "_validator_id", ValidatorIndex(validator_id)) - object.__setattr__(self, "data_root", data_root) - - @property - def validator_id(self) -> ValidatorIndex: - """The validator who produced the signature.""" - return self._validator_id - - class AggregationError(Exception): """Raised when signature aggregation or verification fails.""" diff --git a/tests/lean_spec/helpers/builders.py b/tests/lean_spec/helpers/builders.py index 2e048a67..7e186d49 100644 --- a/tests/lean_spec/helpers/builders.py +++ b/tests/lean_spec/helpers/builders.py @@ -31,7 +31,7 @@ from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.containers.state import Validators from lean_spec.subspecs.containers.validator import ValidatorIndex, ValidatorIndices -from lean_spec.subspecs.forkchoice import Store +from lean_spec.subspecs.forkchoice import GossipSignatureEntry, Store from lean_spec.subspecs.koalabear import Fp from lean_spec.subspecs.networking import PeerId from lean_spec.subspecs.networking.peer import PeerInfo @@ -41,7 +41,7 @@ from lean_spec.subspecs.sync.block_cache import BlockCache from lean_spec.subspecs.sync.peer_manager import PeerManager from lean_spec.subspecs.sync.service import SyncService -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.subspecs.xmss.constants import TARGET_CONFIG from lean_spec.subspecs.xmss.containers import Signature from lean_spec.subspecs.xmss.types import ( @@ -389,15 +389,15 @@ def make_store_with_gossip_signatures( validator_id, attestation_slot, ) - data_root = attestation_data.data_root_bytes() gossip_signatures = { - SignatureKey(vid, data_root): key_manager.sign_attestation_data(vid, attestation_data) - for vid in attesting_validators + attestation_data: { + GossipSignatureEntry(vid, key_manager.sign_attestation_data(vid, attestation_data)) + for vid in attesting_validators + }, } store = store.model_copy( update={ "gossip_signatures": gossip_signatures, - "attestation_data_by_root": {data_root: attestation_data}, } ) return store, attestation_data diff --git a/tests/lean_spec/subspecs/containers/test_state_aggregation.py b/tests/lean_spec/subspecs/containers/test_state_aggregation.py index aa4153ae..bb6b15fb 100644 --- a/tests/lean_spec/subspecs/containers/test_state_aggregation.py +++ b/tests/lean_spec/subspecs/containers/test_state_aggregation.py @@ -11,8 +11,8 @@ from lean_spec.subspecs.containers.checkpoint import Checkpoint from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.containers.validator import ValidatorIndex, ValidatorIndices +from lean_spec.subspecs.forkchoice import GossipSignatureEntry from lean_spec.subspecs.ssz.hash import hash_tree_root -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey from tests.lean_spec.helpers import ( make_aggregated_proof, make_attestation_data_simple, @@ -30,12 +30,14 @@ def test_aggregated_signatures_prefers_full_gossip_payload( Slot(2), make_bytes32(3), make_bytes32(4), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(2)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(i), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(i), att_data - ) - for i in range(2) + att_data: { + GossipSignatureEntry( + ValidatorIndex(i), + container_key_manager.sign_attestation_data(ValidatorIndex(i), att_data), + ) + for i in range(2) + } } results = state.aggregate_gossip_signatures( @@ -57,7 +59,7 @@ def test_aggregated_signatures_prefers_full_gossip_payload( public_keys = [container_key_manager.get_public_key(ValidatorIndex(i)) for i in range(2)] aggregated_proofs[0].verify( public_keys=public_keys, - message=data_root, + message=att_data.data_root_bytes(), slot=att_data.slot, ) @@ -72,11 +74,13 @@ def test_aggregate_signatures_splits_when_needed( Slot(3), make_bytes32(5), make_bytes32(6), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(3)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(0), att_data - ) + att_data: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data), + ) + } } block_proof = make_aggregated_proof( @@ -84,8 +88,7 @@ def test_aggregate_signatures_splits_when_needed( ) aggregated_payloads = { - SignatureKey(ValidatorIndex(1), data_root): [block_proof], - SignatureKey(ValidatorIndex(2), data_root): [block_proof], + att_data: {block_proof}, } gossip_results = state.aggregate_gossip_signatures( @@ -116,7 +119,7 @@ def test_aggregate_signatures_splits_when_needed( if participants == [ValidatorIndex(0)]: proof.verify( public_keys=[container_key_manager.get_public_key(ValidatorIndex(0))], - message=data_root, + message=att_data.data_root_bytes(), slot=att_data.slot, ) @@ -142,7 +145,7 @@ def test_build_block_collects_valid_available_attestations( data_root = att_data.data_root_bytes() proof = make_aggregated_proof(container_key_manager, [ValidatorIndex(0)], att_data) - aggregated_payloads = {SignatureKey(ValidatorIndex(0), data_root): [proof]} + aggregated_payloads = {att_data: {proof}} block, post_state, aggregated_atts, aggregated_proofs = state.build_block( slot=Slot(1), @@ -240,22 +243,27 @@ def test_aggregated_signatures_with_multiple_data_groups( Attestation(validator_id=ValidatorIndex(3), data=att_data2), ] - data_root1 = att_data1.data_root_bytes() - data_root2 = att_data2.data_root_bytes() - gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root1): ( - container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data1) - ), - SignatureKey(ValidatorIndex(1), data_root1): ( - container_key_manager.sign_attestation_data(ValidatorIndex(1), att_data1) - ), - SignatureKey(ValidatorIndex(2), data_root2): ( - container_key_manager.sign_attestation_data(ValidatorIndex(2), att_data2) - ), - SignatureKey(ValidatorIndex(3), data_root2): ( - container_key_manager.sign_attestation_data(ValidatorIndex(3), att_data2) - ), + att_data1: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data1), + ), + GossipSignatureEntry( + ValidatorIndex(1), + container_key_manager.sign_attestation_data(ValidatorIndex(1), att_data1), + ), + }, + att_data2: { + GossipSignatureEntry( + ValidatorIndex(2), + container_key_manager.sign_attestation_data(ValidatorIndex(2), att_data2), + ), + GossipSignatureEntry( + ValidatorIndex(3), + container_key_manager.sign_attestation_data(ValidatorIndex(3), att_data2), + ), + }, } results = state.aggregate_gossip_signatures( @@ -290,22 +298,21 @@ def test_aggregated_signatures_falls_back_to_block_payload( Slot(11), make_bytes32(28), make_bytes32(29), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(2)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(0), att_data - ) + att_data: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data), + ) + } } block_proof = make_aggregated_proof( container_key_manager, [ValidatorIndex(0), ValidatorIndex(1)], att_data ) - aggregated_payloads = { - SignatureKey(ValidatorIndex(0), data_root): [block_proof], - SignatureKey(ValidatorIndex(1), data_root): [block_proof], - } + aggregated_payloads = {att_data: {block_proof}} gossip_results = state.aggregate_gossip_signatures( attestations, @@ -329,7 +336,7 @@ def test_aggregated_signatures_falls_back_to_block_payload( if participants == [ValidatorIndex(0)]: proof.verify( public_keys=[container_key_manager.get_public_key(ValidatorIndex(0))], - message=data_root, + message=att_data.data_root_bytes(), slot=att_data.slot, ) @@ -378,7 +385,6 @@ def test_build_block_state_root_valid_when_signatures_split( target=target, source=source, ) - data_root = att_data.data_root_bytes() attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(3)] @@ -387,11 +393,7 @@ def test_build_block_state_root_valid_when_signatures_split( fallback_proof = make_aggregated_proof( container_key_manager, [ValidatorIndex(1), ValidatorIndex(2)], att_data ) - aggregated_payloads = { - SignatureKey(ValidatorIndex(0), data_root): [proof_0], - SignatureKey(ValidatorIndex(1), data_root): [fallback_proof], - SignatureKey(ValidatorIndex(2), data_root): [fallback_proof], - } + aggregated_payloads = {att_data: {proof_0, fallback_proof}} block, post_state, aggregated_atts, _ = pre_state.build_block( slot=Slot(1), @@ -447,7 +449,6 @@ def test_greedy_selects_proof_with_maximum_overlap( Slot(12), make_bytes32(61), make_bytes32(62), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(4)] - data_root = att_data.data_root_bytes() proof_a = make_aggregated_proof( container_key_manager, [ValidatorIndex(0), ValidatorIndex(1)], att_data @@ -459,12 +460,7 @@ def test_greedy_selects_proof_with_maximum_overlap( ) proof_c = make_aggregated_proof(container_key_manager, [ValidatorIndex(3)], att_data) - aggregated_payloads: dict[SignatureKey, list[AggregatedSignatureProof]] = { - SignatureKey(ValidatorIndex(0), data_root): [proof_a], - SignatureKey(ValidatorIndex(1), data_root): [proof_a, proof_b], - SignatureKey(ValidatorIndex(2), data_root): [proof_b], - SignatureKey(ValidatorIndex(3), data_root): [proof_b, proof_c], - } + aggregated_payloads = {att_data: {proof_a, proof_b, proof_c}} aggregated_atts, aggregated_proofs = state.select_aggregated_proofs( attestations, @@ -505,25 +501,25 @@ def test_greedy_stops_when_no_useful_proofs_remain( Slot(13), make_bytes32(71), make_bytes32(72), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(5)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(0), att_data - ), - SignatureKey(ValidatorIndex(1), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(1), att_data - ), + att_data: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data), + ), + GossipSignatureEntry( + ValidatorIndex(1), + container_key_manager.sign_attestation_data(ValidatorIndex(1), att_data), + ), + } } proof_23 = make_aggregated_proof( container_key_manager, [ValidatorIndex(2), ValidatorIndex(3)], att_data ) - aggregated_payloads = { - SignatureKey(ValidatorIndex(2), data_root): [proof_23], - SignatureKey(ValidatorIndex(3), data_root): [proof_23], - } + aggregated_payloads = {att_data: {proof_23}} gossip_results = state.aggregate_gossip_signatures( attestations, @@ -576,12 +572,14 @@ def test_greedy_handles_overlapping_proof_chains( Slot(14), make_bytes32(81), make_bytes32(82), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(5)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(0), att_data - ), + att_data: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data), + ), + } } proof_a = make_aggregated_proof( @@ -594,12 +592,7 @@ def test_greedy_handles_overlapping_proof_chains( container_key_manager, [ValidatorIndex(3), ValidatorIndex(4)], att_data ) - aggregated_payloads = { - SignatureKey(ValidatorIndex(1), data_root): [proof_a], - SignatureKey(ValidatorIndex(2), data_root): [proof_a, proof_b], - SignatureKey(ValidatorIndex(3), data_root): [proof_b, proof_c], - SignatureKey(ValidatorIndex(4), data_root): [proof_c], - } + aggregated_payloads = {att_data: {proof_a, proof_b, proof_c}} gossip_results = state.aggregate_gossip_signatures( attestations, @@ -646,16 +639,13 @@ def test_greedy_single_validator_proofs( Slot(15), make_bytes32(91), make_bytes32(92), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(3)] - data_root = att_data.data_root_bytes() proofs = [ make_aggregated_proof(container_key_manager, [ValidatorIndex(i)], att_data) for i in range(3) ] - aggregated_payloads = { - SignatureKey(ValidatorIndex(i), data_root): [proofs[i]] for i in range(3) - } + aggregated_payloads = {att_data: set(proofs)} aggregated_atts, aggregated_proofs = state.select_aggregated_proofs( attestations, @@ -705,22 +695,21 @@ def test_validator_in_both_gossip_and_fallback_proof( Slot(16), make_bytes32(101), make_bytes32(102), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(2)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(0), att_data - ), + att_data: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data), + ), + } } fallback_proof = make_aggregated_proof( container_key_manager, [ValidatorIndex(0), ValidatorIndex(1)], att_data ) - aggregated_payloads = { - SignatureKey(ValidatorIndex(0), data_root): [fallback_proof], - SignatureKey(ValidatorIndex(1), data_root): [fallback_proof], - } + aggregated_payloads = {att_data: {fallback_proof}} gossip_results = state.aggregate_gossip_signatures( attestations, @@ -746,7 +735,11 @@ def test_validator_in_both_gossip_and_fallback_proof( for proof in aggregated_proofs: participants = proof.participants.to_validator_indices() public_keys = [container_key_manager.get_public_key(vid) for vid in participants] - proof.verify(public_keys=public_keys, message=data_root, slot=att_data.slot) + proof.verify( + public_keys=public_keys, + message=att_data.data_root_bytes(), + slot=att_data.slot, + ) def test_gossip_none_and_aggregated_payloads_none( @@ -804,7 +797,7 @@ def test_aggregated_payloads_only_no_gossip( att_data, ) - aggregated_payloads = {SignatureKey(ValidatorIndex(i), data_root): [proof] for i in range(3)} + aggregated_payloads = {att_data: {proof}} aggregated_atts, aggregated_proofs = state.select_aggregated_proofs( attestations, @@ -846,21 +839,21 @@ def test_proof_with_extra_validators_beyond_needed( Slot(19), make_bytes32(131), make_bytes32(132), source=source ) attestations = [Attestation(validator_id=ValidatorIndex(i), data=att_data) for i in range(2)] - data_root = att_data.data_root_bytes() gossip_signatures = { - SignatureKey(ValidatorIndex(0), data_root): container_key_manager.sign_attestation_data( - ValidatorIndex(0), att_data - ), + att_data: { + GossipSignatureEntry( + ValidatorIndex(0), + container_key_manager.sign_attestation_data(ValidatorIndex(0), att_data), + ), + } } proof = make_aggregated_proof( container_key_manager, [ValidatorIndex(i) for i in range(4)], att_data ) - aggregated_payloads = { - SignatureKey(ValidatorIndex(1), data_root): [proof], - } + aggregated_payloads = {att_data: {proof}} gossip_results = state.aggregate_gossip_signatures( attestations, diff --git a/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py b/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py index 3233743f..751dc03c 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py +++ b/tests/lean_spec/subspecs/forkchoice/test_compute_block_weights.py @@ -8,7 +8,7 @@ from lean_spec.subspecs.containers.validator import ValidatorIndex, ValidatorIndices from lean_spec.subspecs.forkchoice import Store from lean_spec.subspecs.ssz.hash import hash_tree_root -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types.byte_arrays import ByteListMiB from tests.lean_spec.helpers import make_bytes32, make_signed_block @@ -61,11 +61,9 @@ def test_linear_chain_weight_accumulates_upward(base_store: Store) -> None: target=Checkpoint(root=block2_root, slot=Slot(2)), source=Checkpoint(root=genesis_root, slot=Slot(0)), ) - data_root = att_data.data_root_bytes() - proof = _make_empty_proof([ValidatorIndex(0)]) aggregated_payloads = { - SignatureKey(ValidatorIndex(0), data_root): [proof], + att_data: {proof}, } store = base_store.model_copy( @@ -74,7 +72,6 @@ def test_linear_chain_weight_accumulates_upward(base_store: Store) -> None: "states": new_states, "head": block2_root, "latest_known_aggregated_payloads": aggregated_payloads, - "attestation_data_by_root": {data_root: att_data}, } ) @@ -110,12 +107,10 @@ def test_multiple_attestations_accumulate(base_store: Store) -> None: target=Checkpoint(root=block1_root, slot=Slot(1)), source=Checkpoint(root=genesis_root, slot=Slot(0)), ) - data_root = att_data.data_root_bytes() proof = _make_empty_proof([ValidatorIndex(0), ValidatorIndex(1)]) aggregated_payloads = { - SignatureKey(ValidatorIndex(0), data_root): [proof], - SignatureKey(ValidatorIndex(1), data_root): [proof], + att_data: {proof}, } store = base_store.model_copy( @@ -124,7 +119,6 @@ def test_multiple_attestations_accumulate(base_store: Store) -> None: "states": new_states, "head": block1_root, "latest_known_aggregated_payloads": aggregated_payloads, - "attestation_data_by_root": {data_root: att_data}, } ) diff --git a/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py b/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py index ff3c1ebb..334d137e 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py +++ b/tests/lean_spec/subspecs/forkchoice/test_store_attestations.py @@ -17,7 +17,8 @@ from lean_spec.subspecs.containers.checkpoint import Checkpoint from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.containers.validator import ValidatorIndex, ValidatorIndices -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.subspecs.forkchoice import GossipSignatureEntry +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types import Bytes32, Uint64 from tests.lean_spec.helpers import ( TEST_VALIDATOR_ID, @@ -37,16 +38,14 @@ def test_on_block_processes_multi_validator_aggregations(key_manager: XmssKeyMan attestation_slot = Slot(1) attestation_data = base_store.produce_attestation_data(attestation_slot) - data_root = attestation_data.data_root_bytes() participants = [ValidatorIndex(1), ValidatorIndex(2)] proof = make_aggregated_proof(key_manager, participants, attestation_data) - aggregated_payloads = {SignatureKey(vid, data_root): [proof] for vid in participants} + aggregated_payloads = {attestation_data: {proof}} producer_store = base_store.model_copy( update={ - "attestation_data_by_root": {data_root: attestation_data}, "latest_known_aggregated_payloads": aggregated_payloads, } ) @@ -79,18 +78,18 @@ def test_on_block_preserves_immutability_of_aggregated_payloads( # First block with attestations from validators 1 and 2 attestation_slot_1 = Slot(1) attestation_data_1 = base_store.produce_attestation_data(attestation_slot_1) - data_root_1 = attestation_data_1.data_root_bytes() gossip_sigs_1 = { - SignatureKey(validator_id, data_root_1): key_manager.sign_attestation_data( - validator_id, attestation_data_1 - ) - for validator_id in (ValidatorIndex(1), ValidatorIndex(2)) + attestation_data_1: { + GossipSignatureEntry( + validator_id, key_manager.sign_attestation_data(validator_id, attestation_data_1) + ) + for validator_id in (ValidatorIndex(1), ValidatorIndex(2)) + }, } producer_store_1 = base_store.model_copy( update={ - "attestation_data_by_root": {data_root_1: attestation_data_1}, "gossip_signatures": gossip_sigs_1, } ) @@ -103,18 +102,18 @@ def test_on_block_preserves_immutability_of_aggregated_payloads( # Second block with attestations for the SAME validators attestation_slot_2 = Slot(2) attestation_data_2 = store_after_block_1.produce_attestation_data(attestation_slot_2) - data_root_2 = attestation_data_2.data_root_bytes() gossip_sigs_2 = { - SignatureKey(validator_id, data_root_2): key_manager.sign_attestation_data( - validator_id, attestation_data_2 - ) - for validator_id in (ValidatorIndex(1), ValidatorIndex(2)) + attestation_data_2: { + GossipSignatureEntry( + validator_id, key_manager.sign_attestation_data(validator_id, attestation_data_2) + ) + for validator_id in (ValidatorIndex(1), ValidatorIndex(2)) + }, } producer_store_2 = store_after_block_1.model_copy( update={ - "attestation_data_by_root": {data_root_2: attestation_data_2}, "gossip_signatures": gossip_sigs_2, } ) @@ -178,9 +177,7 @@ def test_same_subnet_stores_signature(self, key_manager: XmssKeyManager) -> None ) # Verify signature does NOT exist before calling the method - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(attester_validator, data_root) - assert sig_key not in store.gossip_signatures, ( + assert attestation_data not in store.gossip_signatures, ( "Precondition: signature should not exist before calling method" ) @@ -194,7 +191,8 @@ def test_same_subnet_stores_signature(self, key_manager: XmssKeyManager) -> None ) # Verify signature NOW exists after calling the method - assert sig_key in updated_store.gossip_signatures, ( + sigs = updated_store.gossip_signatures.get(attestation_data, set()) + assert any(entry.validator_id == attester_validator for entry in sigs), ( "Signature from same-subnet validator should be stored" ) @@ -229,9 +227,8 @@ def test_cross_subnet_ignores_signature(self, key_manager: XmssKeyManager) -> No ) # Verify signature was NOT stored - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(attester_validator, data_root) - assert sig_key not in updated_store.gossip_signatures, ( + sigs = updated_store.gossip_signatures.get(attestation_data, set()) + assert not any(entry.validator_id == attester_validator for entry in sigs), ( "Signature from different-subnet validator should NOT be stored" ) @@ -264,18 +261,17 @@ def test_non_aggregator_never_stores_signature(self, key_manager: XmssKeyManager ) # Verify signature was NOT stored even though same subnet - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(attester_validator, data_root) - assert sig_key not in updated_store.gossip_signatures, ( + sigs = updated_store.gossip_signatures.get(attestation_data, set()) + assert not any(entry.validator_id == attester_validator for entry in sigs), ( "Non-aggregator should never store gossip signatures" ) - def test_attestation_data_always_stored(self, key_manager: XmssKeyManager) -> None: + def test_cross_subnet_does_not_create_gossip_entry(self, key_manager: XmssKeyManager) -> None: """ - Attestation data is stored regardless of aggregator status or subnet. + Cross-subnet attestation does not create a gossip_signatures entry. - The attestation_data_by_root map is always updated for later reference, - even when the signature itself is filtered out. + When the attester is in a different subnet, no entry is created + for that attestation data in gossip_signatures. """ current_validator = ValidatorIndex(0) attester_validator = ValidatorIndex(1) # Different subnet @@ -298,19 +294,11 @@ def test_attestation_data_always_stored(self, key_manager: XmssKeyManager) -> No is_aggregator=True, ) - # Verify signature was NOT stored (cross-subnet filtered) - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(attester_validator, data_root) - assert sig_key not in updated_store.gossip_signatures, ( - "Signature should NOT be stored for cross-subnet validator" + # Verify no gossip entry was created for this attestation data + assert attestation_data not in updated_store.gossip_signatures, ( + "Cross-subnet attestation should not create a gossip_signatures entry" ) - # Verify attestation data WAS stored even though signature wasn't - assert data_root in updated_store.attestation_data_by_root, ( - "Attestation data should always be stored" - ) - assert updated_store.attestation_data_by_root[data_root] == attestation_data - class TestOnGossipAggregatedAttestation: """ @@ -324,7 +312,7 @@ def test_valid_proof_stored_correctly(self, key_manager: XmssKeyManager) -> None Valid aggregated attestation is verified and stored. The proof should be stored in latest_new_aggregated_payloads - keyed by each participating validator's SignatureKey. + keyed by attestation data. """ participants = [ValidatorIndex(1), ValidatorIndex(2)] @@ -354,21 +342,19 @@ def test_valid_proof_stored_correctly(self, key_manager: XmssKeyManager) -> None updated_store = store.on_gossip_aggregated_attestation(signed_aggregated) - # Verify proof is stored for each participant - for vid in participants: - sig_key = SignatureKey(vid, data_root) - assert sig_key in updated_store.latest_new_aggregated_payloads, ( - f"Proof should be stored for validator {vid}" - ) - proofs = updated_store.latest_new_aggregated_payloads[sig_key] - assert len(proofs) == 1 - assert proofs[0] == proof + # Verify proof is stored keyed by attestation data + assert attestation_data in updated_store.latest_new_aggregated_payloads, ( + "Proof should be stored for this attestation data" + ) + proofs = updated_store.latest_new_aggregated_payloads[attestation_data] + assert len(proofs) == 1 + assert proof in proofs - def test_attestation_data_stored_by_root(self, key_manager: XmssKeyManager) -> None: + def test_attestation_data_used_as_key(self, key_manager: XmssKeyManager) -> None: """ - Attestation data is stored in attestation_data_by_root. + Attestation data is used directly as the key in aggregated payloads. - This allows later reconstruction of attestations from proofs. + Proofs are accessible by looking up the attestation data. """ participants = [ValidatorIndex(1)] @@ -397,8 +383,8 @@ def test_attestation_data_stored_by_root(self, key_manager: XmssKeyManager) -> N updated_store = store.on_gossip_aggregated_attestation(signed_aggregated) - assert data_root in updated_store.attestation_data_by_root - assert updated_store.attestation_data_by_root[data_root] == attestation_data + assert attestation_data in updated_store.latest_new_aggregated_payloads + assert proof in updated_store.latest_new_aggregated_payloads[attestation_data] def test_invalid_proof_rejected(self, key_manager: XmssKeyManager) -> None: """ @@ -484,9 +470,8 @@ def test_multiple_proofs_accumulate(self, key_manager: XmssKeyManager) -> None: SignedAggregatedAttestation(data=attestation_data, proof=proof_2) ) - # Validator 1 should have BOTH proofs - sig_key = SignatureKey(ValidatorIndex(1), data_root) - stored_proofs = store.latest_new_aggregated_payloads[sig_key] + # Both proofs should be stored under the same attestation data + stored_proofs = store.latest_new_aggregated_payloads[attestation_data] assert len(stored_proofs) == 2 assert proof_1 in stored_proofs, "First proof should be stored" assert proof_2 in stored_proofs, "Second proof should be stored" @@ -521,15 +506,12 @@ def test_aggregates_gossip_signatures_into_proof(self, key_manager: XmssKeyManag # Perform aggregation updated_store, _ = store.aggregate_committee_signatures() - # Verify proofs were created and stored - data_root = attestation_data.data_root_bytes() - for vid in attesting_validators: - sig_key = SignatureKey(vid, data_root) - assert sig_key in updated_store.latest_new_aggregated_payloads, ( - f"Aggregated proof should be stored for validator {vid}" - ) - proofs = updated_store.latest_new_aggregated_payloads[sig_key] - assert len(proofs) >= 1, "At least one proof should exist" + # Verify proofs were created and stored keyed by attestation data + assert attestation_data in updated_store.latest_new_aggregated_payloads, ( + "Aggregated proof should be stored for this attestation data" + ) + proofs = updated_store.latest_new_aggregated_payloads[attestation_data] + assert len(proofs) >= 1, "At least one proof should exist" def test_aggregated_proof_is_valid(self, key_manager: XmssKeyManager) -> None: """ @@ -549,9 +531,8 @@ def test_aggregated_proof_is_valid(self, key_manager: XmssKeyManager) -> None: updated_store, _ = store.aggregate_committee_signatures() - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) - proof = updated_store.latest_new_aggregated_payloads[sig_key][0] + proofs = updated_store.latest_new_aggregated_payloads[attestation_data] + proof = next(iter(proofs)) # Extract participants from the proof participants = proof.participants.to_validator_indices() @@ -560,7 +541,7 @@ def test_aggregated_proof_is_valid(self, key_manager: XmssKeyManager) -> None: # Verify the proof is valid proof.verify( public_keys=public_keys, - message=data_root, + message=attestation_data.data_root_bytes(), slot=attestation_data.slot, ) @@ -604,39 +585,25 @@ def test_multiple_attestation_data_grouped_separately( source=att_data_1.source, ) - data_root_1 = att_data_1.data_root_bytes() - data_root_2 = att_data_2.data_root_bytes() - # Validators 1 attests to data_1, validator 2 attests to data_2 + sig_1 = key_manager.sign_attestation_data(ValidatorIndex(1), att_data_1) + sig_2 = key_manager.sign_attestation_data(ValidatorIndex(2), att_data_2) gossip_signatures = { - SignatureKey(ValidatorIndex(1), data_root_1): key_manager.sign_attestation_data( - ValidatorIndex(1), att_data_1 - ), - SignatureKey(ValidatorIndex(2), data_root_2): key_manager.sign_attestation_data( - ValidatorIndex(2), att_data_2 - ), - } - - attestation_data_by_root = { - data_root_1: att_data_1, - data_root_2: att_data_2, + att_data_1: {GossipSignatureEntry(ValidatorIndex(1), sig_1)}, + att_data_2: {GossipSignatureEntry(ValidatorIndex(2), sig_2)}, } store = base_store.model_copy( update={ "gossip_signatures": gossip_signatures, - "attestation_data_by_root": attestation_data_by_root, } ) updated_store, _ = store.aggregate_committee_signatures() - # Verify both validators have separate proofs - sig_key_1 = SignatureKey(ValidatorIndex(1), data_root_1) - sig_key_2 = SignatureKey(ValidatorIndex(2), data_root_2) - - assert sig_key_1 in updated_store.latest_new_aggregated_payloads - assert sig_key_2 in updated_store.latest_new_aggregated_payloads + # Verify both attestation data have separate proofs + assert att_data_1 in updated_store.latest_new_aggregated_payloads + assert att_data_2 in updated_store.latest_new_aggregated_payloads class TestTickIntervalAggregation: @@ -674,10 +641,7 @@ def test_interval_2_triggers_aggregation_for_aggregator( updated_store, _ = store.tick_interval(has_proposal=False, is_aggregator=True) # Verify aggregation was performed - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) - - assert sig_key in updated_store.latest_new_aggregated_payloads, ( + assert attestation_data in updated_store.latest_new_aggregated_payloads, ( "Aggregation should occur at interval 2 for aggregators" ) @@ -705,10 +669,7 @@ def test_interval_2_skips_aggregation_for_non_aggregator( updated_store, _ = store.tick_interval(has_proposal=False, is_aggregator=False) # Verify aggregation was NOT performed - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) - - assert sig_key not in updated_store.latest_new_aggregated_payloads, ( + assert attestation_data not in updated_store.latest_new_aggregated_payloads, ( "Aggregation should NOT occur for non-aggregators" ) @@ -727,9 +688,6 @@ def test_other_intervals_do_not_trigger_aggregation(self, key_manager: XmssKeyMa attesting_validators=attesting_validators, ) - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) - # Test intervals 0, 1, 3, 4 (skip 2) non_aggregation_intervals = [0, 1, 3, 4] @@ -743,7 +701,7 @@ def test_other_intervals_do_not_trigger_aggregation(self, key_manager: XmssKeyMa updated_store, _ = test_store.tick_interval(has_proposal=False, is_aggregator=True) - assert sig_key not in updated_store.latest_new_aggregated_payloads, ( + assert attestation_data not in updated_store.latest_new_aggregated_payloads, ( f"Aggregation should NOT occur at interval {target_interval}" ) @@ -816,23 +774,23 @@ def test_gossip_to_aggregation_to_storage(self, key_manager: XmssKeyManager) -> ) # Verify signatures were stored + sigs = store.gossip_signatures.get(attestation_data, set()) for vid in attesting_validators: - sig_key = SignatureKey(vid, data_root) - assert sig_key in store.gossip_signatures, f"Signature for {vid} should be stored" + assert any(entry.validator_id == vid for entry in sigs), ( + f"Signature for {vid} should be stored" + ) # Step 2: Advance to interval 2 (aggregation interval) store = store.model_copy(update={"time": Uint64(1)}) store, _ = store.tick_interval(has_proposal=False, is_aggregator=True) # Step 3: Verify aggregated proofs were created - for vid in attesting_validators: - sig_key = SignatureKey(vid, data_root) - assert sig_key in store.latest_new_aggregated_payloads, ( - f"Aggregated proof for {vid} should exist after interval 2" - ) + assert attestation_data in store.latest_new_aggregated_payloads, ( + "Aggregated proofs should exist after interval 2" + ) # Step 4: Verify the proof is valid - proof = store.latest_new_aggregated_payloads[SignatureKey(ValidatorIndex(1), data_root)][0] + proof = next(iter(store.latest_new_aggregated_payloads[attestation_data])) participants = proof.participants.to_validator_indices() public_keys = [key_manager.get_public_key(vid) for vid in participants] diff --git a/tests/lean_spec/subspecs/forkchoice/test_store_pruning.py b/tests/lean_spec/subspecs/forkchoice/test_store_pruning.py index 8df9cd5c..a69128e0 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_store_pruning.py +++ b/tests/lean_spec/subspecs/forkchoice/test_store_pruning.py @@ -3,8 +3,8 @@ from lean_spec.subspecs.containers.attestation import AggregationBits from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.containers.validator import ValidatorIndex, ValidatorIndices -from lean_spec.subspecs.forkchoice import Store -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.subspecs.forkchoice import GossipSignatureEntry, Store +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types import ByteListMiB, Bytes32 from tests.lean_spec.helpers import ( make_attestation_data, @@ -26,27 +26,24 @@ def test_prunes_entries_with_target_at_finalized(pruning_store: Store) -> None: source_slot=Slot(0), source_root=Bytes32.zero(), ) - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) # Set up store with attestation data and finalized slot at 5 store = store.model_copy( update={ - "attestation_data_by_root": {data_root: attestation_data}, - "gossip_signatures": {sig_key: make_mock_signature()}, + "gossip_signatures": { + attestation_data: {GossipSignatureEntry(ValidatorIndex(1), make_mock_signature())}, + }, "latest_finalized": make_checkpoint(root_seed=255, slot=5), } ) # Verify data exists before pruning - assert data_root in store.attestation_data_by_root - assert sig_key in store.gossip_signatures + assert attestation_data in store.gossip_signatures # Prune should remove entries where target.slot <= finalized.slot pruned_store = store.prune_stale_attestation_data() - assert data_root not in pruned_store.attestation_data_by_root - assert sig_key not in pruned_store.gossip_signatures + assert attestation_data not in pruned_store.gossip_signatures def test_prunes_entries_with_target_before_finalized(pruning_store: Store) -> None: @@ -61,27 +58,24 @@ def test_prunes_entries_with_target_before_finalized(pruning_store: Store) -> No source_slot=Slot(0), source_root=Bytes32.zero(), ) - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) # Set up store with finalized slot at 5 (greater than target.slot) store = store.model_copy( update={ - "attestation_data_by_root": {data_root: attestation_data}, - "gossip_signatures": {sig_key: make_mock_signature()}, + "gossip_signatures": { + attestation_data: {GossipSignatureEntry(ValidatorIndex(1), make_mock_signature())}, + }, "latest_finalized": make_checkpoint(root_seed=255, slot=5), } ) # Verify data exists before pruning - assert data_root in store.attestation_data_by_root - assert sig_key in store.gossip_signatures + assert attestation_data in store.gossip_signatures # Prune should remove entries where target.slot <= finalized.slot pruned_store = store.prune_stale_attestation_data() - assert data_root not in pruned_store.attestation_data_by_root - assert sig_key not in pruned_store.gossip_signatures + assert attestation_data not in pruned_store.gossip_signatures def test_keeps_entries_with_target_after_finalized(pruning_store: Store) -> None: @@ -96,32 +90,28 @@ def test_keeps_entries_with_target_after_finalized(pruning_store: Store) -> None source_slot=Slot(5), source_root=make_bytes32(2), ) - data_root = attestation_data.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) # Set up store with finalized slot at 5 (less than target.slot) store = store.model_copy( update={ - "attestation_data_by_root": {data_root: attestation_data}, - "gossip_signatures": {sig_key: make_mock_signature()}, + "gossip_signatures": { + attestation_data: {GossipSignatureEntry(ValidatorIndex(1), make_mock_signature())}, + }, "latest_finalized": make_checkpoint(root_seed=255, slot=5), } ) # Verify data exists before pruning - assert data_root in store.attestation_data_by_root - assert sig_key in store.gossip_signatures + assert attestation_data in store.gossip_signatures # Prune should keep entries where target.slot > finalized.slot pruned_store = store.prune_stale_attestation_data() - assert data_root in pruned_store.attestation_data_by_root - assert sig_key in pruned_store.gossip_signatures - assert pruned_store.attestation_data_by_root[data_root] == attestation_data + assert attestation_data in pruned_store.gossip_signatures def test_prunes_related_structures_together(pruning_store: Store) -> None: - """Verify all four data structures are pruned atomically.""" + """Verify all three data structures are pruned atomically.""" store = pruning_store # Create stale attestation data @@ -132,8 +122,6 @@ def test_prunes_related_structures_together(pruning_store: Store) -> None: source_slot=Slot(0), source_root=Bytes32.zero(), ) - stale_root = stale_attestation.data_root_bytes() - stale_key = SignatureKey(ValidatorIndex(1), stale_root) # Create fresh attestation data fresh_attestation = make_attestation_data( @@ -143,8 +131,6 @@ def test_prunes_related_structures_together(pruning_store: Store) -> None: source_slot=Slot(5), source_root=make_bytes32(255), ) - fresh_root = fresh_attestation.data_root_bytes() - fresh_key = SignatureKey(ValidatorIndex(2), fresh_root) # Create mock aggregated proof (empty proof data for testing) mock_proof = AggregatedSignatureProof( @@ -157,101 +143,58 @@ def test_prunes_related_structures_together(pruning_store: Store) -> None: # Set up store with both stale and fresh entries in all structures store = store.model_copy( update={ - "attestation_data_by_root": { - stale_root: stale_attestation, - fresh_root: fresh_attestation, - }, "gossip_signatures": { - stale_key: make_mock_signature(), - fresh_key: make_mock_signature(), + stale_attestation: {GossipSignatureEntry(ValidatorIndex(1), make_mock_signature())}, + fresh_attestation: {GossipSignatureEntry(ValidatorIndex(2), make_mock_signature())}, }, "latest_new_aggregated_payloads": { - stale_key: [mock_proof], - fresh_key: [mock_proof], + stale_attestation: {mock_proof}, + fresh_attestation: {mock_proof}, }, "latest_known_aggregated_payloads": { - stale_key: [mock_proof], - fresh_key: [mock_proof], + stale_attestation: {mock_proof}, + fresh_attestation: {mock_proof}, }, "latest_finalized": make_checkpoint(root_seed=255, slot=5), } ) # Verify all data exists before pruning - assert stale_root in store.attestation_data_by_root - assert stale_key in store.gossip_signatures - assert stale_key in store.latest_new_aggregated_payloads - assert stale_key in store.latest_known_aggregated_payloads - assert fresh_root in store.attestation_data_by_root - assert fresh_key in store.gossip_signatures - assert fresh_key in store.latest_new_aggregated_payloads - assert fresh_key in store.latest_known_aggregated_payloads + assert stale_attestation in store.gossip_signatures + assert stale_attestation in store.latest_new_aggregated_payloads + assert stale_attestation in store.latest_known_aggregated_payloads + assert fresh_attestation in store.gossip_signatures + assert fresh_attestation in store.latest_new_aggregated_payloads + assert fresh_attestation in store.latest_known_aggregated_payloads pruned_store = store.prune_stale_attestation_data() # Stale entries should be removed from all structures - assert stale_root not in pruned_store.attestation_data_by_root - assert stale_key not in pruned_store.gossip_signatures - assert stale_key not in pruned_store.latest_new_aggregated_payloads - assert stale_key not in pruned_store.latest_known_aggregated_payloads + assert stale_attestation not in pruned_store.gossip_signatures + assert stale_attestation not in pruned_store.latest_new_aggregated_payloads + assert stale_attestation not in pruned_store.latest_known_aggregated_payloads # Fresh entries should be preserved in all structures - assert fresh_root in pruned_store.attestation_data_by_root - assert fresh_key in pruned_store.gossip_signatures - assert fresh_key in pruned_store.latest_new_aggregated_payloads - assert fresh_key in pruned_store.latest_known_aggregated_payloads - - -def test_returns_self_when_nothing_to_prune(pruning_store: Store) -> None: - """Verify optimization returns same instance when no pruning needed.""" - store = pruning_store - - # Create fresh attestation data (target.slot > finalized.slot) - fresh_attestation = make_attestation_data( - slot=Slot(10), - target_slot=Slot(10), - target_root=make_bytes32(1), - source_slot=Slot(5), - source_root=make_bytes32(2), - ) - data_root = fresh_attestation.data_root_bytes() - sig_key = SignatureKey(ValidatorIndex(1), data_root) + assert fresh_attestation in pruned_store.gossip_signatures + assert fresh_attestation in pruned_store.latest_new_aggregated_payloads + assert fresh_attestation in pruned_store.latest_known_aggregated_payloads - # Set up store with only fresh entries and finalized slot at 5 - store = store.model_copy( - update={ - "attestation_data_by_root": {data_root: fresh_attestation}, - "gossip_signatures": {sig_key: make_mock_signature()}, - "latest_finalized": make_checkpoint(root_seed=255, slot=5), - } - ) - # Verify data exists before pruning - assert data_root in store.attestation_data_by_root - assert sig_key in store.gossip_signatures - - pruned_store = store.prune_stale_attestation_data() - - # Should return the same instance (identity check) - assert pruned_store is store - - -def test_handles_empty_attestation_data(pruning_store: Store) -> None: - """Verify pruning works correctly when attestation_data_by_root is empty.""" +def test_handles_empty_gossip_signatures(pruning_store: Store) -> None: + """Verify pruning works correctly when gossip_signatures is empty.""" store = pruning_store - # Ensure store has empty attestation data - assert len(store.attestation_data_by_root) == 0 + # Ensure store has empty gossip signatures + assert len(store.gossip_signatures) == 0 - # Pruning should not fail and should return same instance + # Pruning should not fail pruned_store = store.prune_stale_attestation_data() - assert pruned_store is store - assert len(pruned_store.attestation_data_by_root) == 0 + assert len(pruned_store.gossip_signatures) == 0 -def test_prunes_multiple_validators_same_data_root(pruning_store: Store) -> None: - """Verify pruning removes entries for multiple validators with same data root.""" +def test_prunes_multiple_validators_same_attestation_data(pruning_store: Store) -> None: + """Verify pruning removes entries for multiple validators with same attestation data.""" store = pruning_store # Create stale attestation data @@ -262,34 +205,28 @@ def test_prunes_multiple_validators_same_data_root(pruning_store: Store) -> None source_slot=Slot(0), source_root=Bytes32.zero(), ) - data_root = stale_attestation.data_root_bytes() # Multiple validators signed the same attestation data - sig_key_1 = SignatureKey(ValidatorIndex(1), data_root) - sig_key_2 = SignatureKey(ValidatorIndex(2), data_root) - store = store.model_copy( update={ - "attestation_data_by_root": {data_root: stale_attestation}, "gossip_signatures": { - sig_key_1: make_mock_signature(), - sig_key_2: make_mock_signature(), + stale_attestation: { + GossipSignatureEntry(ValidatorIndex(1), make_mock_signature()), + GossipSignatureEntry(ValidatorIndex(2), make_mock_signature()), + }, }, "latest_finalized": make_checkpoint(root_seed=255, slot=5), } ) # Verify data exists before pruning - assert data_root in store.attestation_data_by_root - assert sig_key_1 in store.gossip_signatures - assert sig_key_2 in store.gossip_signatures + assert stale_attestation in store.gossip_signatures + assert len(store.gossip_signatures[stale_attestation]) == 2 pruned_store = store.prune_stale_attestation_data() - # Both validators' signatures should be removed - assert data_root not in pruned_store.attestation_data_by_root - assert sig_key_1 not in pruned_store.gossip_signatures - assert sig_key_2 not in pruned_store.gossip_signatures + # All validators' signatures should be removed (whole entry pruned) + assert stale_attestation not in pruned_store.gossip_signatures def test_mixed_stale_and_fresh_entries(pruning_store: Store) -> None: @@ -308,16 +245,14 @@ def test_mixed_stale_and_fresh_entries(pruning_store: Store) -> None: for i in range(1, 11) # Slots 1-10 ] - attestation_data_map = {att.data_root_bytes(): att for att in attestations} gossip_sigs = { - SignatureKey(ValidatorIndex(i), att.data_root_bytes()): make_mock_signature() + att: {GossipSignatureEntry(ValidatorIndex(i), make_mock_signature())} for i, att in enumerate(attestations, start=1) } # Finalized at slot 5 means slots 1-5 are stale, 6-10 are fresh store = store.model_copy( update={ - "attestation_data_by_root": attestation_data_map, "gossip_signatures": gossip_sigs, "latest_finalized": make_checkpoint(root_seed=255, slot=5), } @@ -325,14 +260,14 @@ def test_mixed_stale_and_fresh_entries(pruning_store: Store) -> None: # Verify all data exists before pruning for att in attestations: - assert att.data_root_bytes() in store.attestation_data_by_root + assert att in store.gossip_signatures pruned_store = store.prune_stale_attestation_data() # Entries with target.slot <= 5 should be pruned (slots 1-5) for att in attestations[:5]: - assert att.data_root_bytes() not in pruned_store.attestation_data_by_root + assert att not in pruned_store.gossip_signatures # Entries with target.slot > 5 should be kept (slots 6-10) for att in attestations[5:]: - assert att.data_root_bytes() in pruned_store.attestation_data_by_root + assert att in pruned_store.gossip_signatures diff --git a/tests/lean_spec/subspecs/forkchoice/test_validator.py b/tests/lean_spec/subspecs/forkchoice/test_validator.py index a2f3f66d..bda94cd4 100644 --- a/tests/lean_spec/subspecs/forkchoice/test_validator.py +++ b/tests/lean_spec/subspecs/forkchoice/test_validator.py @@ -15,9 +15,9 @@ ValidatorIndex, ) from lean_spec.subspecs.containers.slot import Slot -from lean_spec.subspecs.forkchoice import Store +from lean_spec.subspecs.forkchoice import GossipSignatureEntry, Store from lean_spec.subspecs.ssz.hash import hash_tree_root -from lean_spec.subspecs.xmss.aggregation import SignatureKey +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types import Bytes32, Uint64 from tests.lean_spec.helpers import TEST_VALIDATOR_ID, make_aggregated_proof, make_store @@ -82,29 +82,27 @@ def test_produce_block_with_attestations( signature=key_manager.sign_attestation_data(ValidatorIndex(6), data_6), ) - data_root_5 = signed_5.data.data_root_bytes() - data_root_6 = signed_6.data.data_root_bytes() - proof_5 = make_aggregated_proof(key_manager, [ValidatorIndex(5)], signed_5.data) proof_6 = make_aggregated_proof(key_manager, [ValidatorIndex(6)], signed_6.data) - sig_key_5 = SignatureKey(ValidatorIndex(5), data_root_5) - sig_key_6 = SignatureKey(ValidatorIndex(6), data_root_6) + # Build payloads keyed by attestation data. + # If data_5 == data_6 (same slot/head/target/source), they share a key. + known_payloads: dict[AttestationData, set[AggregatedSignatureProof]] = {} + known_payloads.setdefault(signed_5.data, set()).add(proof_5) + known_payloads.setdefault(signed_6.data, set()).add(proof_6) + + gossip_sigs = {} + gossip_sigs.setdefault(signed_5.data, set()).add( + GossipSignatureEntry(ValidatorIndex(5), signed_5.signature) + ) + gossip_sigs.setdefault(signed_6.data, set()).add( + GossipSignatureEntry(ValidatorIndex(6), signed_6.signature) + ) sample_store = sample_store.model_copy( update={ - "latest_known_aggregated_payloads": { - sig_key_5: [proof_5], - sig_key_6: [proof_6], - }, - "attestation_data_by_root": { - data_root_5: signed_5.data, - data_root_6: signed_6.data, - }, - "gossip_signatures": { - sig_key_5: signed_5.signature, - sig_key_6: signed_6.signature, - }, + "latest_known_aggregated_payloads": known_payloads, + "gossip_signatures": gossip_sigs, } ) @@ -183,7 +181,6 @@ def test_produce_block_empty_attestations(self, sample_store: Store) -> None: sample_store = sample_store.model_copy( update={ "latest_known_aggregated_payloads": {}, - "attestation_data_by_root": {}, } ) @@ -220,15 +217,14 @@ def test_produce_block_state_consistency( signature=key_manager.sign_attestation_data(ValidatorIndex(7), data_7), ) - data_root_7 = signed_7.data.data_root_bytes() proof_7 = make_aggregated_proof(key_manager, [ValidatorIndex(7)], signed_7.data) - sig_key_7 = SignatureKey(ValidatorIndex(7), data_root_7) sample_store = sample_store.model_copy( update={ - "latest_known_aggregated_payloads": {sig_key_7: [proof_7]}, - "attestation_data_by_root": {data_root_7: signed_7.data}, - "gossip_signatures": {sig_key_7: signed_7.signature}, + "latest_known_aggregated_payloads": {signed_7.data: {proof_7}}, + "gossip_signatures": { + signed_7.data: {GossipSignatureEntry(ValidatorIndex(7), signed_7.signature)}, + }, } ) diff --git a/tests/lean_spec/subspecs/validator/test_service.py b/tests/lean_spec/subspecs/validator/test_service.py index 7886f9ab..7019db76 100644 --- a/tests/lean_spec/subspecs/validator/test_service.py +++ b/tests/lean_spec/subspecs/validator/test_service.py @@ -27,7 +27,7 @@ from lean_spec.subspecs.validator import ValidatorRegistry, ValidatorService from lean_spec.subspecs.validator.registry import ValidatorEntry from lean_spec.subspecs.xmss import TARGET_SIGNATURE_SCHEME -from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof, SignatureKey +from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof from lean_spec.types import Bytes32, Uint64 from tests.lean_spec.helpers import TEST_VALIDATOR_ID, MockNetworkRequester, make_store @@ -721,12 +721,11 @@ async def test_block_includes_pending_attestations( slot=attestation_data.slot, ) - aggregated_payloads = {SignatureKey(vid, data_root): [proof] for vid in participants} + aggregated_payloads = {attestation_data: {proof}} - # Update store with attestation data and aggregated payloads + # Update store with aggregated payloads updated_store = store.model_copy( update={ - "attestation_data_by_root": {data_root: attestation_data}, "latest_known_aggregated_payloads": aggregated_payloads, } )