diff --git a/src/lean_spec/subspecs/containers/checkpoint.py b/src/lean_spec/subspecs/containers/checkpoint.py index 43d82807..279d3bad 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 17810b2f..3de1e759 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 00000000..fa8b6e22 --- /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