fix: align process_attestations() with intended specs (leanEthereum#160)

* fix: align process_attestations() with intended specs

* refactor: justifications processing becomes pure functions

* fix: slot typing

* fix: in progress fix of justification reorg test

* fix: stray code must go

* refactor: align justifications with leanEthereum#165

* fix: linting

* fix: update code comments

* fix: add extra attestations

* fix: linting

* fix: linting

* fix: type count

* Update packages/testing/src/consensus_testing/test_types/store_checks.py

Co-authored-by: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>

* fix: docstring

* fix: lowercase justified

---------

Co-authored-by: Thomas Coratger <60488569+tcoratger@users.noreply.github.com>

Author: unnawut

packages/testing/src/consensus_testing/test_types/store_checks.py

@@ -128,9 +128,6 @@ class StoreChecks(CamelModel):
     latest_justified_root: Bytes32 | None = None
     """Expected latest justified checkpoint root."""
 
-    latest_finalized_slot: Slot | None = None
-    """Expected latest finalized checkpoint slot."""
-
     latest_justified_root_label: str | None = None
     """
     Expected latest justified checkpoint root by label reference.
@@ -140,9 +137,21 @@ class StoreChecks(CamelModel):
     and validate the latest justified checkpointroot matches.
     """
 
+    latest_finalized_slot: Slot | None = None
+    """Expected latest finalized checkpoint slot."""
+
     latest_finalized_root: Bytes32 | None = None
     """Expected latest finalized checkpoint root."""
 
+    latest_finalized_root_label: str | None = None
+    """
+    Expected latest finalized checkpoint root by label reference.
+
+    Alternative to `latest_justified_root` that uses the block label system.
+    The framework will resolve this label to the actual block root
+    and validate the latest finalized checkpoint root matches.
+    """
+
     safe_target: Bytes32 | None = None
     """Expected safe target root."""
 

src/lean_spec/subspecs/containers/state/state.py

@@ -22,6 +22,7 @@
 from ..checkpoint import Checkpoint
 from ..config import Config
 from ..slot import Slot
+from .helpers import flatten_justifications_map, get_justifications_map
 from .types import (
     HistoricalBlockHashes,
     JustificationRoots,
@@ -336,61 +337,97 @@ def process_attestations(
         State
             A new state with updated justification/finalization.
         """
-        # Start with current justifications and finalization state.
-        justified_slots = list(self.justified_slots)
+        # Get justifications, justified slots and historical block hashes are already up to
+        # date as per the processing in process_block_header
+        justifications = get_justifications_map(
+            justifications_roots=self.justifications_roots,
+            justifications_validators=self.justifications_validators,
+            validator_count=self.validators.count,
+        )
+
+        # Track state changes to be applied at the end
         latest_justified = self.latest_justified
         latest_finalized = self.latest_finalized
+        justified_slots = list(self.justified_slots)
 
         # Process each attestation in the block.
         for attestation in attestations:
             attestation_data = attestation.data
             source = attestation_data.source
             target = attestation_data.target
 
-            # Validate that this is a reasonable attestation (source comes before target).
-            if source.slot.as_int() >= target.slot.as_int():
-                continue  # Skip invalid attestation
-
-            # Check if source checkpoint is justified.
-            source_slot_int = source.slot.as_int()
-            target_slot_int = target.slot.as_int()
-
-            # Ensure we have enough justified slots history.
-            if source_slot_int < len(justified_slots):
-                source_is_justified = justified_slots[source_slot_int]
-            else:
-                continue  # Source is too far in the past
-
-            # If source is justified, consider justifying the target.
-            if (
-                source_is_justified
-                and target_slot_int < len(justified_slots)
-                and justified_slots[target_slot_int]
-            ):
-                # Target is already justified, check for finalization.
-                if (
-                    source.slot.as_int() + 1 == target.slot.as_int()
-                    and latest_justified.slot.as_int() < target.slot.as_int()
+            # Ignore attestations whose source is not already justified,
+            # or whose target is not in the history, or whose target is not a
+            # valid justifiable slot
+            source_slot = source.slot.as_int()
+            target_slot = target.slot.as_int()
+
+            # Source slot must be justified
+            if not justified_slots[source_slot]:
+                continue
+
+            # Target slot must not be already justified
+            # This condition is missing in 3sf mini but has been added here because
+            # we don't want to re-introduce the target again for remaining votes if
+            # the slot is already justified and its tracking already cleared out
+            # from justifications map
+            if justified_slots[target_slot]:
+                continue
+
+            # Source root must match the state's historical block hashes
+            if source.root != self.historical_block_hashes[source_slot]:
+                continue
+
+            # Target root must match the state's historical block hashes
+            if target.root != self.historical_block_hashes[target_slot]:
+                continue
+
+            # Target slot must be after source slot
+            if target.slot <= source.slot:
+                continue
+
+            # Target slot must be justifiable after the latest finalized slot
+            if not target.slot.is_justifiable_after(self.latest_finalized.slot):
+                continue
+
+            # Track attempts to justify new hashes
+            if target.root not in justifications:
+                justifications[target.root] = [Boolean(False)] * self.validators.count
+
+            validator_id = attestation.validator_id.as_int()
+            if not justifications[target.root][validator_id]:
+                justifications[target.root][validator_id] = Boolean(True)
+
+            count = sum(bool(justified) for justified in justifications[target.root])
+
+            # If 2/3 attested to the same new valid hash to justify
+            # in 3sf mini this is strict equality, but we have updated it to >=
+            # also have modified it from count >= (2 * state.config.num_validators) // 3
+            # to prevent integer division which could lead to less than 2/3 of validators
+            # justifying specially if the num_validators is low in testing scenarios
+            if 3 * count >= (2 * self.validators.count):
+                latest_justified = target
+                justified_slots[target_slot] = True
+                del justifications[target.root]
+
+                # Finalization: if the target is the next valid justifiable
+                # hash after the source
+                if not any(
+                    Slot(slot).is_justifiable_after(self.latest_finalized.slot)
+                    for slot in range(source_slot + 1, target_slot)
                 ):
-                    # Consecutive justified checkpoints -> finalize the source.
                     latest_finalized = source
-                    latest_justified = target
-
-            else:
-                # Try to justify the target if source is justified.
-                if source_is_justified:
-                    # Ensure justified_slots is long enough, then mark the target slot.
-                    while len(justified_slots) <= target_slot_int:
-                        justified_slots.append(Boolean(False))
-                    justified_slots[target_slot_int] = Boolean(True)
 
-                    # Update latest_justified if this target is newer.
-                    if target.slot.as_int() > latest_justified.slot.as_int():
-                        latest_justified = target
+        # Flatten and set updated justifications back to the state
+        justifications_roots, justifications_validators = flatten_justifications_map(
+            justifications, self.validators.count
+        )
 
         # Return the updated state.
         return self.model_copy(
             update={
+                "justifications_roots": justifications_roots,
+                "justifications_validators": justifications_validators,
                 "justified_slots": self.justified_slots.__class__(data=justified_slots),
                 "latest_justified": latest_justified,
                 "latest_finalized": latest_finalized,

tests/consensus/devnet/fc/test_fork_choice_reorgs.py

@@ -777,22 +777,25 @@ def test_reorg_on_newly_justified_slot(
 
     Scenario
     --------
-    Two forks compete across multiple justifiable slots. Fork choice must
-    correctly handle reorgs while respecting justification rules.
+    Two forks compete. Fork A is heavier and longer, but Fork B manages to
+    become justified. Fork choice must switch to the justified fork regardless
+    of weight/length.
 
     - Slot 1: Base
-    - Slot 2: Fork A (competing with B)
-    - Slot 2: Fork B (competing with A)
-    - Slot 3: Fork A (Fork A now has depth 2 - becomes head)
-    - Slot 4: Fork A (Fork A now has depth 3 - still head)
-    - Slot 5: Fork B with enough justifications for Slot 2 → triggers reorg → Fork B becomes head
+    - Slots 2-4: Fork A extends (becomes head with depth 3)
+    - Slot 5: Fork B appears (descending from Base, skipping slots 2-4)
+    - Slot 6: Fork B extends. This block contains enough attestations to
+              justify Fork B at Slot 5.
 
     Expected Behavior
     -----------------
-    1. Fork A and Fork B initially have equal weight at Slot 2
-    2. Fork A takes lead at Slot 3 and Slot 4
-    2. At Slot 5, enough attestations justified Fork B at Slot 2, Fork A permanently non-canonical
-    3. Fork B becomes head at Slot 5 due to justification of Fork B at Slot 2
+    1. Fork A takes the lead initially (Slots 2-4) as the heaviest chain.
+    2. Fork B appears at Slot 5 but is initially lighter.
+    3. At Slot 6, the new block includes attestations that justify Fork B at Slot 5.
+    4. The justified checkpoint updates to Slot 5 (fork_b_1).
+    5. Fork A is immediately discarded because it does not descend from the new
+       justified checkpoint (Fork A is on a branch from Slot 1).
+    6. Fork B becomes the canonical head.
 
     Why This Matters
     ----------------
@@ -806,6 +809,10 @@ def test_reorg_on_newly_justified_slot(
     - Safety guarantees maintained during reorgs
     """
     fork_choice_test(
+        # Using 9 validators: 3 for Fork A and 6 for Fork B to achieve 2/3rd for Fork B
+        anchor_state=generate_pre_state(
+            validators=Validators(data=[Validator(pubkey=Bytes52.zero()) for _ in range(9)])
+        ),
         steps=[
             # Common base at slot 1
             BlockStep(
@@ -816,63 +823,99 @@ def test_reorg_on_newly_justified_slot(
                 ),
             ),
             # Fork A: slot 2
+            # Fork A is the heaviest chain (1 block from justified slot)
             BlockStep(
-                block=BlockSpec(slot=Slot(2), parent_label="base", label="fork_a_2"),
+                block=BlockSpec(slot=Slot(2), parent_label="base", label="fork_a_1"),
                 checks=StoreChecks(
                     head_slot=Slot(2),
-                    head_root_label="fork_a_2",
+                    head_root_label="fork_a_1",
                 ),
             ),
-            # Fork B: slot 2 (competing)
+            # Fork A: slot 3
+            # Fork A is the heaviest chain (2 blocks from justified slot)
             BlockStep(
-                block=BlockSpec(slot=Slot(2), parent_label="base", label="fork_b_2"),
+                block=BlockSpec(slot=Slot(3), parent_label="fork_a_1", label="fork_a_2"),
                 checks=StoreChecks(
-                    head_slot=Slot(2),
-                    head_root_label="fork_b_2",
+                    head_slot=Slot(3),
+                    head_root_label="fork_a_2",
                 ),
             ),
-            # Fork A: slot 3 (extends first, takes lead)
+            # Fork A: slot 4
+            # Fork A is the heaviest chain (3 blocks from justified slot)
             BlockStep(
-                block=BlockSpec(slot=Slot(3), parent_label="fork_a_2", label="fork_a_3"),
+                block=BlockSpec(slot=Slot(4), parent_label="fork_a_2", label="fork_a_3"),
                 checks=StoreChecks(
-                    head_slot=Slot(3),
-                    head_root_label="fork_a_3",  # Fork A leads (2 blocks vs 1)
+                    head_slot=Slot(4),
+                    head_root_label="fork_a_3",
                 ),
             ),
-            # Fork A: slot 4 (extends, still head)
+            # Fork B: slot 5 (first block of fork B)
+            # Fork A is still the heaviest chain (3 blocks from justified slot)
             BlockStep(
-                block=BlockSpec(slot=Slot(4), parent_label="fork_a_3", label="fork_a_4"),
+                block=BlockSpec(slot=Slot(5), parent_label="base", label="fork_b_1"),
                 checks=StoreChecks(
                     head_slot=Slot(4),
-                    head_root_label="fork_a_4",  # Fork A leads (3 blocks vs 1)
+                    head_root_label="fork_a_3",
                 ),
             ),
-            # Fork B: slot 5 (enough attestations justifying Slot 2)
+            # Fork B: slot 6
+            # Validator 5 justified fork_b_1 in slot 5
+            # Validator 6 justifying fork_b_2 in slot 6
+            # Add extra justifications on fork_b_1 from validator 0, 1, 7, 8
+            # This makes fork_b_1 justified by 2/3rd of validators: 0, 1, 5, 6, 7, 8
             BlockStep(
                 block=BlockSpec(
-                    slot=Slot(5),
-                    parent_label="fork_a_4",
-                    label="fork_b_5",
+                    slot=Slot(6),
+                    parent_label="fork_b_1",
+                    label="fork_b_2",
                     attestations=[
-                        # The proposer validator_id is 5 % 4 = 1, so adding attestations for 2 and 3
                         SignedAttestationSpec(
-                            validator_id=ValidatorIndex(2),
-                            slot=Slot(4),
-                            target_slot=Slot(2),
-                            target_root_label="fork_b_2",
+                            validator_id=ValidatorIndex(0),
+                            slot=Slot(5),
+                            target_slot=Slot(5),
+                            target_root_label="fork_b_1",
+                        ),
+                        SignedAttestationSpec(
+                            validator_id=ValidatorIndex(1),
+                            slot=Slot(5),
+                            target_slot=Slot(5),
+                            target_root_label="fork_b_1",
+                        ),
+                        # fork_b_1 should be able to justify without extra attestations
+                        # from validator 5 and 6 but the test is failing without these
+                        # two attestations below because block proposer's attestations
+                        # are not being counted towards justification
+                        SignedAttestationSpec(
+                            validator_id=ValidatorIndex(5),
+                            slot=Slot(5),
+                            target_slot=Slot(5),
+                            target_root_label="fork_b_1",
                         ),
                         SignedAttestationSpec(
-                            validator_id=ValidatorIndex(3),
-                            slot=Slot(4),
-                            target_slot=Slot(2),
-                            target_root_label="fork_b_2",
+                            validator_id=ValidatorIndex(6),
+                            slot=Slot(5),
+                            target_slot=Slot(5),
+                            target_root_label="fork_b_1",
+                        ),
+                        SignedAttestationSpec(
+                            validator_id=ValidatorIndex(7),
+                            slot=Slot(5),
+                            target_slot=Slot(5),
+                            target_root_label="fork_b_1",
+                        ),
+                        SignedAttestationSpec(
+                            validator_id=ValidatorIndex(8),
+                            slot=Slot(5),
+                            target_slot=Slot(5),
+                            target_root_label="fork_b_1",
                         ),
                     ],
                 ),
                 checks=StoreChecks(
-                    head_slot=Slot(5),
-                    latest_justified_slot=Slot(2),
-                    head_root_label="fork_b_5",  # Fork B leads as Fork B at Slot 2 is justified
+                    head_slot=Slot(6),
+                    head_root_label="fork_b_2",  # Fork B now leads as fork_b_1 is justified
+                    latest_justified_slot=Slot(5),
+                    latest_justified_root_label="fork_b_1",
                 ),
             ),
         ],