undercooked/tests/sim/test_rng.py

"""`sim/rng.py` is golden-vector locked — the replay contract starts here (A21).

The numbers below are not "whatever the code returns today": they are splitmix64
as published (Steele/Lea/Venc 2014, Algorithm 5), written down once and then
never moved. If a future change to the generator breaks them, a committed replay
golden would silently describe a different game, and that is exactly what a
replay is supposed to catch.

Beyond the vectors, the tests pin the properties the kernel relies on: the
counter wraps at 2**64, `next_float` is a half-open [0,1) draw that is actually
spread out, `range_int` hits both of its ends, `fork` is reproducible and never
disturbs its parent, and nothing in the module is shared global state.
"""

from __future__ import annotations

import pytest

from undercooked.sim.rng import GOLDEN_GAMMA, MASK64, Rng, mix64

pytestmark = pytest.mark.unit

# --- published splitmix64 vectors ---------------------------------------------------


def test_seed_zero_u64_vector() -> None:
    rng = Rng(0)
    assert [rng.next_u64() for _ in range(4)] == [
        16294208416658607535,
        7960286522194355700,
        487617019471545679,
        17909611376780542444,
    ]


def test_round_seed_vector() -> None:
    # 0x5EED0401 is the `random_seed` fixture value: what a real round draws first.
    rng = Rng(0x5EED_0401)
    assert [rng.next_u64() for _ in range(2)] == [
        966088850652082199,
        13654240339074911480,
    ]


def test_next_float_vector() -> None:
    rng = Rng(0)
    assert rng.next_float() == pytest.approx(0.8833108082136426, abs=0.0)
    assert rng.next_float() == pytest.approx(0.43152799704850997, abs=0.0)
    assert rng.next_float() == pytest.approx(0.026433771592597743, abs=0.0)


def test_range_int_vector() -> None:
    rng = Rng(0)
    assert [rng.range_int(0, 5) for _ in range(8)] == [1, 0, 1, 4, 1, 0, 5, 2]


def test_choice_vector() -> None:
    rng = Rng(1)
    assert [rng.choice(["a", "b", "c", "d"]) for _ in range(6)] == [
        "b",
        "d",
        "c",
        "d",
        "b",
        "a",
    ]


def test_fork_vector() -> None:
    assert Rng(0).fork(0).state == 5197578548964807871
    assert Rng(0).fork(7).state == 8409816773569330625
    assert Rng(0).fork(8).state == 15655543297350242377
    assert Rng(0).fork(0).next_u64() == 6235967106033911276


def test_mix64_is_the_published_finalizer() -> None:
    assert mix64(0) == 0  # the fixed point `fork` steps around
    assert mix64(GOLDEN_GAMMA) == 16294208416658607535  # == Rng(0).next_u64()


# --- the counter --------------------------------------------------------------------


@pytest.mark.parametrize("seed", [0, -1, -2**64, 2**64, 2**65 + 7, MASK64])
def test_any_seed_is_reduced_to_64_bits(seed: int) -> None:
    rng = Rng(seed)
    assert 0 <= rng.state <= MASK64


def test_seeds_that_agree_mod_2p64_agree() -> None:
    assert Rng(-1).state == MASK64
    assert Rng(-1).next_u64() == Rng(MASK64).next_u64()


def test_counter_wraps_and_never_leaves_64_bits() -> None:
    rng = Rng(MASK64 - 2)
    for _ in range(6):
        rng.next_u64()
        assert 0 <= rng.state <= MASK64


def test_zero_seed_is_not_a_fixed_point() -> None:
    rng = Rng(0)
    first = rng.next_u64()
    assert rng.state != 0
    assert first != rng.next_u64()


# --- draw distributions ---------------------------------------------------------------


def test_next_float_is_half_open_over_many_draws() -> None:
    rng = Rng(0xC0FFEE)
    draws = [rng.next_float() for _ in range(10_000)]
    assert all(0.0 <= draw < 1.0 for draw in draws)


