Merge pull request #13417 from InKryption/rand-deterministic-indexing

std.Random: add functions with explicit index type

lib/std/rand.zig

@@ -61,14 +61,41 @@ pub const Random = struct {
     }
 
     /// Returns a random value from an enum, evenly distributed.
-    pub fn enumValue(r: Random, comptime EnumType: type) EnumType {
+    ///
+    /// Note that this will not yield consistent results across all targets
+    /// due to dependence on the representation of `usize` as an index.
+    /// See `enumValueWithIndex` for further commentary.
+    pub inline fn enumValue(r: Random, comptime EnumType: type) EnumType {
+        return r.enumValueWithIndex(EnumType, usize);
+    }
+
+    /// Returns a random value from an enum, evenly distributed.
+    ///
+    /// An index into an array of all named values is generated using the
+    /// specified `Index` type to determine the return value.
+    /// This allows for results to be independent of `usize` representation.
+    ///
+    /// Prefer `enumValue` if this isn't important.
+    ///
+    /// See `uintLessThan`, which this function uses in most cases,
+    /// for commentary on the runtime of this function.
+    pub fn enumValueWithIndex(r: Random, comptime EnumType: type, comptime Index: type) EnumType {
         comptime assert(@typeInfo(EnumType) == .Enum);
 
         // We won't use int -> enum casting because enum elements can have
         //  arbitrary values.  Instead we'll randomly pick one of the type's values.
-        const values = std.enums.values(EnumType);
-        const index = r.uintLessThan(usize, values.len);
-        return values[index];
+        const values = comptime std.enums.values(EnumType);
+        comptime assert(values.len > 0); // can't return anything
+        comptime assert(maxInt(Index) >= values.len - 1); // can't access all values
+        comptime if (values.len == 1) return values[0];
+
+        const index = if (comptime values.len - 1 == maxInt(Index))
+            r.int(Index)
+        else
+            r.uintLessThan(Index, values.len);
+
+        const MinInt = MinArrayIndex(Index);
+        return values[@intCast(MinInt, index)];
     }
 
     /// Returns a random int `i` such that `minInt(T) <= i <= maxInt(T)`.
@@ -323,14 +350,37 @@ pub const Random = struct {
     }
 
     /// Shuffle a slice into a random order.
-    pub fn shuffle(r: Random, comptime T: type, buf: []T) void {
+    ///
+    /// Note that this will not yield consistent results across all targets
+    /// due to dependence on the representation of `usize` as an index.
+    /// See `shuffleWithIndex` for further commentary.
+    pub inline fn shuffle(r: Random, comptime T: type, buf: []T) void {
+        r.shuffleWithIndex(T, buf, usize);
+    }
+
+    /// Shuffle a slice into a random order, using an index of a
+    /// specified type to maintain distribution across targets.
+    /// Asserts the index type can represent `buf.len`.
+    ///
+    /// Indexes into the slice are generated using the specified `Index`
+    /// type, which determines distribution properties. This allows for
+    /// results to be independent of `usize` representation.
+    ///
+    /// Prefer `shuffle` if this isn't important.
+    ///
+    /// See `intRangeLessThan`, which this function uses,
+    /// for commentary on the runtime of this function.
+    pub fn shuffleWithIndex(r: Random, comptime T: type, buf: []T, comptime Index: type) void {
+        const MinInt = MinArrayIndex(Index);
         if (buf.len < 2) {
             return;
         }
 
-        var i: usize = 0;
-        while (i < buf.len - 1) : (i += 1) {
-            const j = r.intRangeLessThan(usize, i, buf.len);
+        // `i <= j < max <= maxInt(MinInt)`
+        const max = @intCast(MinInt, buf.len);
+        var i: MinInt = 0;
+        while (i < max - 1) : (i += 1) {
+            const j = @intCast(MinInt, r.intRangeLessThan(Index, i, max));
             mem.swap(T, &buf[i], &buf[j]);
         }
     }
@@ -370,6 +420,13 @@ pub const Random = struct {
 
         unreachable;
     }
+
+    /// Returns the smallest of `Index` and `usize`.
+    fn MinArrayIndex(comptime Index: type) type {
+        const index_info = @typeInfo(Index).Int;
+        assert(index_info.signedness == .unsigned);
+        return if (index_info.bits >= @typeInfo(usize).Int.bits) usize else Index;
+    }
 };
 
 /// Convert a random integer 0 <= random_int <= maxValue(T),