core: better block production and usage in fixture (#218)

* core: better block production and usage in fixture

* small touchups about attestations

* small fix

* small doc touchups

* small touchups

* better doc for _resolve_parent_root

* some doc fix

* fix doc

* doc fixes

* add docstring back

* fix linter

* fix comment

Author: tcoratger

packages/testing/src/consensus_testing/test_fixtures/fork_choice.py

@@ -290,14 +290,10 @@ def _build_block_from_spec(
         """
         Build a full SignedBlockWithAttestation from a lightweight BlockSpec.
 
-        Builds blocks via state transition dry-run, similar to state transition tests,
-        but also creates a proper proposer attestation for fork choice.
-        This mimics what a local block builder would do.
-
-        TODO: We cannot use Store.produce_block_with_signatures() because it has
-        side effects (adds block to store at lines 556-559 of store.py). If the spec
-        is refactored to separate block production from store updates, we should use
-        that method instead. Until then, this manual approach is necessary.
+        This method combines:
+            - spec logic (via the state block building logic),
+            - test-specific logic (label resolution and signing),
+        to produce a complete signed block.
 
         Parameters
         ----------
@@ -307,87 +303,32 @@ def _build_block_from_spec(
             The fork choice store (used to get head state and latest justified).
         block_registry : dict[str, Block]
             Registry of labeled blocks for fork creation.
+        key_manager : XmssKeyManager
+            Key manager for signing attestations.
 
         Returns:
         -------
         SignedBlockWithAttestation
             A complete signed block ready for processing.
         """
-        # Determine proposer
-        if spec.proposer_index is None:
-            validator_count = store.states[store.head].validators.count
-            proposer_index = Uint64(int(spec.slot) % int(validator_count))
-        else:
-            proposer_index = spec.proposer_index
-
-        # Resolve parent block if parent_label is specified
-        if spec.parent_label is not None:
-            if spec.parent_label not in block_registry:
-                raise ValueError(
-                    f"parent_label '{spec.parent_label}' not found - "
-                    f"available labels: {list(block_registry.keys())}"
-                )
-            parent_block = block_registry[spec.parent_label]
-            parent_root = hash_tree_root(parent_block)
-
-            # Get state at the parent block
-            if parent_root not in store.states:
-                raise ValueError(
-                    f"parent_label '{spec.parent_label}' (root=0x{parent_root.hex()[:16]}...) "
-                    f"has no state in store - cannot build on this fork"
-                )
-            parent_state = store.states[parent_root]
-
-            # Advance state to the new block's slot
-            temp_state = parent_state.process_slots(spec.slot)
-        else:
-            # Default: build on current head
-            head_state = store.states[store.head]
-            temp_state = head_state.process_slots(spec.slot)
-            parent_root = hash_tree_root(temp_state.latest_block_header)
-
-        # Prepare attestations from spec if provided
-        attestations = []
-        attestation_signatures = []
-        if spec.attestations is not None:
-            for attestation in spec.attestations:
-                if isinstance(attestation, SignedAttestationSpec):
-                    # Use the parent state's latest_justified for source checkpoint
-                    parent_state = store.states[parent_root]
-                    signed_attestation = self._build_signed_attestation_from_spec(
-                        attestation, block_registry, parent_state
-                    )
-                    # Extract the Attestation message and signature
-                    attestations.append(signed_attestation.message)
-                    attestation_signatures.append(signed_attestation.signature)
-                else:
-                    # Already a SignedAttestation, extract the message
-                    attestations.append(attestation.message)
-                    attestation_signatures.append(attestation.signature)
-
-        # Build block with collected attestations
-        body = BlockBody(attestations=Attestations(data=attestations))
-
-        # Create temporary block for dry-run
-        temp_block = Block(
-            slot=spec.slot,
-            proposer_index=proposer_index,
-            parent_root=parent_root,
-            state_root=Bytes32.zero(),
-            body=body,
+        # Determine proposer index
+        proposer_index = spec.proposer_index or Uint64(
+            int(spec.slot) % store.states[store.head].validators.count
         )
 
-        # Process to get correct state root
-        post_state = temp_state.process_block(temp_block)
-        correct_state_root = hash_tree_root(post_state)
+        # Resolve parent root from label or default to head
+        parent_root = self._resolve_parent_root(spec, store, block_registry)
 
-        # Create final block
-        final_block = Block(
+        # Build attestations from spec
+        attestations = self._build_attestations_from_spec(spec, store, block_registry, parent_root)
+
+        # Use State.build_block for core block building (pure spec logic)
+        parent_state = store.states[parent_root]
+        final_block, _, _, _ = parent_state.build_block(
             slot=spec.slot,
             proposer_index=proposer_index,
             parent_root=parent_root,
-            state_root=correct_state_root,
-            body=body,
+            attestations=attestations,
         )
 
         # Create proposer attestation for this block
@@ -398,17 +339,13 @@ def _build_block_from_spec(
                 slot=spec.slot,
                 head=Checkpoint(root=block_root, slot=spec.slot),
                 target=Checkpoint(root=block_root, slot=spec.slot),
-                # Use the anchor block as source for genesis case
-                source=Checkpoint(root=parent_root, slot=temp_state.latest_block_header.slot),
+                source=Checkpoint(root=parent_root, slot=parent_state.latest_block_header.slot),
             ),
         )
 
         # Sign all attestations and the proposer attestation
-        signature_list = []
-        for attestation in final_block.body.attestations:
-            signature_list.append(key_manager.sign_attestation(attestation))
-        proposer_attestation_signature = key_manager.sign_attestation(proposer_attestation)
-        signature_list.append(proposer_attestation_signature)
+        signature_list = [key_manager.sign_attestation(att) for att in attestations]
+        signature_list.append(key_manager.sign_attestation(proposer_attestation))
 
         return SignedBlockWithAttestation(
             message=BlockWithAttestation(
@@ -418,6 +355,72 @@ def _build_block_from_spec(
             signature=BlockSignatures(data=signature_list),
         )
 
+    def _resolve_parent_root(
+        self,
+        spec: BlockSpec,
+        store: Store,
+        block_registry: dict[str, Block],
+    ) -> Bytes32:
+        """
+        Resolve parent root from BlockSpec.
+            - If parent_label is specified, look it up in the registry.
+            - Otherwise, default to the current head's parent.
+        """
+        # Fast path: no label means build on current head.
+        if not (label := spec.parent_label):
+            return store.head
+
+        # Label was provided: look up the block in the registry.
+        if not (parent_block := block_registry.get(label)):
+            raise ValueError(f"Parent label '{label}' not found. Available: {list(block_registry)}")
+
+        # Compute the SSZ root of the parent block.
+        #
+        # This root serves as both:
+        # - The key to look up the parent's post-state in the store
+        # - The value to place in the new block's `parent_root` field
+        parent_root = hash_tree_root(parent_block)
+
+        # Verify the parent's state exists in the store.
+        #
+        # Building a block requires the parent's post-state to:
+        # - Advance slots via `process_slots()`
+        # - Apply the new block via `process_block()`
+        #
+        # If the state is missing, we cannot proceed.
+        if parent_root not in store.states:
+            raise ValueError(
+                f"Parent '{label}' (root=0x{parent_root.hex()[:16]}...) "
+                "has no state in store - cannot build on this fork"
+            )
+
+        return parent_root
+
+    def _build_attestations_from_spec(
+        self,
+        spec: BlockSpec,
+        store: Store,
+        block_registry: dict[str, Block],
+        parent_root: Bytes32,
+    ) -> list[Attestation]:
+        """Build attestations list from BlockSpec."""
+        if spec.attestations is None:
+            return []
+
+        parent_state = store.states[parent_root]
+        attestations = []
+
+        for att_spec in spec.attestations:
+            if isinstance(att_spec, SignedAttestationSpec):
+                signed_att = self._build_signed_attestation_from_spec(
+                    att_spec, block_registry, parent_state
+                )
+                attestations.append(signed_att.message)
+            else:
+                attestations.append(att_spec.message)
+
+        return attestations
+
     def _build_signed_attestation_from_spec(
         self,
         spec: SignedAttestationSpec,
@@ -473,8 +476,7 @@ def _build_signed_attestation_from_spec(
             message=attestation,
             signature=(
                 spec.signature
-                if spec.signature is not None
-                else Signature(
+                or Signature(
                     path=HashTreeOpening(siblings=HashDigestList(data=[])),
                     rho=Randomness(data=[Fp(0) for _ in range(PROD_CONFIG.RAND_LEN_FE)]),
                     hashes=HashDigestList(data=[]),

src/lean_spec/subspecs/containers/state/state.py

@@ -6,6 +6,8 @@
 justified and finalized.
 """
 
+from typing import TYPE_CHECKING, AbstractSet, Iterable
+
 from lean_spec.subspecs.ssz.hash import hash_tree_root
 from lean_spec.types import (
     ZERO_HASH,
@@ -16,6 +18,10 @@
     is_proposer,
 )
 
+from ..attestation import Attestation, SignedAttestation
+
+if TYPE_CHECKING:
+    from lean_spec.subspecs.xmss.containers import Signature
 from ..block import Block, BlockBody, BlockHeader
 from ..block.types import Attestations
 from ..checkpoint import Checkpoint
@@ -515,3 +521,92 @@ def state_transition(self, block: Block, valid_signatures: bool = True) -> "Stat
             raise AssertionError("Invalid block state root")
 
         return new_state
+
+    def build_block(
+        self,
+        slot: Slot,
+        proposer_index: Uint64,
+        parent_root: Bytes32,
+        attestations: list[Attestation] | None = None,
+        available_signed_attestations: Iterable[SignedAttestation] | None = None,
+        known_block_roots: AbstractSet[Bytes32] | None = None,
+    ) -> tuple[Block, "State", list[Attestation], list["Signature"]]:
+        """
+        Build a valid block on top of this state.
+
+        Computes the post-state and creates a block with the correct state root.
+
+        If `available_signed_attestations` and `known_block_roots` are provided,
+        performs fixed-point attestation collection: iteratively adds valid
+        attestations until no more can be included. This is necessary because
+        processing attestations may update the justified checkpoint, which may
+        make additional attestations valid.
+
+        Args:
+            slot: Target slot for the block.
+            proposer_index: Validator index of the proposer.
+            parent_root: Root of the parent block.
+            attestations: Initial attestations to include.
+            available_signed_attestations: Pool of attestations to collect from.
+            known_block_roots: Set of known block roots for attestation validation.
+
+        Returns:
+            Tuple of (Block, post-State, collected attestations, signatures).
+        """
+        # Initialize empty attestation set for iterative collection
+        attestations = list(attestations or [])
+        signatures: list[Signature] = []
+
+        # Iteratively collect valid attestations using fixed-point algorithm
+        #
+        # Continue until no new attestations can be added to the block.
+        # This ensures we include the maximal valid attestation set.
+        while True:
+            # Create candidate block with current attestation set
+            candidate_block = Block(
+                slot=slot,
+                proposer_index=proposer_index,
+                parent_root=parent_root,
+                state_root=Bytes32.zero(),
+                body=BlockBody(attestations=Attestations(data=attestations)),
+            )
+
+            # Apply state transition to get the post-block state
+            post_state = self.process_slots(slot).process_block(candidate_block)
+
+            # No attestation source provided: done after computing post_state
+            if available_signed_attestations is None or known_block_roots is None:
+                break
+
+            # Find new valid attestations matching post-state justification
+            new_attestations: list[Attestation] = []
+            new_signatures: list[Signature] = []
+
+            for signed_attestation in available_signed_attestations:
+                data = signed_attestation.message.data
+
+                # Skip if target block is unknown
+                if data.head.root not in known_block_roots:
+                    continue
+
+                # Skip if attestation source does not match post-state's latest justified
+                if data.source != post_state.latest_justified:
+                    continue
+
+                # Add attestation if not already included
+                if signed_attestation.message not in attestations:
+                    new_attestations.append(signed_attestation.message)
+                    new_signatures.append(signed_attestation.signature)
+
+            # Fixed point reached: no new attestations found
+            if not new_attestations:
+                break
+
+            # Add new attestations and continue iteration
+            attestations.extend(new_attestations)
+            signatures.extend(new_signatures)
+
+        # Store the post state root in the block
+        final_block = candidate_block.model_copy(update={"state_root": hash_tree_root(post_state)})
+
+        return final_block, post_state, attestations, signatures