Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 27 additions & 11 deletions src/lean_spec/spec/forks/lstar/fork_choice.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,9 @@ def create_store( # type: ignore[override] # ty: ignore[invalid-method-overrid
def prune_stale_attestation_data(self, store: LstarStore) -> LstarStore:
"""Remove attestation data that can no longer influence fork choice.

An attestation becomes stale when its target checkpoint falls at or before
An attestation becomes stale when its head checkpoint falls at or before
the finalized slot. Such attestations cannot affect chain selection since
the target is already finalized.
the head is already finalized.

Pruning removes all attestation-related data:

Expand All @@ -114,22 +114,22 @@ def prune_stale_attestation_data(self, store: LstarStore) -> LstarStore:
"""
# Filter out stale entries from all attestation-related mappings.
#
# Each mapping is keyed by attestation data, so we check membership by slot
# against the finalized slot.
# Each mapping is keyed by attestation data, so we check the attested
# head slot against the finalized slot.
store.attestation_signatures = {
attestation_data: signatures
for attestation_data, signatures in store.attestation_signatures.items()
if attestation_data.target.slot > store.latest_finalized.slot
if attestation_data.head.slot > store.latest_finalized.slot
}
store.latest_new_aggregated_payloads = {
attestation_data: proofs
for attestation_data, proofs in store.latest_new_aggregated_payloads.items()
if attestation_data.target.slot > store.latest_finalized.slot
if attestation_data.head.slot > store.latest_finalized.slot
}
store.latest_known_aggregated_payloads = {
attestation_data: proofs
for attestation_data, proofs in store.latest_known_aggregated_payloads.items()
if attestation_data.target.slot > store.latest_finalized.slot
if attestation_data.head.slot > store.latest_finalized.slot
}
return store

Expand Down Expand Up @@ -467,7 +467,12 @@ def compute_block_weights(self, store: LstarStore) -> dict[Bytes32, int]:
for every ancestor above the finalized slot.
"""
attestations = self.extract_attestations_from_aggregated_payloads(
store, store.latest_known_aggregated_payloads
store,
{
attestation_data: proofs
for attestation_data, proofs in store.latest_known_aggregated_payloads.items()
if attestation_data.head.slot > store.latest_finalized.slot
},
)

weights = self._accumulate_ancestor_weights(
Expand Down Expand Up @@ -547,9 +552,16 @@ def update_head(self, store: LstarStore) -> LstarStore:
1. Latest justified checkpoint as the starting root
2. LMD-GHOST fork choice rule (heaviest subtree by attestation weight)
"""
# Extract attestations from known aggregated payloads
# Extract fork-choice relevant attestations from known aggregated payloads.
# A vote whose head is at or below finalization has no weight above the
# finalized boundary.
attestations = self.extract_attestations_from_aggregated_payloads(
store, store.latest_known_aggregated_payloads
store,
{
attestation_data: proofs
for attestation_data, proofs in store.latest_known_aggregated_payloads.items()
if attestation_data.head.slot > store.latest_finalized.slot
},
)

# Run LMD-GHOST fork choice algorithm.
Expand Down Expand Up @@ -636,7 +648,11 @@ def update_safe_target(self, store: LstarStore) -> LstarStore:
# "Known" is excluded by design.
attestations = self.extract_attestations_from_aggregated_payloads(
store,
store.latest_new_aggregated_payloads,
{
attestation_data: proofs
for attestation_data, proofs in store.latest_new_aggregated_payloads.items()
if attestation_data.head.slot > store.latest_finalized.slot
},
)

# Run LMD GHOST with the supermajority threshold.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,89 @@ def _make_empty_proof(participants: list[ValidatorIndex]) -> SingleMessageAggreg
)


def _add_two_block_chain(base_store: Store) -> tuple[Bytes32, Bytes32, Bytes32]:
genesis_root = base_store.head

block1 = make_signed_block(
slot=Slot(1),
proposer_index=ValidatorIndex(0),
parent_root=genesis_root,
state_root=make_bytes32(10),
)
block1_root = hash_tree_root(block1.block)

block2 = make_signed_block(
slot=Slot(2),
proposer_index=ValidatorIndex(1),
parent_root=block1_root,
state_root=make_bytes32(20),
)
block2_root = hash_tree_root(block2.block)

base_store.blocks = {
**base_store.blocks,
block1_root: block1.block,
block2_root: block2.block,
}
genesis_state = base_store.states[genesis_root]
base_store.states = {
**base_store.states,
block1_root: genesis_state,
block2_root: genesis_state,
}
base_store.head = block2_root

return genesis_root, block1_root, block2_root


