diff --git a/src/lean_spec/subspecs/containers/state/state.py b/src/lean_spec/subspecs/containers/state/state.py index e34e317d..6e3c5668 100644 --- a/src/lean_spec/subspecs/containers/state/state.py +++ b/src/lean_spec/subspecs/containers/state/state.py @@ -421,13 +421,18 @@ def process_attestations( finalized_slot = latest_finalized.slot justified_slots = self.justified_slots - # Map roots to their slots for pruning when finalization advances. - # Only track roots after the finalized boundary; earlier roots are pruned. + # Map roots to all slots where they appear. + # + # Missed slots produce duplicate zero hashes in history. + # During pruning, we must check all slot occurrences for each root. + # A single-slot mapping would fail when iterating over slots. start_slot = int(finalized_slot) + 1 - root_to_slot = { - self.historical_block_hashes[i]: Slot(i) - for i in range(start_slot, len(self.historical_block_hashes)) - } + root_to_slots: dict[Bytes32, list[Slot]] = {} + for i in range(start_slot, len(self.historical_block_hashes)): + root = self.historical_block_hashes[i] + if root not in root_to_slots: + root_to_slots[root] = [] + root_to_slots[root].append(Slot(i)) # Process each attestation independently # @@ -550,17 +555,19 @@ def process_attestations( # Rebase/prune justification tracking across the new finalized boundary. # - # The state stores `justified_slots` starting at (finalized_slot + 1), + # The state stores justified slot flags starting at (finalized_slot + 1), # so when finalization advances by `delta`, we drop the first `delta` bits. # - # We also prune any pending `justifications` whose slots are now finalized. + # We also prune any pending justifications whose slots are now finalized. + # A root may appear at multiple slots; keep the justification if ANY + # slot for that root is still unfinalized (conservative approach). delta = int(finalized_slot - old_finalized_slot) if delta > 0: justified_slots = justified_slots.shift_window(delta) justifications = { root: votes for root, votes in justifications.items() - if root_to_slot.get(root, Slot(0)) > finalized_slot + if any(slot > finalized_slot for slot in root_to_slots.get(root, [])) } # Convert the vote structure back into SSZ format diff --git a/tests/lean_spec/subspecs/containers/test_state_justified_slots.py b/tests/lean_spec/subspecs/containers/test_state_justified_slots.py index 4e4d4cb9..2169e11b 100644 --- a/tests/lean_spec/subspecs/containers/test_state_justified_slots.py +++ b/tests/lean_spec/subspecs/containers/test_state_justified_slots.py @@ -12,7 +12,12 @@ from lean_spec.subspecs.containers.checkpoint import Checkpoint from lean_spec.subspecs.containers.slot import Slot from lean_spec.subspecs.containers.state import State -from lean_spec.types import Uint64 +from lean_spec.subspecs.containers.state.types import ( + HistoricalBlockHashes, + JustificationRoots, + JustificationValidators, +) +from lean_spec.types import ZERO_HASH, Boolean, Uint64 from tests.lean_spec.helpers import make_aggregated_attestation, make_block, make_validators @@ -97,3 +102,126 @@ def test_is_slot_justified_raises_on_out_of_bounds() -> None: State.generate_genesis(Uint64(0), make_validators(1)).justified_slots.is_slot_justified( Slot(0), Slot(1) ) + + +def test_duplicate_roots_in_root_to_slots_mapping() -> None: + """ + Verify duplicate block roots are tracked correctly for pruning decisions. + + Missed slots produce empty block hashes (zeros). + Multiple missed slots create duplicate entries in the history. + + When finalization advances, pending justifications must be pruned. + The pruning logic needs to know which slots each root appears at. + + The root-to-slots mapping must store all slots where each root appears. + Otherwise, iteration during pruning fails. + + Test strategy: + + 1. Build a chain with zeros at two slots (simulating missed blocks) + 2. Add a pending justification that should survive pruning + 3. Trigger finalization to run the pruning logic + 4. Verify the pending justification survives correctly + """ + # Two of three validators form a supermajority. + state = State.generate_genesis(genesis_time=Uint64(0), validators=make_validators(3)) + + # Phase 1: Build a chain and justify slot 1. + # + # We need an existing justified checkpoint before we can test pruning. + + state = state.process_slots(Slot(1)) + block_1 = make_block(state, Slot(1), attestations=[]) + state = state.process_block(block_1) + + state = state.process_slots(Slot(2)) + block_2 = make_block(state, Slot(2), attestations=[]) + source_0 = Checkpoint(root=block_1.parent_root, slot=Slot(0)) + target_1 = Checkpoint(root=block_2.parent_root, slot=Slot(1)) + att_0_to_1 = make_aggregated_attestation( + participant_ids=[0, 1], + attestation_slot=Slot(2), + source=source_0, + target=target_1, + ) + block_2 = make_block(state, Slot(2), attestations=[att_0_to_1]) + state = state.process_block(block_2) + + assert state.latest_finalized.slot == Slot(0) + assert state.latest_justified.slot == Slot(1) + + # Phase 2: Extend chain to populate more history entries. + # + # We need enough slots to inject duplicate roots later. + + state = state.process_slots(Slot(3)) + block_3 = make_block(state, Slot(3), attestations=[]) + state = state.process_block(block_3) + + state = state.process_slots(Slot(4)) + block_4 = make_block(state, Slot(4), attestations=[]) + state = state.process_block(block_4) + + state = state.process_slots(Slot(5)) + block_5 = make_block(state, Slot(5), attestations=[]) + state = state.process_block_header(block_5) + + # Phase 3: Inject duplicate roots to simulate missed blocks. + # + # Missed blocks leave zeros in the history. + # Multiple missed blocks create the same root at different slots. + # The pruning logic must handle this case correctly. + + slot_3_root = state.historical_block_hashes[3] + modified_hashes = list(state.historical_block_hashes.data) + modified_hashes[2] = ZERO_HASH + modified_hashes[4] = ZERO_HASH + + # Register a pending justification for slot 3. + # + # This justification should survive pruning because slot 3 + # comes after the finalized boundary. + pending_votes = [Boolean(True), Boolean(False), Boolean(False)] + + state = state.model_copy( + update={ + "historical_block_hashes": HistoricalBlockHashes(data=modified_hashes), + "justifications_roots": JustificationRoots(data=[slot_3_root]), + "justifications_validators": JustificationValidators(data=pending_votes), + } + ) + + # Sanity check: zeros at slots 2 and 4, real root at slot 3. + assert state.historical_block_hashes[2] == ZERO_HASH + assert state.historical_block_hashes[4] == ZERO_HASH + assert state.historical_block_hashes[3] == slot_3_root + + # Phase 4: Trigger finalization to exercise pruning. + # + # This attestation justifies slot 2 and finalizes slot 1. + # Finalization triggers pruning of stale justifications. + + source_1 = Checkpoint(root=state.historical_block_hashes[1], slot=Slot(1)) + target_2 = Checkpoint(root=ZERO_HASH, slot=Slot(2)) + att_1_to_2 = make_aggregated_attestation( + participant_ids=[0, 1], + attestation_slot=Slot(5), + source=source_1, + target=target_2, + ) + + # Processing this attestation runs the pruning logic. + # + # Pruning iterates over all slots for each root in history. + # Duplicate roots must map to multiple slots, not just one. + state = state.process_attestations([att_1_to_2]) + + # Verify finalization succeeded. + assert state.latest_finalized.slot == Slot(1) + assert state.latest_justified.slot == Slot(2) + + # The pending justification for slot 3 must survive. + # + # Slot 3 is beyond the finalized boundary, so pruning keeps it. + assert slot_3_root in list(state.justifications_roots)