forkchoice helpers: rm get_latest_justified (leanEthereum#162)

* forkchoice helpers: rm get_latest_justified

* clean up

Author: tcoratger

src/lean_spec/subspecs/forkchoice/__init__.py

@@ -7,12 +7,10 @@
 
 from .helpers import (
     get_fork_choice_head,
-    get_latest_justified,
 )
 from .store import Store
 
 __all__ = [
     "Store",
     "get_fork_choice_head",
-    "get_latest_justified",
 ]

src/lean_spec/subspecs/forkchoice/helpers.py

@@ -4,13 +4,11 @@
 Pure functions implementing the LMD GHOST forkchoice rule and related utilities.
 """
 
-from typing import Dict, Optional
+from typing import Dict
 
 from lean_spec.subspecs.containers import (
     Block,
-    Checkpoint,
     SignedAttestation,
-    State,
 )
 from lean_spec.types import Bytes32, ValidatorIndex
 
@@ -69,23 +67,3 @@ def get_fork_choice_head(
 
         # Choose best child: most attestations, then lexicographically highest hash
         current = max(children, key=lambda x: (attestation_weights.get(x, 0), x))
-
-
-def get_latest_justified(states: Dict[Bytes32, "State"]) -> Optional[Checkpoint]:
-    """
-    Find the justified checkpoint with the highest slot.
-
-    Args:
-        states: All known states indexed by hash.
-
-    Returns:
-        Latest justified checkpoint, or None if no states.
-    """
-    if not states:
-        return None
-
-    # Find state with maximum justified slot
-    latest_state = max(states.values(), key=lambda s: s.latest_justified.slot)
-
-    # Return latest justified checkpoint from that state
-    return latest_state.latest_justified

src/lean_spec/subspecs/forkchoice/store.py

@@ -43,7 +43,7 @@
 )
 from lean_spec.types.container import Container
 
-from .helpers import get_fork_choice_head, get_latest_justified
+from .helpers import get_fork_choice_head
 
 
 class Store(Container):
@@ -398,18 +398,27 @@ def update_head(self) -> "Store":
             New Store with updated head, latest_justified, and latest_finalized.
 
         """
-        # Compute latest justified checkpoint from all known states
+        # Find the Latest Justified Checkpoint
         #
-        # Scans every state in the store to find the checkpoint with the
-        # highest slot that has achieved justification.
-        latest_justified = get_latest_justified(self.states)
-
-        # Preserve current justified checkpoint if no newer one found
+        # We must first determine the anchor point for our fork choice algorithm.
+        # This anchor is the justified checkpoint (a block root and slot) with the
+        # highest slot number known across *all* known states.
+        #
+        # We find this by:
+        # a) Scanning all known states.
+        # b) Finding the state that contains the justified checkpoint with the
+        #    highest slot number.
+        # c) Extracting that specific checkpoint object to use as our anchor.
         #
-        # This handles the case where no states have updated justification.
-        # Maintains the previous justified checkpoint rather than losing it.
-        if latest_justified is None:
-            latest_justified = self.latest_justified
+        # If there are no states to scan (e.g., at initialization), the
+        # operation would fail. In this case, we fall back to using the
+        # store's currently recorded justified checkpoint, preserving the
+        # last known good anchor.
+        latest_justified = (
+            max(self.states.values(), key=lambda s: s.latest_justified.slot).latest_justified
+            if self.states
+            else self.latest_justified
+        )
 
         # Run LMD-GHOST fork choice algorithm
         #

tests/lean_spec/subspecs/forkchoice/test_helpers.py

@@ -1,15 +1,14 @@
 """Tests for pure forkchoice helper functions."""
 
-from typing import TYPE_CHECKING, Dict, Type
+from typing import TYPE_CHECKING, Dict
 
 import pytest
 
-from lean_spec.subspecs.containers import Block, BlockBody, Checkpoint, State
+from lean_spec.subspecs.containers import Block, BlockBody, Checkpoint
 from lean_spec.subspecs.containers.block import Attestations
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.forkchoice.helpers import (
     get_fork_choice_head,
-    get_latest_justified,
 )
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import Bytes32, Uint64, ValidatorIndex
@@ -142,56 +141,3 @@ def test_get_fork_choice_head_multiple_attestations(
         )
 
         assert head == target_hash
-
-
-class TestLatestJustifiedFunction:
-    """Test the pure get_latest_justified helper function."""
-
-    def test_get_latest_justified_empty(self) -> None:
-        """Test get_latest_justified with empty states."""
-        result = get_latest_justified({})
-        assert result is None
-
-    def test_get_latest_justified_single_state(self, mock_state_factory: Type["MockState"]) -> None:
-        """Test get_latest_justified with a single state."""
-        checkpoint = Checkpoint(root=Bytes32(b"test" + b"\x00" * 28), slot=Slot(5))
-
-        states: Dict[Bytes32, State] = {
-            Bytes32(b"state1" + b"\x00" * 26): mock_state_factory(checkpoint),
-        }
-
-        result = get_latest_justified(states)
-        assert result == checkpoint
-
-    def test_get_latest_justified_multiple_states(
-        self,
-        mock_state_factory: Type["MockState"],
-    ) -> None:
-        """Test get_latest_justified when states have different slots."""
-        checkpoint1 = Checkpoint(root=Bytes32(b"test1" + b"\x00" * 27), slot=Slot(10))
-        checkpoint2 = Checkpoint(
-            root=Bytes32(b"test2" + b"\x00" * 27), slot=Slot(20)
-        )  # Higher slot
-
-        states: Dict[Bytes32, State] = {
-            Bytes32(b"state1" + b"\x00" * 26): mock_state_factory(checkpoint1),
-            Bytes32(b"state2" + b"\x00" * 26): mock_state_factory(checkpoint2),
-        }
-
-        result = get_latest_justified(states)
-        assert result == checkpoint2  # Should return the one with higher slot
-
-    def test_get_latest_justified_tie_breaking(self, mock_state_factory: Type["MockState"]) -> None:
-        """Test get_latest_justified when slots tie."""
-        checkpoint1 = Checkpoint(root=Bytes32(b"test1" + b"\x00" * 27), slot=Slot(10))
-        checkpoint2 = Checkpoint(root=Bytes32(b"test2" + b"\x00" * 27), slot=Slot(10))
-
-        states: Dict[Bytes32, State] = {
-            Bytes32(b"state1" + b"\x00" * 26): mock_state_factory(checkpoint1),
-            Bytes32(b"state2" + b"\x00" * 26): mock_state_factory(checkpoint2),
-        }
-
-        result = get_latest_justified(states)
-        # Should return one of them consistently
-        assert result in [checkpoint1, checkpoint2]
-        assert result.slot == Slot(10)