fix(forkchoice): correct LMD-GHOST orphan-skip, tighten checkpoint advance, assert start_root

src/lean_spec/forks/lstar/spec.py

@@ -1287,14 +1287,14 @@ def on_block(
 
             # Propagate checkpoint advances from the post-state.
             #
-            # Keep the checkpoint with the higher slot.
-            # On slot ties, prefer the store's own checkpoint.
+            # A candidate replaces the store's checkpoint only when its slot is strictly higher.
+            # On slot ties the store's view stays authoritative.
             #
-            # The store's checkpoint is pinned to the anchor at init and only
-            # moves forward via real justification/finalization events.
-            # On ties the store's view is authoritative.
-            latest_justified = max(store.latest_justified, post_state.latest_justified)
-            latest_finalized = max(store.latest_finalized, post_state.latest_finalized)
+            # Why: the store's checkpoint is pinned at init.
+            # It advances only on real justification or finalization events.
+            # An incoming tie must not silently swap roots.
+            latest_justified = store.latest_justified.advance_to(post_state.latest_justified)
+            latest_finalized = store.latest_finalized.advance_to(post_state.latest_finalized)
 
             store = store.model_copy(
                 update={
@@ -1416,6 +1416,10 @@ def _compute_lmd_ghost_head(
         When two branches have equal weight, the one with the lexicographically
         larger hash is chosen to break ties.
         """
+        # Invariant: the anchor must be a block the store already knows.
+        # A loud failure here beats a cryptic missing-key error deep in the weight loop.
+        assert start_root in store.blocks, f"start_root {start_root.hex()} not in store.blocks"
+
         # Remember the slot of the anchor once and reuse it during the walk.
         #
         # This avoids repeated lookups inside the inner loop.
@@ -1439,17 +1443,15 @@ def _compute_lmd_ghost_head(
                 weights[current_root] += 1
                 current_root = store.blocks[current_root].parent_root
 
-        # Build the adjacency tree (parent -> children).
+        # Build the parent -> children adjacency.
         #
-        # We use a defaultdict to avoid checking if keys exist.
+        # Genesis blocks land in the bucket keyed by the zero hash.
+        # That bucket is never consulted.
+        # The walk anchors at the latest justified root and only descends.
         children_map: dict[Bytes32, list[Bytes32]] = defaultdict(list)
 
         for root, block in store.blocks.items():
-            # 1. Structural check: skip blocks without parents (e.g., purely genesis/orphans)
-            if not block.parent_root:
-                continue
-
-            # 2. Heuristic check: prune branches early if they lack sufficient weight
+            # Prune low-weight branches early when a threshold is set.
             if min_score > 0 and weights[root] < min_score:
                 continue
 
@@ -1948,10 +1950,12 @@ def produce_block_with_signatures(
         # Update checkpoints from post-state.
         #
         # Locally produced blocks bypass normal block processing.
-        # We must manually propagate any checkpoint advances.
-        # Higher slots indicate more recent justified/finalized states.
-        latest_justified = max(final_post_state.latest_justified, store.latest_justified)
-        latest_finalized = max(final_post_state.latest_finalized, store.latest_finalized)
+        # Checkpoint advances must be propagated manually here.
+        #
+        # Tie semantics mirror the block-import path.
+        # A candidate needs a strictly higher slot to replace the store's view.
+        latest_justified = store.latest_justified.advance_to(final_post_state.latest_justified)
+        latest_finalized = store.latest_finalized.advance_to(final_post_state.latest_finalized)
 
         # Persist block and state immutably.
         new_store = store.model_copy(

src/lean_spec/types/checkpoint.py

@@ -1,9 +1,11 @@
 """
-Checkpoint Container.
+Checkpoint container.
 
-A checkpoint marks a specific moment in the chain. It combines a block
-identifier with a slot number. Checkpoints are used for justification and
-finalization.
+A checkpoint marks a specific moment in the chain.
+
+It combines a block identifier with a slot number.
+
+Checkpoints are used for justification and finalization.
 """
 
 from lean_spec.types.byte_arrays import Bytes32
@@ -27,3 +29,13 @@ def __lt__(self, other: "Checkpoint") -> bool:
             return NotImplemented
         # Slot drives the order; equal slots leave the pair incomparable.
         return self.slot < other.slot
+
+    def advance_to(self, candidate: "Checkpoint") -> "Checkpoint":
+        """
+        Return the later of two checkpoints, keeping self on a slot tie.
+
+        Forward-only progression for justified and finalized checkpoints.
+
+        The candidate replaces the receiver only when its slot is strictly higher.
+        """
+        return candidate if candidate.slot > self.slot else self

tests/lean_spec/forks/lstar/forkchoice/test_compute_block_weights.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
+import pytest
+
 from lean_spec.forks.lstar import Store
 from lean_spec.forks.lstar.containers.attestation import AttestationData
 from lean_spec.forks.lstar.spec import LstarSpec
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof
-from lean_spec.types import Checkpoint, Slot, ValidatorIndex, ValidatorIndices
+from lean_spec.types import Bytes32, Checkpoint, Slot, ValidatorIndex, ValidatorIndices
 from lean_spec.types.byte_arrays import ByteListMiB
 from tests.lean_spec.helpers import make_bytes32, make_signed_block
 
@@ -125,3 +127,15 @@ def test_multiple_attestations_accumulate(spec: LstarSpec, base_store: Store) ->
 
     # Both validators contribute weight to block1
     assert weights == {block1_root: 2}
+
+
+def test_compute_lmd_ghost_head_rejects_unknown_start_root(
+    spec: LstarSpec, base_store: Store
+) -> None:
+    """An anchor missing from the store fails the head-walk invariant loudly."""
+    # A 32-byte root that is guaranteed not to be in the store.
+    unknown_root = Bytes32(b"\xee" * 32)
+    assert unknown_root not in base_store.blocks
+
+    with pytest.raises(AssertionError, match="not in store.blocks"):
+        spec._compute_lmd_ghost_head(base_store, start_root=unknown_root, attestations={})

tests/lean_spec/forks/lstar/forkchoice/test_validator.py

@@ -407,7 +407,10 @@ def test_produce_block_missing_parent_state(self, spec: LstarSpec) -> None:
             validator_id=TEST_VALIDATOR_ID,
         )
 
-        with pytest.raises(KeyError):  # Missing head in get_proposal_head
+        # The forkchoice head walk asserts that the justified root is known.
+        # Calling produce_block on a store whose latest_justified.root is
+        # missing from blocks must fail loudly with that invariant message.
+        with pytest.raises(AssertionError, match="not in store.blocks"):
             spec.produce_block_with_signatures(store, Slot(1), ValidatorIndex(1))
 
     def test_validator_operations_invalid_parameters(

tests/lean_spec/subspecs/containers/test_checkpoint.py

@@ -66,3 +66,32 @@ def test_max_keeps_first_argument_on_slot_tie() -> None:
     b = Checkpoint(root=ROOT_B, slot=Slot(5))
     assert max(a, b) == a
     assert max(b, a) == b
+
+
+def test_advance_to_returns_candidate_on_higher_slot() -> None:
+    """A candidate at a strictly higher slot replaces the receiver."""
+    current = Checkpoint(root=ROOT_A, slot=Slot(3))
+    candidate = Checkpoint(root=ROOT_B, slot=Slot(4))
+    assert current.advance_to(candidate) == candidate
+
+
+def test_advance_to_keeps_self_on_lower_slot() -> None:
+    """A candidate at a lower slot is ignored."""
+    current = Checkpoint(root=ROOT_A, slot=Slot(4))
+    candidate = Checkpoint(root=ROOT_B, slot=Slot(3))
+    assert current.advance_to(candidate) == current
+
+
+def test_advance_to_keeps_self_on_slot_tie() -> None:
+    """On a slot tie the receiver wins regardless of root."""
+    current = Checkpoint(root=ROOT_A, slot=Slot(7))
+    candidate = Checkpoint(root=ROOT_B, slot=Slot(7))
+    assert current.advance_to(candidate) == current
+    # Symmetric: the receiver of the call always wins on a tie.
+    assert candidate.advance_to(current) == candidate
+
+
+def test_advance_to_is_idempotent() -> None:
+    """Calling against the same checkpoint returns the receiver unchanged."""
+    cp = Checkpoint(root=ROOT_A, slot=Slot(2))
+    assert cp.advance_to(cp) == cp