From 1b0e1167766589ab7ccd74f96323a2ed0ecb0d2b Mon Sep 17 00:00:00 2001 From: Thomas Coratger <60488569+tcoratger@users.noreply.github.com> Date: Sun, 26 Apr 2026 22:19:14 +0200 Subject: [PATCH] refactor: order Checkpoint by slot to drop max() lambda noise MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Checkpoint.__lt__ ordered by slot so the four max(...) calls in the forkchoice store can drop their `key=lambda c: c.slot` argument. Tie-break behavior is preserved: equal-slot checkpoints are incomparable, so max() keeps its first argument — same semantics each call site already relied on. Adds unit tests covering strict/reverse/equal ordering, anti-reflexivity, the NotImplemented sentinel for foreign types, the resulting TypeError on the < operator, and the two max() integration cases (different slots, equal slots). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../subspecs/containers/checkpoint.py | 8 +++ src/lean_spec/subspecs/forkchoice/store.py | 16 ++--- .../subspecs/containers/test_checkpoint.py | 70 +++++++++++++++++++ 3 files changed, 82 insertions(+), 12 deletions(-) create mode 100644 tests/lean_spec/subspecs/containers/test_checkpoint.py diff --git a/src/lean_spec/subspecs/containers/checkpoint.py b/src/lean_spec/subspecs/containers/checkpoint.py index 43d82807a..279d3bad3 100644 --- a/src/lean_spec/subspecs/containers/checkpoint.py +++ b/src/lean_spec/subspecs/containers/checkpoint.py @@ -19,3 +19,11 @@ class Checkpoint(Container): slot: Slot """The slot number of the checkpoint's block.""" + + def __lt__(self, other: "Checkpoint") -> bool: + """Order checkpoints by slot.""" + # Foreign types: defer to Python's reflected fallback. + if not isinstance(other, Checkpoint): + return NotImplemented + # Slot drives the order; equal slots leave the pair incomparable. + return self.slot < other.slot diff --git a/src/lean_spec/subspecs/forkchoice/store.py b/src/lean_spec/subspecs/forkchoice/store.py index 17810b2f3..3de1e7590 100644 --- a/src/lean_spec/subspecs/forkchoice/store.py +++ b/src/lean_spec/subspecs/forkchoice/store.py @@ -521,12 +521,8 @@ def on_block( # The store's checkpoint is pinned to the anchor at init and only # moves forward via real justification/finalization events. # On ties the store's view is authoritative. - latest_justified = max( - self.latest_justified, post_state.latest_justified, key=lambda c: c.slot - ) - latest_finalized = max( - self.latest_finalized, post_state.latest_finalized, key=lambda c: c.slot - ) + latest_justified = max(self.latest_justified, post_state.latest_justified) + latest_finalized = max(self.latest_finalized, post_state.latest_finalized) store = self.model_copy( update={ @@ -1333,12 +1329,8 @@ def produce_block_with_signatures( # Locally produced blocks bypass normal block processing. # We must manually propagate any checkpoint advances. # Higher slots indicate more recent justified/finalized states. - latest_justified = max( - final_post_state.latest_justified, store.latest_justified, key=lambda c: c.slot - ) - latest_finalized = max( - final_post_state.latest_finalized, store.latest_finalized, key=lambda c: c.slot - ) + latest_justified = max(final_post_state.latest_justified, store.latest_justified) + latest_finalized = max(final_post_state.latest_finalized, store.latest_finalized) # Persist block and state immutably. new_store = store.model_copy( diff --git a/tests/lean_spec/subspecs/containers/test_checkpoint.py b/tests/lean_spec/subspecs/containers/test_checkpoint.py new file mode 100644 index 000000000..fa8b6e227 --- /dev/null +++ b/tests/lean_spec/subspecs/containers/test_checkpoint.py @@ -0,0 +1,70 @@ +"""Unit tests for Checkpoint ordering.""" + +from __future__ import annotations + +import pytest + +from lean_spec.subspecs.containers.checkpoint import Checkpoint +from lean_spec.subspecs.containers.slot import Slot +from lean_spec.types import Bytes32 + +# Two distinct roots to verify ordering ignores root content. +ROOT_A = Bytes32(b"\xa0" * 32) +ROOT_B = Bytes32(b"\xb0" * 32) + + +def test_lt_returns_true_for_lower_slot() -> None: + """Lower slot is less than higher slot.""" + low = Checkpoint(root=ROOT_A, slot=Slot(1)) + high = Checkpoint(root=ROOT_A, slot=Slot(2)) + assert (low < high) is True + + +def test_lt_returns_false_for_higher_slot() -> None: + """Higher slot is not less than lower slot.""" + low = Checkpoint(root=ROOT_A, slot=Slot(1)) + high = Checkpoint(root=ROOT_A, slot=Slot(2)) + assert (high < low) is False + + +def test_lt_returns_false_for_equal_slots_with_different_roots() -> None: + """Equal slots are incomparable regardless of root.""" + a = Checkpoint(root=ROOT_A, slot=Slot(7)) + b = Checkpoint(root=ROOT_B, slot=Slot(7)) + assert (a < b) is False + assert (b < a) is False + + +def test_lt_returns_false_for_identical_checkpoint() -> None: + """Checkpoint is never less than itself.""" + cp = Checkpoint(root=ROOT_A, slot=Slot(3)) + assert (cp < cp) is False + + +def test_lt_returns_not_implemented_for_non_checkpoint() -> None: + """Direct dunder call returns NotImplemented for foreign types.""" + cp = Checkpoint(root=ROOT_A, slot=Slot(1)) + assert cp.__lt__(42) is NotImplemented # type: ignore[arg-type] + + +def test_lt_raises_typeerror_when_compared_with_non_checkpoint() -> None: + """Operator < raises TypeError after the reflected fallback fails.""" + cp = Checkpoint(root=ROOT_A, slot=Slot(1)) + with pytest.raises(TypeError): + _ = cp < 42 # type: ignore[operator] + + +def test_max_returns_higher_slot_regardless_of_argument_order() -> None: + """max selects the higher-slot checkpoint regardless of argument order.""" + low = Checkpoint(root=ROOT_A, slot=Slot(1)) + high = Checkpoint(root=ROOT_B, slot=Slot(2)) + assert max(low, high) == high + assert max(high, low) == high + + +def test_max_keeps_first_argument_on_slot_tie() -> None: + """max returns the first argument on slot ties.""" + a = Checkpoint(root=ROOT_A, slot=Slot(5)) + b = Checkpoint(root=ROOT_B, slot=Slot(5)) + assert max(a, b) == a + assert max(b, a) == b