refactor: some fork choice logic to help advance chain with blocks

- forkchoice was not advancing chain with blocks, made some minor
  changes as compared to consensus-specs that may not be correct
  but help show how to write forkchoice tests for vector generation.

Author: fselmo

packages/tests/src/lean_spec_tests/base_types.py

@@ -25,20 +25,3 @@ class CamelModel(BaseModel):
     def copy(self: Self, **kwargs: Any) -> Self:
         """Create a copy of the model with the updated fields that are validated."""
         return self.__class__(**(self.model_dump(exclude_unset=True) | kwargs))
-
-    @staticmethod
-    def json_encoder(obj: Any) -> Any:
-        """
-        Custom JSON encoder for leanSpec types.
-
-        Converts:
-        - Bytes32 and other BaseBytes types → "0x..." hex string
-        - Uint64 and other BaseUint types → int
-        """
-        # Check if it's a bytes type (has hex() method)
-        if hasattr(obj, "hex") and callable(obj.hex):
-            return f"0x{obj.hex()}"
-        # Check if it's a uint type (subclass of int but not bool)
-        if isinstance(obj, int) and not isinstance(obj, bool):
-            return int(obj)
-        raise TypeError(f"Object of type {type(obj)} is not JSON serializable")

packages/tests/src/lean_spec_tests/pytest_plugins/filler.py

@@ -8,7 +8,6 @@
 
 from lean_spec.subspecs.containers import State
 from lean_spec.types import Uint64
-from lean_spec_tests.base_types import CamelModel
 from lean_spec_tests.forks import Fork, get_fork_by_name, get_forks
 from lean_spec_tests.test_fixtures import BaseConsensusFixture
 
@@ -127,7 +126,6 @@ def write_fixtures(self) -> None:
                     all_tests,
                     f,
                     indent=4,
-                    default=CamelModel.json_encoder,
                 )
 
 

packages/tests/src/lean_spec_tests/test_fixtures/base.py

@@ -2,6 +2,7 @@
 
 import hashlib
 import json
+from functools import cached_property
 from typing import Any, ClassVar, Dict, Type
 
 from pydantic import Field
@@ -51,7 +52,7 @@ def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
         if cls.format_name:
             BaseConsensusFixture.formats[cls.format_name] = cls
 
-    @property
+    @cached_property
     def json_dict(self) -> Dict[str, Any]:
         """
         Return the JSON representation of the fixture.
@@ -65,7 +66,7 @@ def json_dict(self) -> Dict[str, Any]:
             exclude={"info"},
         )
 
-    @property
+    @cached_property
     def hash(self) -> str:
         """
         Generate a deterministic hash for this fixture.
@@ -77,7 +78,6 @@ def hash(self) -> str:
             self.json_dict,
             sort_keys=True,
             separators=(",", ":"),
-            default=CamelModel.json_encoder,
         )
         h = hashlib.sha256(json_str.encode("utf-8")).hexdigest()
         return f"0x{h}"

src/lean_spec/subspecs/forkchoice/helpers.py

@@ -34,10 +34,6 @@ def get_fork_choice_head(
     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] = {}
 
@@ -49,21 +45,26 @@ def get_fork_choice_head(
                 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
+    # Build children mapping for ALL blocks (not just those above min_score)
+    # This ensures fork choice works even when there are no votes
     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)
+        if block.parent_root:
+            # Only include blocks that have enough votes OR when min_score is 0
+            if min_score == 0 or 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)
+    # Walk down tree, choosing child with most votes (tiebreak by lexicographic hash)
+    # Matches consensus-specs: (weight, root) for tie-breaking
     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))
+        # Choose best child: most votes, then lexicographically highest hash
+        # Note: Removed slot tiebreaker to match consensus-specs
+        current = max(children, key=lambda x: (vote_weights.get(x, 0), x))
 
 
 def get_latest_justified(states: Dict[Bytes32, "State"]) -> Optional[Checkpoint]:

src/lean_spec/subspecs/forkchoice/store.py

@@ -102,13 +102,17 @@ class method acts as a factory for creating a new Store instance.
         anchor_root = hash_tree_root(anchor_block)
         anchor_slot = anchor_block.slot
 