def _add_finalized_fork(base_store: Store) -> tuple[Bytes32, Bytes32, Bytes32, Bytes32]:
genesis_root = base_store.head

block1 = make_signed_block(
slot=Slot(1),
proposer_index=ValidatorIndex(0),
parent_root=genesis_root,
state_root=make_bytes32(10),
)
block1_root = hash_tree_root(block1.block)

child_a = make_signed_block(
slot=Slot(2),
proposer_index=ValidatorIndex(1),
parent_root=block1_root,
state_root=make_bytes32(20),
)
child_a_root = hash_tree_root(child_a.block)

child_b = make_signed_block(
slot=Slot(2),
proposer_index=ValidatorIndex(2),
parent_root=block1_root,
state_root=make_bytes32(30),
)
child_b_root = hash_tree_root(child_b.block)

voted_root, tie_break_root = sorted([child_a_root, child_b_root])

genesis_state = base_store.states[genesis_root]
base_store.blocks = {
**base_store.blocks,
block1_root: block1.block,
child_a_root: child_a.block,
child_b_root: child_b.block,
}
base_store.states = {
**base_store.states,
block1_root: genesis_state,
child_a_root: genesis_state,
child_b_root: genesis_state,
}
base_store.head = tie_break_root
assert tie_break_root > voted_root

return genesis_root, block1_root, voted_root, tie_break_root


