feat(std): add bounded random int to PCG32

guppylang/src/guppylang/std/random.py

@@ -10,6 +10,7 @@
 from typing import no_type_check
 
 from guppylang import guppy
+from guppylang.std.builtins import exit
 from guppylang.std.num import nat
 
 
@@ -21,6 +22,16 @@ def _mask32(value: nat) -> nat:
     return value & uint32_mask
 
 
+@guppy
+@no_type_check
+def _mask64(value: nat) -> nat:
+    # 2**64 - 1: PCG32 keeps its state in a 64-bit word (see pcg32_random_r).
+    # Combine two 32-bit masks; hugr cannot const-fold ``(1 << 64) - 1``.
+    low = _mask32(value)
+    high = _mask32(value >> nat(32))
+    return low | (high << nat(32))
+
+
 @guppy
 @no_type_check
 def _uint32_to_signed(value: nat) -> int:
@@ -45,12 +56,31 @@ class PCG32:
 
         rng = seeded_pcg32(1)
         value = rng.next_int()
+        roll = rng.next_int_bounded(6)
         another = rng.next_int()
     """
 
     _state: nat
     _inc: nat
 
+    @guppy
+    @no_type_check
+    def _next_uint32(self: PCG32) -> nat:
+        """Advance the generator and return the next 32-bit output word."""
+        # LCG multiplier N from the PCG paper / pcg32_random_r (64-bit state, 32-bit
+        # output).
+        pcg32_mult: nat = 6364136223846793005
+        old_state = self._state
+        self._state = _mask64(nat(old_state * pcg32_mult + self._inc))
+        # XSH-RR output permutation: xor-shift then random rotate (see pcg32_random_r).
+        xorshifted = _mask32(
+            ((_mask64(old_state) >> nat(18)) ^ _mask64(old_state)) >> nat(27)
+        )
+        # oldstate >> 59 is at most 31 for 64-bit state; mask for safe hugr shifts.
+        rot = _mask32(_mask64(old_state) >> nat(59)) & nat(31)
+        rot_inv = _mask32(nat(32) - rot) & nat(31)
+        return _mask32((xorshifted >> rot) | (xorshifted << rot_inv))
+
     @guppy
     @no_type_check
     def next_int(self: PCG32) -> int:
@@ -59,22 +89,37 @@ def next_int(self: PCG32) -> int:
         Returns a signed 32-bit integer, matching the shape of
         :py:meth:`guppylang.std.qsystem.random.RNG.random_int`.
         """
-        # LCG multiplier N from the PCG paper / pcg32_random_r (64-bit state, 32-bit
-        # output).
-        pcg32_mult: nat = 6364136223846793005
-        old_state = self._state
-        self._state = nat(old_state * pcg32_mult + self._inc)
-        # XSH-RR output permutation: xor-shift then random rotate (see pcg32_random_r).
-        xorshifted = _mask32(((old_state >> nat(18)) ^ old_state) >> nat(27))
-        rot = _mask32(old_state >> nat(59))
-        rot_inv = _mask32((~rot + nat(1)) & nat(31))
-        output = _mask32((xorshifted >> rot) | (xorshifted << rot_inv))
-        return _uint32_to_signed(output)
+        return _uint32_to_signed(self._next_uint32())
+
+    @guppy
+    @no_type_check
+    def next_int_bounded(self: PCG32, bound: nat) -> int:
+        """Generate a uniformly random integer in ``[0, bound)``.
+
+        Uses rejection sampling (``pcg32_boundedrand_r`` from the PCG reference) to
+        avoid the bias introduced by ``next_int() % bound``.
+
+        Args:
+            bound: Upper bound (exclusive); must be positive and less than ``2**31``.
+
+        Returns:
+            A value in ``[0, bound)``, matching the shape of
+            :py:meth:`guppylang.std.qsystem.random.RNG.random_int_bounded`.
+        """
+        if bound == nat(0):
+            exit("PCG32.next_int_bounded: bound must be positive")
+        # uint32_t threshold = -bound % bound in pcg32_boundedrand_r.
+        two_to_32: nat = nat(1) << nat(32)
+        threshold = _mask32(two_to_32 - _mask32(bound)) % bound
+        while True:
+            r = self._next_uint32()
+            if r >= threshold:
+                return int(r % bound)
 
 
 @guppy
 @no_type_check
