Skip to content
174 changes: 88 additions & 86 deletions packages/testing/src/consensus_testing/test_fixtures/fork_choice.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
----------
Expand All @@ -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
Expand All @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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=[]),
Expand Down
95 changes: 95 additions & 0 deletions src/lean_spec/subspecs/containers/state/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Loading
Loading