def test_genesis_only_store_returns_empty_weights(spec: LstarSpec, base_store: Store) -> None:
"""A genesis-only store with no attestations has no block weights."""
assert spec.compute_block_weights(base_store) == {}
Expand Down Expand Up @@ -81,6 +164,131 @@ def test_linear_chain_weight_accumulates_upward(spec: LstarSpec, base_store: Sto
assert weights == {block2_root: 1, block1_root: 1}


def test_stale_latest_message_does_not_mask_fresh_weight(
spec: LstarSpec, base_store: Store
) -> None:
"""A stale post-finalization message must not hide a validator's fresh vote."""
genesis_root = base_store.head

block1 = make_signed_block(
slot=Slot(1),
proposer_index=ValidatorIndex(0),
parent_root=genesis_root,
state_root=make_bytes32(10),
)
block1_root = hash_tree_root(block1.block)

block2 = make_signed_block(
slot=Slot(2),
proposer_index=ValidatorIndex(1),
parent_root=block1_root,
state_root=make_bytes32(20),
)
block2_root = hash_tree_root(block2.block)

base_store.blocks = {
**base_store.blocks,
block1_root: block1.block,
block2_root: block2.block,
}
genesis_state = base_store.states[genesis_root]
base_store.states = {
**base_store.states,
block1_root: genesis_state,
block2_root: genesis_state,
}
base_store.head = block2_root
base_store.latest_finalized = Checkpoint(root=block1_root, slot=Slot(1))

fresh_vote = AttestationData(
slot=Slot(2),
head=Checkpoint(root=block2_root, slot=Slot(2)),
target=Checkpoint(root=block2_root, slot=Slot(2)),
source=Checkpoint(root=block1_root, slot=Slot(1)),
)
stale_vote = AttestationData(
slot=Slot(3),
head=Checkpoint(root=block1_root, slot=Slot(1)),
target=Checkpoint(root=block1_root, slot=Slot(1)),
source=Checkpoint(root=genesis_root, slot=Slot(0)),
)

base_store.latest_known_aggregated_payloads = {
fresh_vote: {_make_empty_proof([ValidatorIndex(0)])},
stale_vote: {_make_empty_proof([ValidatorIndex(0)])},
}

weights = spec.compute_block_weights(base_store)

assert weights == {block2_root: 1}
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test is a great anchor for the bug you're fixing — it's exactly the "stale message masks a fresh one" case, so I'd keep it as-is. 👍

To round things out, here's a companion test vector I'd suggest adding right alongside it. It encodes the mirror-image situation from my inline note: a vote whose head is above finalized but whose target sits at the finalized slot. It should still earn weight.

Heads up: on the branch as it stands today this test fails (compute_block_weights returns {}), which is precisely the behavior I was flagging. Once the filter keys on head.slot, it passes — so it doubles as the regression guard that keeps this from creeping back later. 🙂

def test_fresh_head_with_finalized_target_still_counts(
    spec: LstarSpec, base_store: Store
) -> None:
    """A vote with a head above finalized still adds weight when its target sits at the finalized slot."""
    genesis_root = base_store.head

    block1 = make_signed_block(
        slot=Slot(1),
        proposer_index=ValidatorIndex(0),
        parent_root=genesis_root,
        state_root=make_bytes32(10),
    )
    block1_root = hash_tree_root(block1.block)

    block2 = make_signed_block(
        slot=Slot(2),
        proposer_index=ValidatorIndex(1),
        parent_root=block1_root,
        state_root=make_bytes32(20),
    )
    block2_root = hash_tree_root(block2.block)

    base_store.blocks = {
        **base_store.blocks,
        block1_root: block1.block,
        block2_root: block2.block,
    }
    genesis_state = base_store.states[genesis_root]
    base_store.states = {
        **base_store.states,
        block1_root: genesis_state,
        block2_root: genesis_state,
    }
    base_store.head = block2_root
    base_store.latest_finalized = Checkpoint(root=block1_root, slot=Slot(1))

    # Head is above the finalized slot, so the upward walk reaches block2 and adds weight.
    # Target sits exactly at the finalized slot, which is still a valid, justifiable target.
    fresh_vote_with_old_target = AttestationData(
        slot=Slot(2),
        head=Checkpoint(root=block2_root, slot=Slot(2)),
        target=Checkpoint(root=block1_root, slot=Slot(1)),
        source=Checkpoint(root=genesis_root, slot=Slot(0)),
    )
    base_store.latest_known_aggregated_payloads = {
        fresh_vote_with_old_target: {_make_empty_proof([ValidatorIndex(0)])},
    }

    weights = spec.compute_block_weights(base_store)

    assert weights == {block2_root: 1}

And if you wanted to take the maintainer's earlier "maybe a test vector?" nudge all the way to cross-client land, both of these (the masking case and this mirror case) would make lovely consensus fillers under tests/consensus/lstar/fc/, modeled on test_store_pruning.py — build the finalized chain with BlockStep, drop the aggregates in with GossipAggregatedAttestationStep, then assert the chosen head with StoreChecks. Totally optional, just flagging the path. 🙏



def test_fresh_head_with_finalized_target_still_counts(spec: LstarSpec, base_store: Store) -> None:
"""A vote with head above finalization still counts when its target is finalized."""
genesis_root, block1_root, block2_root = _add_two_block_chain(base_store)
base_store.latest_finalized = Checkpoint(root=block1_root, slot=Slot(1))

fresh_vote_with_finalized_target = AttestationData(
slot=Slot(2),
head=Checkpoint(root=block2_root, slot=Slot(2)),
target=Checkpoint(root=block1_root, slot=Slot(1)),
source=Checkpoint(root=genesis_root, slot=Slot(0)),
)
base_store.latest_known_aggregated_payloads = {
fresh_vote_with_finalized_target: {_make_empty_proof([ValidatorIndex(0)])},
}

weights = spec.compute_block_weights(base_store)

assert weights == {block2_root: 1}


def test_finalized_target_vote_can_drive_head_selection(spec: LstarSpec, base_store: Store) -> None:
"""A weighted head above finalization wins over the zero-weight tie break branch."""
genesis_root, block1_root, voted_root, tie_break_root = _add_finalized_fork(base_store)
finalized = Checkpoint(root=block1_root, slot=Slot(1))
base_store.latest_justified = finalized
base_store.latest_finalized = finalized

vote = AttestationData(
slot=Slot(2),
head=Checkpoint(root=voted_root, slot=Slot(2)),
target=finalized,
source=Checkpoint(root=genesis_root, slot=Slot(0)),
)
base_store.latest_known_aggregated_payloads = {
vote: {_make_empty_proof([ValidatorIndex(0)])},
}

store = spec.update_head(base_store)

assert tie_break_root != voted_root
assert store.head == voted_root


def test_finalized_target_vote_can_advance_safe_target(spec: LstarSpec, base_store: Store) -> None:
"""Safe target counts new-pool votes whose head is above finalization."""
genesis_root, block1_root, block2_root = _add_two_block_chain(base_store)
finalized = Checkpoint(root=block1_root, slot=Slot(1))
base_store.latest_justified = finalized
base_store.latest_finalized = finalized
base_store.safe_target = block1_root

vote = AttestationData(
slot=Slot(2),
head=Checkpoint(root=block2_root, slot=Slot(2)),
target=finalized,
source=Checkpoint(root=genesis_root, slot=Slot(0)),
)
base_store.latest_new_aggregated_payloads = {
vote: {_make_empty_proof([ValidatorIndex(0), ValidatorIndex(1)])},
}

store = spec.update_safe_target(base_store)

assert store.safe_target == block2_root


def test_multiple_attestations_accumulate(spec: LstarSpec, base_store: Store) -> None:
"""Multiple validators attesting to the same head accumulate weight."""
genesis_root = base_store.head
Expand Down
Loading
Loading