def test_next_float_is_spread_out() -> None:
    rng = Rng(0xC0FFEE)
    buckets = [0] * 10
    for _ in range(10_000):
        buckets[int(rng.next_float() * 10)] += 1
    assert all(850 <= count <= 1150 for count in buckets), buckets
    mean_bucket = sum(count * (index + 0.5) for index, count in enumerate(buckets)) / 10_000
    assert 4.5 <= mean_bucket <= 5.5


def test_next_float_uses_the_whole_range() -> None:
    rng = Rng(424_242)
    draws = [rng.next_float() for _ in range(10_000)]
    assert min(draws) < 0.01
    assert max(draws) > 0.99


def test_range_int_stays_in_bounds_and_reaches_both_ends() -> None:
    rng = Rng(2024)
    draws = [rng.range_int(3, 9) for _ in range(5_000)]
    assert set(draws) == {3, 4, 5, 6, 7, 8, 9}


def test_range_int_of_a_single_value_never_moves() -> None:
    rng = Rng(9)
    assert [rng.range_int(-2, -2) for _ in range(4)] == [-2, -2, -2, -2]


def test_range_int_rejects_an_empty_range() -> None:
    with pytest.raises(ValueError, match=r"empty integer range 5\.\.2"):
        Rng(0).range_int(5, 2)


def test_choice_only_returns_members() -> None:
    rng = Rng(7)
    flavors = ["tomato", "onion", "lettuce"]
    assert {rng.choice(flavors) for _ in range(500)} == set(flavors)


def test_choice_rejects_an_empty_sequence() -> None:
    with pytest.raises(ValueError, match="empty sequence"):
        Rng(0).choice([])


# --- streams are independent, and fork does not disturb the parent ----------------------


def test_two_generators_from_one_seed_agree() -> None:
    left, right = Rng(1234), Rng(1234)
    assert [left.next_u64() for _ in range(32)] == [right.next_u64() for _ in range(32)]


def test_a_generator_is_not_shared_state() -> None:
    first, second = Rng(5), Rng(5)
    for _ in range(4):
        first.next_u64()
    assert second.state == 5
    assert second.next_u64() == Rng(5).next_u64()


def test_fork_is_reproducible_and_independent() -> None:
    parent = Rng(0xC1)
    child_a, child_b = parent.fork(3), parent.fork(3)
    assert child_a == child_b
    assert [child_a.next_u64() for _ in range(8)] == [child_b.next_u64() for _ in range(8)]


def test_fork_leaves_the_parent_stream_untouched() -> None:
    parent = Rng(0xC1)
    before = parent.state
    for salt in range(6):
        parent.fork(salt)
    assert parent.state == before
    assert parent.next_u64() == Rng(0xC1).next_u64()


@pytest.mark.parametrize(
    ("salt", "differs_from_zero"),
    [(0, False), (1, True), (2, True), (7, True), (-1, True), (MASK64, True)],
)
def test_the_salt_selects_the_stream(salt: int, differs_from_zero: bool) -> None:
    parent = Rng(0)
    other = parent.fork(salt)
    assert (other.state != parent.fork(0).state) is differs_from_zero


def test_fork_zero_is_not_the_parent_stream() -> None:
    # mix64(0) == 0, so a naive `state ^ mix64(salt)` would clone the parent here.
    assert Rng(0).fork(0).state != 0
    assert Rng(0).fork(0).next_u64() != Rng(0).next_u64()


def test_nested_forks_are_reproducible() -> None:
    outer = Rng(77)
    deep_a = outer.fork(1).fork(2).next_u64()
    deep_b = Rng(77).fork(1).fork(2).next_u64()
    assert deep_a == deep_b


def test_fork_accepts_any_integer_salt() -> None:
    assert Rng(3).fork(-42).state == Rng(3).fork(-42 + (1 << 64)).state