forkchoice: implement new validator methods (leanEthereum#54)

* forkchoice: implement new validator methods

* fix tests

* cleanup

Author: tcoratger

src/lean_spec/subspecs/containers/state.py

@@ -8,7 +8,7 @@
 from lean_spec.subspecs.chain import config as chainconfig
 from lean_spec.subspecs.ssz.constants import ZERO_HASH
 from lean_spec.subspecs.ssz.hash import hash_tree_root
-from lean_spec.types import Boolean, Bytes32, Container, Uint64, ValidatorIndex
+from lean_spec.types import Boolean, Bytes32, Container, Uint64, ValidatorIndex, is_proposer
 from lean_spec.types import List as SSZList
 from lean_spec.types.bitfields import Bitlist
 
@@ -108,7 +108,7 @@ def is_proposer(self, validator_index: ValidatorIndex) -> bool:
         The proposer selection follows a simple round-robin mechanism based on the
         slot number and the total number of validators.
         """
-        return self.slot % self.config.num_validators == validator_index
+        return is_proposer(validator_index, Uint64(self.slot.as_int()), self.config.num_validators)
 
     def get_justifications(self) -> Dict[Bytes32, List[Boolean]]:
         """
@@ -545,7 +545,8 @@ def process_attestations(
             # Track a unique vote for (target_root, validator_id) only
             if target_root not in justifications:
                 # Initialize a fresh bitvector for this target root (all False).
-                justifications[target_root] = [Boolean(False)] * self.config.num_validators.as_int()
+                limit = DEVNET_CONFIG.validator_registry_limit.as_int()
+                justifications[target_root] = [Boolean(False)] * limit
 
             validator_id = vote.validator_id.as_int()
             if not justifications[target_root][validator_id]:

src/lean_spec/subspecs/containers/vote.py

@@ -1,12 +1,13 @@
 """Vote Containers."""
 
 from lean_spec.subspecs.containers.slot import Slot
-from lean_spec.types import Bytes32, StrictBaseModel, Uint64
+from lean_spec.types import Bytes32, Uint64
+from lean_spec.types.container import Container
 
 from .checkpoint import Checkpoint
 
 
-class Vote(StrictBaseModel):
+class Vote(Container):
     """Represents a validator's vote for chain head."""
 
     validator_id: Uint64
@@ -25,7 +26,7 @@ class Vote(StrictBaseModel):
     """The last justified checkpoint known to the validator."""
 
 
-class SignedVote(StrictBaseModel):
+class SignedVote(Container):
     """A container for a vote and its corresponding signature."""
 
     data: Vote

src/lean_spec/subspecs/forkchoice/store.py

@@ -14,14 +14,16 @@
 )
 from lean_spec.subspecs.containers import (
     Block,
+    BlockBody,
     Checkpoint,
     Config,
     SignedVote,
     State,
+    Vote,
 )
 from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.ssz.hash import hash_tree_root
-from lean_spec.types import Bytes32, Uint64, ValidatorIndex
+from lean_spec.types import Bytes32, Uint64, ValidatorIndex, is_proposer
 from lean_spec.types.container import Container
 
 from .helpers import get_fork_choice_head, get_latest_justified
@@ -362,3 +364,152 @@ def get_vote_target(self) -> Checkpoint:
 
         target_block = self.blocks[target_block_root]
         return Checkpoint(root=hash_tree_root(target_block), slot=target_block.slot)
+
+    def produce_block(self, slot: Slot, validator_index: ValidatorIndex) -> Block:
+        """
+        Produce a new block for the given slot and validator.
+
+        Algorithm Overview:
+        1. Validate proposer authorization for the target slot
+        2. Get the current chain head as the parent block
+        3. Iteratively build attestation set:
+           - Create candidate block with current attestations
+           - Apply state transition (slot advancement + block processing)
+           - Find new valid attestations matching post-state requirements
+           - Continue until no new attestations can be added
+        4. Finalize block with computed state root and store it
+
+        Args:
+            slot: Target slot number for block production
+            validator_index: Index of validator authorized to propose this block
+
+        Returns:
+            Complete block with maximal attestation set and valid state root
+
+        Raises:
+            AssertionError: If validator lacks proposer authorization for slot
+        """
+        # Validate proposer authorization for this slot
+        if not is_proposer(validator_index, slot, self.config.num_validators):
+            msg = f"Validator {validator_index} is not the proposer for slot {slot}"
+            raise AssertionError(msg)
+
+        # Get parent block and state to build upon
+        head_root = self.get_proposal_head(slot)
+        head_state = self.states[head_root]
+
+        # Initialize empty attestation set for iterative collection
+        attestations: list[SignedVote] = []
+
+        # Iteratively collect valid attestations using fixed-point algorithm
+        #
+        # Continue until no new attestations can be added to the block
+        while True:
+            # Create candidate block with current attestation set
+            candidate_block = Block(
+                slot=slot,
+                proposer_index=validator_index,
+                parent_root=head_root,
+                state_root=Bytes32.zero(),  # Temporary; updated after state computation
+                body=BlockBody(attestations=attestations),
+            )
+
+            # Apply state transition to get the post-block state
+            # First advance state to target slot, then process the block
+            advanced_state = head_state.process_slots(slot)
+            post_state = advanced_state.process_block(candidate_block)
+
+            # Find new valid attestations matching post-state justification
+            new_attestations: list[SignedVote] = []
+            for validator_id, checkpoint in self.latest_known_votes.items():
+                # Skip if target block is unknown in our store
+                if checkpoint.root not in self.blocks:
+                    continue
+
+                # Create attestation with post-state's latest justified as source
+                vote = Vote(
+                    validator_id=validator_id,
+                    slot=checkpoint.slot,
+                    head=checkpoint,
+                    target=checkpoint,
+                    source=post_state.latest_justified,
+                )
+                signed_vote = SignedVote(data=vote, signature=Bytes32.zero())
+
+                # Include if not already in attestation set
+                if signed_vote not in attestations:
+                    new_attestations.append(signed_vote)
+
+            # Fixed point reached: no new attestations found
+            if not new_attestations:
+                break
+
+            # Add new attestations and continue iteration
+            attestations.extend(new_attestations)
+
+        # Create final block with all collected attestations
+        final_state = head_state.process_slots(slot)
+        final_block = Block(
+            slot=slot,
+            proposer_index=validator_index,
+            parent_root=head_root,
+            state_root=Bytes32.zero(),  # Will be updated with computed hash
+            body=BlockBody(attestations=attestations),
+        )
+
+        # Apply state transition to get final post-state and compute state root
+        final_post_state = final_state.process_block(final_block)
+        finalized_block = final_block.model_copy(
+            update={"state_root": hash_tree_root(final_post_state)}
+        )
+
+        # Store block and state in forkchoice store
+        block_hash = hash_tree_root(finalized_block)
+        self.blocks[block_hash] = finalized_block
+        self.states[block_hash] = final_post_state
+
+        return finalized_block
+
+    def produce_attestation_vote(self, slot: Slot, validator_index: ValidatorIndex) -> Vote:
+        """
+        Produce an attestation vote for the given slot and validator.
+
+        This method constructs a Vote object according to the lean protocol
+        specification for attestation voting. The vote represents the
+        validator's view of the chain state and their choice for the
+        next justified checkpoint.
+
+        The algorithm:
+        1. Get the current head block for proposal
+        2. Calculate the appropriate vote target using current forkchoice state
+        3. Use the store's latest justified checkpoint as the vote source
+        4. Construct and return the complete Vote object
+
+        Args:
+            slot: The slot for which to produce the attestation vote.
+            validator_index: The validator index producing the vote.
+
+        Returns:
+            A fully constructed Vote object ready for signing and broadcast.
+        """
+        # Get the head block the validator sees for this slot
+        head_root = self.get_proposal_head(slot)
+        head_checkpoint = Checkpoint(
+            root=head_root,
+            slot=self.blocks[head_root].slot,
+        )
+
+        # Calculate the target checkpoint for this vote
+        #
+        # This uses the store's current forkchoice state to determine
+        # the appropriate attestation target
+        target_checkpoint = self.get_vote_target()
+
+        # Create the vote using current forkchoice state
+        return Vote(
+            validator_id=validator_index,
+            slot=slot,
+            head=head_checkpoint,
+            target=target_checkpoint,
+            source=self.latest_justified,
+        )

src/lean_spec/types/__init__.py

@@ -7,14 +7,15 @@
 from .collections import List, Vector
 from .container import Container
 from .uint import Uint64
-from .validator import ValidatorIndex
+from .validator import ValidatorIndex, is_proposer
 
 __all__ = [
     "Uint64",
     "BasisPoint",
     "Bytes32",
     "StrictBaseModel",
     "ValidatorIndex",
+    "is_proposer",
     "List",
     "Vector",
     "Boolean",

src/lean_spec/types/validator.py

@@ -1,6 +1,24 @@
-"""Validator-related type definitions for the specification."""
+"""Validator-related type definitions and utilities for the specification."""
 
 from .uint import Uint64
 
 ValidatorIndex = Uint64
 """A type alias for a validator's index in the registry."""
+
+
+def is_proposer(validator_index: ValidatorIndex, slot: Uint64, num_validators: Uint64) -> bool:
+    """
+    Determine if a validator is the proposer for a given slot.
+
+    Uses round-robin proposer selection based on slot number and total
+    validator count, following the lean protocol specification.
+
+    Args:
+        validator_index: The validator's unique index.
+        slot: The slot number to check proposer assignment for.
+        num_validators: Total number of validators in the registry.
+
+    Returns:
+        True if the validator is the proposer for the slot, False otherwise.
+    """
+    return slot % num_validators == validator_index

tests/lean_spec/subspecs/containers/test_state.py

@@ -190,59 +190,6 @@ def base_state(
     )
 
 
-def test_is_proposer(
-    sample_config: Config,
-    sample_block_header: BlockHeader,
-    sample_checkpoint: Checkpoint,
-) -> None:
-    """
-    Validate round-robin proposer selection.
-
-    Strategy
-    --------
-    - Build states at different slots.
-    - Check that proposer_index == slot % num_validators is recognized as proposer.
-    - Check neighbors are not recognized.
-    """
-
-    def create_state_at_slot(slot: int) -> State:
-        """Local helper: instantiate a State pinned at a given slot."""
-        # Return a State with the given slot and default/empty lists elsewhere.
-        return State(
-            config=sample_config,
-            slot=Slot(slot),
-            latest_block_header=sample_block_header,
-            latest_justified=sample_checkpoint,
-            latest_finalized=sample_checkpoint,
-            historical_block_hashes=[],
-            justified_slots=[],
-            justifications_roots=[],
-            justifications_validators=[],
-        )
-
-    # At slot 0, validator 0 should be the proposer (0 % 10 == 0)
-    state_slot_0 = create_state_at_slot(0)
-    # True for the correct proposer.
-    assert state_slot_0.is_proposer(ValidatorIndex(0)) is True
-    # False for a different validator.
-    assert state_slot_0.is_proposer(ValidatorIndex(1)) is False
-
-    # At slot 7, validator 7 should be the proposer (7 % 10 == 7)
-    state_slot_7 = create_state_at_slot(7)
-    assert state_slot_7.is_proposer(ValidatorIndex(7)) is True
-    assert state_slot_7.is_proposer(ValidatorIndex(8)) is False
-
-    # At slot 10, wrap-around selects validator 0 (10 % 10 == 0)
-    state_slot_10 = create_state_at_slot(10)
-    assert state_slot_10.is_proposer(ValidatorIndex(0)) is True
-    assert state_slot_10.is_proposer(ValidatorIndex(1)) is False
-
-    # At slot 23, wrap-around selects validator 3 (23 % 10 == 3)
-    state_slot_23 = create_state_at_slot(23)
-    assert state_slot_23.is_proposer(ValidatorIndex(3)) is True
-    assert state_slot_23.is_proposer(ValidatorIndex(2)) is False
-
-
 def test_get_justifications_empty(base_state: State) -> None:
     """
     get_justifications: empty input yields empty map.