state: fix bug about duplicate roots in root to slots map (#296)

Author: tcoratger

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

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)