-def seeded_pcg32(seed: int) -> PCG32:
+def seeded_pcg32(seed: nat) -> PCG32:
     """Create a new :py:class:`PCG32` generator from a seed value.
 
     The seed selects one of ``2**63`` possible PCG32 sequences. The same seed always
@@ -86,12 +131,12 @@ def seeded_pcg32(seed: int) -> PCG32:
     """
     # Default initstate from PCG reference examples (e.g. Rosetta Code PCG32 task).
     initstate = nat(42)
-    initseq = nat(seed)
+    initseq = seed
     # PCG requires an odd increment; initseq is the stream/sequence selector.
     inc = nat((initseq << nat(1)) | nat(1))
     # pcg32_srandom_r: advance twice after mixing initstate into state.
     rng = PCG32(nat(0), inc)
     rng.next_int()
-    rng._state += initstate
+    rng._state = _mask64(rng._state + initstate)
     rng.next_int()
     return rng

tests/integration/std/test_random.py

@@ -1,8 +1,11 @@
 from guppylang import guppy
 from guppylang.std.builtins import owned, output
 from guppylang.std.quantum import h, measure, qubit, x
+from guppylang.std.qsystem.utils import get_current_shot
 from guppylang.std.random import seeded_pcg32
 
+from selene_sim.backends.bundled_simulators import ClassicalReplay
+
 
 def test_pcg32_compile(validate) -> None:
     @guppy
@@ -84,31 +87,43 @@ def main() -> None:
         output("first", first)
         output("second", second)
 
-    results = dict(
-        main.emulator(0).coinflip_sim().with_seed(42).run().results[0].entries
-    )
-    assert results == {"first": 1307692281, "second": -444364974}
+        # Repeat with no inner RNG: outer stream must match exactly.
+        outer = seeded_pcg32(1)
+        other_first = outer.next_int()
+        other_second = outer.next_int()
+        output("other_first", other_first)
+        output("other_second", other_second)
 
+    results = main.emulator(0).coinflip_sim().with_seed(42).run().collated_shots()[0]
+    assert results["first"] == results["other_first"] == [1307692281]
+    assert results["second"] == results["other_second"] == [-444364974]
 
-def test_pcg32_independent_streams(validate, run_int_fn) -> None:
-    """Inner and outer RNG streams do not interfere with each other."""
 
-    @guppy
-    def uses_inner_rng() -> int:
-        inner = seeded_pcg32(2)
-        value = inner.next_int()
-        return value
+def test_pcg32_matches_qsystem_random() -> None:
+    """PCG32 and qsystem RNG produce the same values for the same seed."""
 
-    @guppy
-    def main() -> int:
-        outer = seeded_pcg32(1)
-        first = outer.next_int()
-        _ = uses_inner_rng()
-        second = outer.next_int()
-        return first + second
+    from guppylang.std.qsystem.random import RNG
 
