add biased random range api

1 files changed, 68 insertions, 8 deletions

diff --git a/std/rand/index.zig b/std/rand/index.zig
index 0d9e58fd87..43fc43330b 100644
--- a/std/rand/index.zig
+++ b/std/rand/index.zig
@@ -57,6 +57,23 @@ pub const Random = struct {
         return @bitCast(T, unsigned_result);
     }
 
+    /// Constant-time implementation off ::uintLessThan.
+    /// The results of this function may be biased.
+    pub fn uintLessThanBiased(r: *Random, comptime T: type, less_than: T) T {
+        assert(T.is_signed == false);
+        assert(0 < less_than);
+        // Small is typically u32
+        const Small = @IntType(false, @divTrunc(T.bit_count + 31, 32) * 32);
+        // Large is typically u64
+        const Large = @IntType(false, Small.bit_count * 2);
+
+        // adapted from:
+        //   http://www.pcg-random.org/posts/bounded-rands.html
+        //   "Integer Multiplication (Biased)"
+        var x: Small = r.int(Small);
+        var m: Large = Large(x) * Large(less_than);
+        return @intCast(T, m >> Small.bit_count);
+    }
     /// Returns an evenly distributed random unsigned integer `0 <= i < less_than`.
     /// This function assumes that the underlying ::fillFn produces evenly distributed values.
     /// Within this assumption, the runtime of this function is exponentially distributed.
@@ -64,8 +81,7 @@ pub const Random = struct {
     /// the runtime of this function would technically be unbounded.
     /// However, if ::fillFn is backed by any evenly distributed pseudo random number generator,
     /// this function is guaranteed to return.
-    /// If you need deterministic runtime bounds, consider instead using `r.int(T) % less_than`,
-    /// which will usually be biased toward smaller values.
+    /// If you need deterministic runtime bounds, use `::uintLessThanBiased`.
     pub fn uintLessThan(r: *Random, comptime T: type, less_than: T) T {
         assert(T.is_signed == false);
         assert(0 < less_than);
@@ -101,6 +117,16 @@ pub const Random = struct {
         return @intCast(T, m >> Small.bit_count);
     }
 
+    /// Constant-time implementation off ::uintAtMost.
+    /// The results of this function may be biased.
+    pub fn uintAtMostBiased(r: *Random, comptime T: type, at_most: T) T {
+        assert(T.is_signed == false);
+        if (at_most == maxInt(T)) {
+            // have the full range
+            return r.int(T);
+        }
+        return r.uintLessThanBiased(T, at_most + 1);
+    }
     /// Returns an evenly distributed random unsigned integer `0 <= i <= at_most`.
     /// See ::uintLessThan, which this function uses in most cases,
     /// for commentary on the runtime of this function.
@@ -113,6 +139,22 @@ pub const Random = struct {
         return r.uintLessThan(T, at_most + 1);
     }
 
+    /// Constant-time implementation off ::intRangeLessThan.
+    /// The results of this function may be biased.
+    pub fn intRangeLessThanBiased(r: *Random, comptime T: type, at_least: T, less_than: T) T {
+        assert(at_least < less_than);
+        if (T.is_signed) {
+            // Two's complement makes this math pretty easy.
+            const UnsignedT = @IntType(false, T.bit_count);
+            const lo = @bitCast(UnsignedT, at_least);
+            const hi = @bitCast(UnsignedT, less_than);
+            const result = lo +% r.uintLessThanBiased(UnsignedT, hi -% lo);
+            return @bitCast(T, result);
+        } else {
+            // The signed implementation would work fine, but we can use stricter arithmetic operators here.
+            return at_least + r.uintLessThanBiased(T, less_than - at_least);
+        }
+    }
     /// Returns an evenly distributed random integer `at_least <= i < less_than`.
     /// See ::uintLessThan, which this function uses in most cases,
     /// for commentary on the runtime of this function.
@@ -131,6 +173,22 @@ pub const Random = struct {
         }
     }
 
+    /// Constant-time implementation off ::intRangeAtMostBiased.
+    /// The results of this function may be biased.
+    pub fn intRangeAtMostBiased(r: *Random, comptime T: type, at_least: T, at_most: T) T {
+        assert(at_least <= at_most);
+        if (T.is_signed) {
+            // Two's complement makes this math pretty easy.
+            const UnsignedT = @IntType(false, T.bit_count);
+            const lo = @bitCast(UnsignedT, at_least);
+            const hi = @bitCast(UnsignedT, at_most);
+            const result = lo +% r.uintAtMostBiased(UnsignedT, hi -% lo);
+            return @bitCast(T, result);
+        } else {
+            // The signed implementation would work fine, but we can use stricter arithmetic operators here.
+            return at_least + r.uintAtMostBiased(T, at_most - at_least);
+        }
+    }
     /// Returns an evenly distributed random integer `at_least <= i <= at_most`.
     /// See ::uintLessThan, which this function uses in most cases,
     /// for commentary on the runtime of this function.
@@ -149,15 +207,11 @@ pub const Random = struct {
         }
     }
 
-    /// Return a random integer/boolean type.
     /// TODO: deprecated. use ::boolean or ::int instead.
     pub fn scalar(r: *Random, comptime T: type) T {
-        if (T == bool) return r.boolean();
-        return r.int(T);
+        return if (T == bool) r.boolean() else r.int(T);
     }
 
-    /// Return a random integer with even distribution between `start`
-    /// inclusive and `end` exclusive.  `start` must be less than `end`.
     /// TODO: deprecated. renamed to ::intRangeLessThan
     pub fn range(r: *Random, comptime T: type, start: T, end: T) T {
         return r.intRangeLessThan(T, start, end);
@@ -373,6 +427,8 @@ fn testRandomIntAtMost() void {
     assert(r.random.intRangeAtMost(i3, -4, -1) == -1);
     r.next_value = 0xff;
     assert(r.random.intRangeAtMost(i3, -2, 1) == 1);
+
+    assert(r.random.uintAtMost(u0, 0) == 0);
 }
 
 // Generator to extend 64-bit seed values into longer sequences.
@@ -884,12 +940,16 @@ test "Random range" {
 }
 
 fn testRange(r: *Random, start: i8, end: i8) void {
+    testRangeBias(r, start, end, true);
+    testRangeBias(r, start, end, false);
+}
+fn testRangeBias(r: *Random, start: i8, end: i8, biased: bool) void {
     const count = @intCast(usize, i32(end) - i32(start));
     var values_buffer = []bool{false} ** 0x100;
     const values = values_buffer[0..count];
     var i: usize = 0;
     while (i < count) {
-        const value: i32 = r.intRangeLessThan(i8, start, end);
+        const value: i32 = if (biased) r.intRangeLessThanBiased(i8, start, end) else r.intRangeLessThan(i8, start, end);
         const index = @intCast(usize, value - start);
         if (!values[index]) {
             i += 1;