subspecs: add forkchoice subspec with tests (leanEthereum#53)

* subspecs: add forkchoice subspec with tests

* mv constants

* small comment touchup

* cleanup import stuffs

* fix tests with uint

* small fix

* mv get_vote_target to store

* reorganize tests

* rm useless duplicate tests

* use Bytes32.zero()

Author: tcoratger

src/lean_spec/subspecs/chain/config.py

@@ -12,12 +12,18 @@
 
 # --- Time Parameters ---
 
+INTERVALS_PER_SLOT = Uint64(4)
+"""Number of intervals per slot for forkchoice processing."""
+
 SLOT_DURATION_MS: Final = Uint64(4000)
 """The fixed duration of a single slot in milliseconds."""
 
 SECONDS_PER_SLOT: Final = SLOT_DURATION_MS // Uint64(1000)
 """The fixed duration of a single slot in seconds."""
 
+SECONDS_PER_INTERVAL = SECONDS_PER_SLOT // INTERVALS_PER_SLOT
+"""Seconds per forkchoice processing interval."""
+
 JUSTIFICATION_LOOKBACK_SLOTS: Final = Uint64(3)
 """The number of slots to lookback for justification."""
 

src/lean_spec/subspecs/forkchoice/__init__.py

@@ -0,0 +1,18 @@
+"""
+Forkchoice algorithm implementation.
+
+This module implements the LMD GHOST forkchoice algorithm for Ethereum,
+providing the core functionality for determining the canonical chain head.
+"""
+
+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/constants.py

@@ -0,0 +1,6 @@
+"""Forkchoice algorithm constants."""
+
+from lean_spec.types import Bytes32
+
+ZERO_HASH = Bytes32.zero()
+"""All-zero hash used as genesis parent."""

src/lean_spec/subspecs/forkchoice/helpers.py

@@ -0,0 +1,86 @@
+"""
+Core forkchoice algorithms.
+
+Pure functions implementing the LMD GHOST forkchoice rule and related utilities.
+"""
+
+from typing import Dict, Optional
+
+from lean_spec.subspecs.containers import Block, Checkpoint, State
+from lean_spec.types import Bytes32, ValidatorIndex
+
+from .constants import ZERO_HASH
+
+
+def get_fork_choice_head(
+    blocks: Dict[Bytes32, Block],
+    root: Bytes32,
+    latest_votes: Dict[ValidatorIndex, Checkpoint],
+    min_score: int = 0,
+) -> Bytes32:
+    """
+    Use LMD GHOST to find the head block from a given root.
+
+    Args:
+        blocks: All known blocks indexed by hash.
+        root: Starting point root (usually latest justified).
+        latest_votes: Current votes by validator index.
+        min_score: Minimum vote count for block inclusion.
+
+    Returns:
+        Hash of the chosen head block.
+    """
+    # Start at genesis if root is zero hash
+    if root == ZERO_HASH:
+        root = min(blocks.keys(), key=lambda block_hash: blocks[block_hash].slot)
+
+    # If no votes, return the starting root immediately
+    if not latest_votes:
+        return root
+
+    # Count votes for each block (votes for descendants count for ancestors)
+    vote_weights: Dict[Bytes32, int] = {}
+
+    for vote in latest_votes.values():
+        if vote.root in blocks:
+            # Walk up from vote target, incrementing ancestor weights
+            block_hash = vote.root
+            while blocks[block_hash].slot > blocks[root].slot:
+                vote_weights[block_hash] = vote_weights.get(block_hash, 0) + 1
+                block_hash = blocks[block_hash].parent_root
+
+    # Build children mapping for blocks above min score
+    children_map: Dict[Bytes32, list[Bytes32]] = {}
+    for block_hash, block in blocks.items():
+        if block.parent_root and vote_weights.get(block_hash, 0) >= min_score:
+            children_map.setdefault(block.parent_root, []).append(block_hash)
+
+    # Walk down tree, choosing child with most votes (tiebreak by slot, then hash)
+    current = root
+    while True:
+        children = children_map.get(current, [])
+        if not children:
+            return current
+
+        # Choose best child: most votes, then highest slot, then highest hash
+        current = max(children, key=lambda x: (vote_weights.get(x, 0), blocks[x].slot, 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