-    validate(main.compile_function())
-    run_int_fn(main, 1307692281 + -444364974)
+    @guppy
+    def main() -> None:
+        pcg = seeded_pcg32(55555)
+        output("pcg_int", pcg.next_int())
+        output("pcg_bnd2", pcg.next_int_bounded(2))
+        output("pcg_bnd6", pcg.next_int_bounded(6))
+        output("pcg_bnd100", pcg.next_int_bounded(100))
+
+        qsys = RNG(55555)
+        output("qsys_int", qsys.random_int())
+        output("qsys_bnd2", qsys.random_int_bounded(2))
+        output("qsys_bnd6", qsys.random_int_bounded(6))
+        output("qsys_bnd100", qsys.random_int_bounded(100))
+        qsys.discard()
+
+    results = main.emulator(0).coinflip_sim().with_seed(42).run().collated_shots()[0]
+    assert results["pcg_int"] == results["qsys_int"] == [636174845]
+    assert results["pcg_bnd2"] == results["qsys_bnd2"] == [1]
+    assert results["pcg_bnd6"] == results["qsys_bnd6"] == [0]
+    assert results["pcg_bnd100"] == results["qsys_bnd100"] == [27]
 
 
 def test_pcg32_no_interference_with_quantum() -> None:
@@ -140,9 +155,86 @@ def main() -> None:
             value = rng_outer.next_int()
             output("rng0_1", value)
 
-    results = main.emulator(2).coinflip_sim().with_seed(42).run().results[0].entries
-    entries = dict(results)
-    assert entries["rng0_0"] == 1307692281
-    assert entries["rng0_1"] == -444364974
-    if "rng1_0" in entries:
-        assert entries["rng1_0"] == -8000311
+    measurements = [
+        [False, False],
+        [True, True],
+        [True, False],
+        [False, True],
+    ]
+    shots = (
+        main.emulator(2)
+        .with_simulator(ClassicalReplay(measurements=measurements))
+        .with_shots(len(measurements))
+        .run()
+        .collated_shots()
+    )
+    for shot in shots:
+        assert shot["rng0_0"] == [1307692281]
+        assert shot["rng0_1"] == [-444364974]
+    assert "rng1_0" not in shots[0]
+    assert shots[1]["rng1_0"] == [-8000311]
+    assert shots[2]["rng1_0"] == [-8000311]
+    assert "rng1_0" not in shots[3]
+
+
+def test_pcg32_bounded_compile(validate) -> None:
+    @guppy
+    def main() -> int:
+        rng = seeded_pcg32(1)
+        return rng.next_int_bounded(6)
+
+    validate(main.compile_function())
+
+
+def test_pcg32_bounded_sequence_seed_1(run_int_fn) -> None:
+    @guppy
+    def main() -> int:
+        rng = seeded_pcg32(1)
+        coin = rng.next_int_bounded(2)
+        die = rng.next_int_bounded(6)
+        hundred = rng.next_int_bounded(100)
+        return coin + die + hundred
+
+    run_int_fn(main, 1 + 4 + 4)
+
+
+def test_pcg32_bounded_sequence_seed_2(run_int_fn) -> None:
+    @guppy
+    def main() -> int:
+        rng = seeded_pcg32(2)
+        return rng.next_int_bounded(100)
+
+    run_int_fn(main, 85)
+
+
+def test_pcg32_bounded_deterministic_sequence(run_int_fn) -> None:
+    """Bounded draws for initseq=54 match pcg32_boundedrand_r reference."""
+
+    @guppy
+    def main() -> int:
+        rng = seeded_pcg32(54)
+        total = rng.next_int_bounded(2)
+        total += rng.next_int_bounded(6)
+        total += rng.next_int_bounded(100)
+        return total
+
+    run_int_fn(main, 1 + 3 + 24)
+
+
+def test_pcg32_bounded_bound_zero_exits() -> None:
+    """Invalid bound exits the current shot; subsequent shots still run."""
+
+    @guppy
+    def main() -> None:
+        if get_current_shot() == 0:
+            rng = seeded_pcg32(1)
+            _ = rng.next_int_bounded(0)
+        else:
+            rng = seeded_pcg32(1)
+            output("ok", rng.next_int_bounded(6))
+
+    res = main.emulator(0).coinflip_sim().with_shots(2).run()
+    assert res.results[0].entries == [
+        ("exit: PCG32.next_int_bounded: bound must be positive", 1)
+    ]
+    assert res.results[1].entries == [("ok", 3)]