diff --git a/src/lean_spec/spec/forks/lstar/fork_choice.py b/src/lean_spec/spec/forks/lstar/fork_choice.py index deaa3ed13..c0ce86ffd 100644 --- a/src/lean_spec/spec/forks/lstar/fork_choice.py +++ b/src/lean_spec/spec/forks/lstar/fork_choice.py @@ -102,9 +102,9 @@ def create_store( # type: ignore[override] # ty: ignore[invalid-method-overrid 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 + An attestation becomes stale when its head checkpoint falls at or before the finalized slot. Such attestations cannot affect chain selection since - the target is already finalized. + the head is already finalized. Pruning removes all attestation-related data: @@ -114,22 +114,22 @@ def prune_stale_attestation_data(self, store: LstarStore) -> LstarStore: """ # Filter out stale entries from all attestation-related mappings. # - # Each mapping is keyed by attestation data, so we check membership by slot - # against the finalized slot. + # Each mapping is keyed by attestation data, so we check the attested + # head slot against the finalized slot. store.attestation_signatures = { attestation_data: signatures for attestation_data, signatures in store.attestation_signatures.items() - if attestation_data.target.slot > store.latest_finalized.slot + if attestation_data.head.slot > store.latest_finalized.slot } store.latest_new_aggregated_payloads = { attestation_data: proofs for attestation_data, proofs in store.latest_new_aggregated_payloads.items() - if attestation_data.target.slot > store.latest_finalized.slot + if attestation_data.head.slot > store.latest_finalized.slot } store.latest_known_aggregated_payloads = { attestation_data: proofs for attestation_data, proofs in store.latest_known_aggregated_payloads.items() - if attestation_data.target.slot > store.latest_finalized.slot + if attestation_data.head.slot > store.latest_finalized.slot } return store @@ -467,7 +467,12 @@ def compute_block_weights(self, store: LstarStore) -> dict[Bytes32, int]: for every ancestor above the finalized slot. """ attestations = self.extract_attestations_from_aggregated_payloads( - store, store.latest_known_aggregated_payloads + store, + { + attestation_data: proofs + for attestation_data, proofs in store.latest_known_aggregated_payloads.items() + if attestation_data.head.slot > store.latest_finalized.slot + }, ) weights = self._accumulate_ancestor_weights( @@ -547,9 +552,16 @@ def update_head(self, store: LstarStore) -> LstarStore: 1. Latest justified checkpoint as the starting root 2. LMD-GHOST fork choice rule (heaviest subtree by attestation weight) """ - # Extract attestations from known aggregated payloads + # Extract fork-choice relevant attestations from known aggregated payloads. + # A vote whose head is at or below finalization has no weight above the + # finalized boundary. attestations = self.extract_attestations_from_aggregated_payloads( - store, store.latest_known_aggregated_payloads + store, + { + attestation_data: proofs + for attestation_data, proofs in store.latest_known_aggregated_payloads.items() + if attestation_data.head.slot > store.latest_finalized.slot + }, ) # Run LMD-GHOST fork choice algorithm. @@ -636,7 +648,11 @@ def update_safe_target(self, store: LstarStore) -> LstarStore: # "Known" is excluded by design. attestations = self.extract_attestations_from_aggregated_payloads( store, - store.latest_new_aggregated_payloads, + { + attestation_data: proofs + for attestation_data, proofs in store.latest_new_aggregated_payloads.items() + if attestation_data.head.slot > store.latest_finalized.slot + }, ) # Run LMD GHOST with the supermajority threshold. diff --git a/tests/lean_spec/spec/forks/lstar/forkchoice/test_compute_block_weights.py b/tests/lean_spec/spec/forks/lstar/forkchoice/test_compute_block_weights.py index 732263826..d095a17e2 100644 --- a/tests/lean_spec/spec/forks/lstar/forkchoice/test_compute_block_weights.py +++ b/tests/lean_spec/spec/forks/lstar/forkchoice/test_compute_block_weights.py @@ -22,6 +22,89 @@ def _make_empty_proof(participants: list[ValidatorIndex]) -> SingleMessageAggreg ) +def _add_two_block_chain(base_store: Store) -> tuple[Bytes32, Bytes32, Bytes32]: + genesis_root = base_store.head + + block1 = make_signed_block( + slot=Slot(1), + proposer_index=ValidatorIndex(0), + parent_root=genesis_root, + state_root=make_bytes32(10), + ) + block1_root = hash_tree_root(block1.block) + + block2 = make_signed_block( + slot=Slot(2), + proposer_index=ValidatorIndex(1), + parent_root=block1_root, + state_root=make_bytes32(20), + ) + block2_root = hash_tree_root(block2.block) + + base_store.blocks = { + **base_store.blocks, + block1_root: block1.block, + block2_root: block2.block, + } + genesis_state = base_store.states[genesis_root] + base_store.states = { + **base_store.states, + block1_root: genesis_state, + block2_root: genesis_state, + } + base_store.head = block2_root + + return genesis_root, block1_root, block2_root + + +def _add_finalized_fork(base_store: Store) -> tuple[Bytes32, Bytes32, Bytes32, Bytes32]: + genesis_root = base_store.head + + block1 = make_signed_block( + slot=Slot(1), + proposer_index=ValidatorIndex(0), + parent_root=genesis_root, + state_root=make_bytes32(10), + ) + block1_root = hash_tree_root(block1.block) + + child_a = make_signed_block( + slot=Slot(2), + proposer_index=ValidatorIndex(1), + parent_root=block1_root, + state_root=make_bytes32(20), + ) + child_a_root = hash_tree_root(child_a.block) + + child_b = make_signed_block( + slot=Slot(2), + proposer_index=ValidatorIndex(2), + parent_root=block1_root, + state_root=make_bytes32(30), + ) + child_b_root = hash_tree_root(child_b.block) + + voted_root, tie_break_root = sorted([child_a_root, child_b_root]) + + genesis_state = base_store.states[genesis_root] + base_store.blocks = { + **base_store.blocks, + block1_root: block1.block, + child_a_root: child_a.block, + child_b_root: child_b.block, + } + base_store.states = { + **base_store.states, + block1_root: genesis_state, + child_a_root: genesis_state, + child_b_root: genesis_state, + } + base_store.head = tie_break_root + assert tie_break_root > voted_root + + return genesis_root, block1_root, voted_root, tie_break_root + + def test_genesis_only_store_returns_empty_weights(spec: LstarSpec, base_store: Store) -> None: """A genesis-only store with no attestations has no block weights.""" assert spec.compute_block_weights(base_store) == {} @@ -81,6 +164,131 @@ def test_linear_chain_weight_accumulates_upward(spec: LstarSpec, base_store: Sto assert weights == {block2_root: 1, block1_root: 1} +def test_stale_latest_message_does_not_mask_fresh_weight( + spec: LstarSpec, base_store: Store +) -> None: + """A stale post-finalization message must not hide a validator's fresh vote.""" + genesis_root = base_store.head + + block1 = make_signed_block( + slot=Slot(1), + proposer_index=ValidatorIndex(0), + parent_root=genesis_root, + state_root=make_bytes32(10), + ) + block1_root = hash_tree_root(block1.block) + + block2 = make_signed_block( + slot=Slot(2), + proposer_index=ValidatorIndex(1), + parent_root=block1_root, + state_root=make_bytes32(20), + ) + block2_root = hash_tree_root(block2.block) + + base_store.blocks = { + **base_store.blocks, + block1_root: block1.block, + block2_root: block2.block, + } + genesis_state = base_store.states[genesis_root] + base_store.states = { + **base_store.states, + block1_root: genesis_state, + block2_root: genesis_state, + } + base_store.head = block2_root + base_store.latest_finalized = Checkpoint(root=block1_root, slot=Slot(1)) + + fresh_vote = AttestationData( + slot=Slot(2), + head=Checkpoint(root=block2_root, slot=Slot(2)), + target=Checkpoint(root=block2_root, slot=Slot(2)), + source=Checkpoint(root=block1_root, slot=Slot(1)), + ) + stale_vote = AttestationData( + slot=Slot(3), + head=Checkpoint(root=block1_root, slot=Slot(1)), + target=Checkpoint(root=block1_root, slot=Slot(1)), + source=Checkpoint(root=genesis_root, slot=Slot(0)), + ) + + base_store.latest_known_aggregated_payloads = { + fresh_vote: {_make_empty_proof([ValidatorIndex(0)])}, + stale_vote: {_make_empty_proof([ValidatorIndex(0)])}, + } + + weights = spec.compute_block_weights(base_store) + + assert weights == {block2_root: 1} + + +def test_fresh_head_with_finalized_target_still_counts(spec: LstarSpec, base_store: Store) -> None: + """A vote with head above finalization still counts when its target is finalized.""" + genesis_root, block1_root, block2_root = _add_two_block_chain(base_store) + base_store.latest_finalized = Checkpoint(root=block1_root, slot=Slot(1)) + + fresh_vote_with_finalized_target = AttestationData( + slot=Slot(2), + head=Checkpoint(root=block2_root, slot=Slot(2)), + target=Checkpoint(root=block1_root, slot=Slot(1)), + source=Checkpoint(root=genesis_root, slot=Slot(0)), + ) + base_store.latest_known_aggregated_payloads = { + fresh_vote_with_finalized_target: {_make_empty_proof([ValidatorIndex(0)])}, + } + + weights = spec.compute_block_weights(base_store) + + assert weights == {block2_root: 1} + + +def test_finalized_target_vote_can_drive_head_selection(spec: LstarSpec, base_store: Store) -> None: + """A weighted head above finalization wins over the zero-weight tie break branch.""" + genesis_root, block1_root, voted_root, tie_break_root = _add_finalized_fork(base_store) + finalized = Checkpoint(root=block1_root, slot=Slot(1)) + base_store.latest_justified = finalized + base_store.latest_finalized = finalized + + vote = AttestationData( + slot=Slot(2), + head=Checkpoint(root=voted_root, slot=Slot(2)), + target=finalized, + source=Checkpoint(root=genesis_root, slot=Slot(0)), + ) + base_store.latest_known_aggregated_payloads = { + vote: {_make_empty_proof([ValidatorIndex(0)])}, + } + + store = spec.update_head(base_store) + + assert tie_break_root != voted_root + assert store.head == voted_root + + +def test_finalized_target_vote_can_advance_safe_target(spec: LstarSpec, base_store: Store) -> None: + """Safe target counts new-pool votes whose head is above finalization.""" + genesis_root, block1_root, block2_root = _add_two_block_chain(base_store) + finalized = Checkpoint(root=block1_root, slot=Slot(1)) + base_store.latest_justified = finalized + base_store.latest_finalized = finalized + base_store.safe_target = block1_root + + vote = AttestationData( + slot=Slot(2), + head=Checkpoint(root=block2_root, slot=Slot(2)), + target=finalized, + source=Checkpoint(root=genesis_root, slot=Slot(0)), + ) + base_store.latest_new_aggregated_payloads = { + vote: {_make_empty_proof([ValidatorIndex(0), ValidatorIndex(1)])}, + } + + store = spec.update_safe_target(base_store) + + assert store.safe_target == block2_root + + def test_multiple_attestations_accumulate(spec: LstarSpec, base_store: Store) -> None: """Multiple validators attesting to the same head accumulate weight.""" genesis_root = base_store.head diff --git a/tests/lean_spec/spec/forks/lstar/forkchoice/test_store_pruning.py b/tests/lean_spec/spec/forks/lstar/forkchoice/test_store_pruning.py index d4112f482..0ea630cd2 100644 --- a/tests/lean_spec/spec/forks/lstar/forkchoice/test_store_pruning.py +++ b/tests/lean_spec/spec/forks/lstar/forkchoice/test_store_pruning.py @@ -2,7 +2,7 @@ from lean_spec.spec.forks import AggregationBits, Slot, ValidatorIndex from lean_spec.spec.forks.lstar import AttestationSignatureEntry, Store -from lean_spec.spec.forks.lstar.containers import SingleMessageAggregate +from lean_spec.spec.forks.lstar.containers import AttestationData, SingleMessageAggregate from lean_spec.spec.forks.lstar.spec import LstarSpec from lean_spec.spec.ssz import ByteList512KiB, Bytes32 from tests.lean_spec.helpers import ( @@ -13,11 +13,11 @@ ) -def test_prunes_entries_with_target_at_finalized(spec: LstarSpec, pruning_store: Store) -> None: - """Verify entries with target.slot == finalized slot are pruned.""" +def test_prunes_entries_with_head_at_finalized(spec: LstarSpec, pruning_store: Store) -> None: + """Verify entries with head.slot == finalized slot are pruned.""" store = pruning_store - # Create attestation data with target.slot == 5 + # Create attestation data with head.slot == 5 attestation_data = make_attestation_data( slot=Slot(5), target_slot=Slot(5), @@ -35,17 +35,17 @@ def test_prunes_entries_with_target_at_finalized(spec: LstarSpec, pruning_store: # Verify data exists before pruning assert attestation_data in store.attestation_signatures - # Prune should remove entries where target.slot <= finalized.slot + # Prune should remove entries where head.slot <= finalized.slot pruned_store = spec.prune_stale_attestation_data(store) assert attestation_data not in pruned_store.attestation_signatures -def test_prunes_entries_with_target_before_finalized(spec: LstarSpec, pruning_store: Store) -> None: - """Verify entries with target.slot < finalized slot are pruned.""" +def test_prunes_entries_with_head_before_finalized(spec: LstarSpec, pruning_store: Store) -> None: + """Verify entries with head.slot < finalized slot are pruned.""" store = pruning_store - # Create attestation data with target.slot == 3 + # Create attestation data with head.slot == 3 attestation_data = make_attestation_data( slot=Slot(3), target_slot=Slot(3), @@ -54,7 +54,7 @@ def test_prunes_entries_with_target_before_finalized(spec: LstarSpec, pruning_st source_root=Bytes32.zero(), ) - # Set up store with finalized slot at 5 (greater than target.slot) + # Set up store with finalized slot at 5 (greater than head.slot) store.attestation_signatures = { attestation_data: {AttestationSignatureEntry(ValidatorIndex(1), make_mock_signature())}, } @@ -63,17 +63,17 @@ def test_prunes_entries_with_target_before_finalized(spec: LstarSpec, pruning_st # Verify data exists before pruning assert attestation_data in store.attestation_signatures - # Prune should remove entries where target.slot <= finalized.slot + # Prune should remove entries where head.slot <= finalized.slot pruned_store = spec.prune_stale_attestation_data(store) assert attestation_data not in pruned_store.attestation_signatures -def test_keeps_entries_with_target_after_finalized(spec: LstarSpec, pruning_store: Store) -> None: - """Verify entries with target.slot > finalized slot are kept.""" +def test_keeps_entries_with_head_after_finalized(spec: LstarSpec, pruning_store: Store) -> None: + """Verify entries with head.slot > finalized slot are kept.""" store = pruning_store - # Create attestation data with target.slot == 10 + # Create attestation data with head.slot == 10 attestation_data = make_attestation_data( slot=Slot(10), target_slot=Slot(10), @@ -82,7 +82,7 @@ def test_keeps_entries_with_target_after_finalized(spec: LstarSpec, pruning_stor source_root=make_bytes32(2), ) - # Set up store with finalized slot at 5 (less than target.slot) + # Set up store with finalized slot at 5 (less than head.slot) store.attestation_signatures = { attestation_data: {AttestationSignatureEntry(ValidatorIndex(1), make_mock_signature())}, } @@ -91,12 +91,45 @@ def test_keeps_entries_with_target_after_finalized(spec: LstarSpec, pruning_stor # Verify data exists before pruning assert attestation_data in store.attestation_signatures - # Prune should keep entries where target.slot > finalized.slot + # Prune should keep entries where head.slot > finalized.slot pruned_store = spec.prune_stale_attestation_data(store) assert attestation_data in pruned_store.attestation_signatures +def test_keeps_finalized_target_when_head_after_finalized( + spec: LstarSpec, pruning_store: Store +) -> None: + """A finalized target is still useful while the attested head is unfinalized.""" + store = pruning_store + finalized = make_checkpoint(root_seed=1, slot=5) + head = make_checkpoint(root_seed=2, slot=6) + attestation_data = AttestationData( + slot=Slot(6), + head=head, + target=finalized, + source=make_checkpoint(root_seed=0, slot=0), + ) + + placeholder = ByteList512KiB(data=b"") + mock_proof = SingleMessageAggregate( + participants=AggregationBits.from_indices([ValidatorIndex(1)]), + proof=placeholder, + ) + signature = AttestationSignatureEntry(ValidatorIndex(1), make_mock_signature()) + + store.attestation_signatures = {attestation_data: {signature}} + store.latest_new_aggregated_payloads = {attestation_data: {mock_proof}} + store.latest_known_aggregated_payloads = {attestation_data: {mock_proof}} + store.latest_finalized = finalized + + pruned_store = spec.prune_stale_attestation_data(store) + + assert attestation_data in pruned_store.attestation_signatures + assert attestation_data in pruned_store.latest_new_aggregated_payloads + assert attestation_data in pruned_store.latest_known_aggregated_payloads + + def test_prunes_related_structures_together(spec: LstarSpec, pruning_store: Store) -> None: """Verify all three data structures are pruned atomically.""" store = pruning_store @@ -240,10 +273,10 @@ def test_mixed_stale_and_fresh_entries(spec: LstarSpec, pruning_store: Store) -> pruned_store = spec.prune_stale_attestation_data(store) - # Entries with target.slot <= 5 should be pruned (slots 1-5) + # Entries with head.slot <= 5 should be pruned (slots 1-5) for attestation in attestations[:5]: assert attestation not in pruned_store.attestation_signatures - # Entries with target.slot > 5 should be kept (slots 6-10) + # Entries with head.slot > 5 should be kept (slots 6-10) for attestation in attestations[5:]: assert attestation in pruned_store.attestation_signatures