+        # Create checkpoint from anchor block (matches consensus-specs)
+        # The anchor block becomes the initial justified and finalized checkpoint
+        anchor_checkpoint = Checkpoint(root=anchor_root, slot=anchor_slot)
+
         return cls(
             time=Uint64(anchor_slot * INTERVALS_PER_SLOT),
             config=state.config,
             head=anchor_root,
             safe_target=anchor_root,
-            latest_justified=state.latest_justified,
-            latest_finalized=state.latest_finalized,
+            latest_justified=anchor_checkpoint,
+            latest_finalized=anchor_checkpoint,
             blocks={anchor_root: copy.copy(anchor_block)},
             states={anchor_root: copy.copy(state)},
         )

tests/lean_spec/subspecs/forkchoice/test_fork_choice_algorithm.py

@@ -59,8 +59,14 @@ class TestLMDGHOSTAlgorithm:
     """Test the core LMD GHOST fork choice algorithm."""
 
     def test_fork_choice_no_votes(self, sample_blocks: Dict[Bytes32, Block]) -> None:
-        """Test fork choice algorithm with no votes returns the root."""
+        """
+        Test fork choice algorithm with no votes walks to the leaf.
+
+        With no votes, fork choice should walk down the tree and select the
+        leaf block (the furthest descendant), breaking ties by lexicographic hash.
+        """
         root_hash = list(sample_blocks.keys())[0]
+        leaf_hash = list(sample_blocks.keys())[2]  # block_b (slot 2, the leaf)
 
         head = get_fork_choice_head(
             blocks=sample_blocks,
@@ -69,7 +75,7 @@ def test_fork_choice_no_votes(self, sample_blocks: Dict[Bytes32, Block]) -> None
             min_score=0,
         )
 
-        assert head == root_hash
+        assert head == leaf_hash
 
     def test_fork_choice_single_vote(self, sample_blocks: Dict[Bytes32, Block]) -> None:
         """Test fork choice algorithm with a single vote."""
@@ -254,16 +260,17 @@ def test_fork_choice_tie_breaking(self) -> None:
             block_b_hash: block_b,
         }
 
-        # No votes - algorithm returns the starting root (genesis)
+        # No votes - algorithm breaks tie by lexicographically highest hash
         head = get_fork_choice_head(
             blocks=blocks,
             root=genesis_hash,
             latest_votes={},
             min_score=0,
         )
 
-        # Should return the genesis block when no votes exist
-        assert head == genesis_hash
+        # Should return the block with lexicographically highest hash
+        expected_head = max(block_a_hash, block_b_hash)
+        assert head == expected_head
 
     def test_fork_choice_deep_chain(self) -> None:
         """Test fork choice algorithm with a deeper chain."""

tests/lean_spec/subspecs/forkchoice/test_helpers.py

@@ -73,14 +73,15 @@ def test_get_fork_choice_head_with_votes(self, sample_blocks: Dict[Bytes32, Bloc
         assert head == target_hash
 
     def test_get_fork_choice_head_no_votes(self, sample_blocks: Dict[Bytes32, Block]) -> None:
-        """Test get_fork_choice_head with no votes returns the root."""
+        """Test get_fork_choice_head with no votes walks to the leaf."""
         root_hash = list(sample_blocks.keys())[0]
+        leaf_hash = list(sample_blocks.keys())[2]  # block_b is the leaf
 
         head = get_fork_choice_head(
             blocks=sample_blocks, root=root_hash, latest_votes={}, min_score=0
         )
 
-        assert head == root_hash
+        assert head == leaf_hash
 
     def test_get_fork_choice_head_with_min_score(self, sample_blocks: Dict[Bytes32, Block]) -> None:
         """Test get_fork_choice_head respects minimum score."""

tests/lean_spec/subspecs/forkchoice/test_store_lifecycle.py

@@ -139,11 +139,13 @@ def test_store_factory_method(self) -> None:
 
         # Verify initialization
         anchor_root = hash_tree_root(anchor_block)
+        anchor_checkpoint = Checkpoint(root=anchor_root, slot=Slot(0))
         assert store.config == state.config
         assert store.head == anchor_root
         assert store.safe_target == anchor_root
-        assert store.latest_justified == state.latest_justified
-        assert store.latest_finalized == state.latest_finalized
+        # Store uses anchor checkpoint, not state's checkpoint
+        assert store.latest_justified == anchor_checkpoint
+        assert store.latest_finalized == anchor_checkpoint
         assert anchor_root in store.blocks
         assert anchor_root in store.states
 

tests/lean_spec/subspecs/forkchoice/test_validator.py

@@ -371,7 +371,8 @@ def test_multiple_validators_coordination(self, sample_store: Store) -> None:
             assert att.source.root == first_att.source.root
 
         # Validator 2 produces next block for slot 2
-        # Without votes for block1, this will build on genesis (current head)
+        # After processing block1, head should be block1 (fork choice walks the tree)
+        # So block2 will build on block1
         block2 = sample_store.produce_block(Slot(2), ValidatorIndex(2))
 
         # Verify block properties
@@ -383,10 +384,18 @@ def test_multiple_validators_coordination(self, sample_store: Store) -> None:
         assert block1_hash in sample_store.blocks
         assert block2_hash in sample_store.blocks
 
-        # Both blocks should build on genesis (the current head)
-        genesis_hash = sample_store.head
+        # block1 builds on genesis, block2 builds on block1 (current head)
+        # Get the original genesis hash from the store's blocks
+        genesis_hash = min(
+            (
+                root
+                for root in sample_store.blocks.keys()
+                if sample_store.blocks[root].slot == Slot(0)
+            ),
+            key=lambda root: root,
+        )
         assert block1.parent_root == genesis_hash
-        assert block2.parent_root == genesis_hash
+        assert block2.parent_root == block1_hash
 
     def test_validator_edge_cases(self, sample_store: Store) -> None:
         """Test edge cases in validator operations."""

tests/spec_tests/devnet/fork_choice/test_head_selection.py

@@ -5,54 +5,81 @@
     BlockBuilder,
     BlockStep,
     ForkChoiceTestFiller,
+    StoreChecks,
 )
 
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.containers.state import State
+from lean_spec.subspecs.ssz.hash import hash_tree_root
 
 pytestmark = pytest.mark.valid_until("Devnet")
 
 
-def test_process_single_block(
+def test_head_updates_after_single_block(
     fork_choice_test: ForkChoiceTestFiller,
     genesis: State,
 ) -> None:
     """
-    Test that a single block can be processed successfully.
+    Test that head updates correctly after processing a single block.
 
-    This is the simplest fork choice test - just verify that we can
-    process a block through the Store without errors.
+    With no attestations, fork choice should select the latest block
+    on the canonical chain.
     """
     # Create a block at slot 1
     block1 = BlockBuilder(genesis).build(Slot(1))
+    block1_root = hash_tree_root(block1.message)
 
     # Generate the test fixture
-    # Note: anchor_block is auto-generated from anchor_state
-    # We don't add checks here - just verify the block processes successfully
+    # After processing block 1, the head should point to block 1
     fork_choice_test(
         anchor_state=genesis,
         steps=[
-            BlockStep(block=block1),
+            BlockStep(
+                block=block1,
+                checks=StoreChecks(
+                    head_slot=Slot(1),
+                    head_root=block1_root,
+                ),
+            ),
         ],
     )
 
 
-def test_process_two_blocks_sequential(
+def test_head_advances_with_sequential_blocks(
     fork_choice_test: ForkChoiceTestFiller,
     genesis: State,
 ) -> None:
-    """Test processing two sequential blocks."""
+    """
+    Test head selection advances through sequential blocks.
+
+    Each new block should become the new head since there are no forks.
+    """
     # Create blocks
     block1 = BlockBuilder(genesis).build(Slot(1))
+    block1_root = hash_tree_root(block1.message)
+
     state_after_block1 = genesis.state_transition(block1)
     block2 = BlockBuilder(state_after_block1).build(Slot(2))
+    block2_root = hash_tree_root(block2.message)
 
     # Generate the test fixture
-    # Just verify both blocks process successfully
+    # Head should advance from block 1 to block 2
     fork_choice_test(
         anchor_state=genesis,
         steps=[
-            BlockStep(block=block1),
-            BlockStep(block=block2),
+            BlockStep(
+                block=block1,
+                checks=StoreChecks(
+                    head_slot=Slot(1),
+                    head_root=block1_root,
+                ),
+            ),
+            BlockStep(
+                block=block2,
+                checks=StoreChecks(
+                    head_slot=Slot(2),
+                    head_root=block2_root,
+                ),
+            ),
         ],
     )