From 8828fe3a7d5db66ad3caae0c275b4917c6247b5e Mon Sep 17 00:00:00 2001
From: InKryption <inkryption07@gmail.com>
Date: Wed, 2 Nov 2022 18:23:14 +0100
Subject: [PATCH] rand: add shuffleWithIndex

and reimplement shuffle in terms of it. This allows the caller to specify an index type of a fixed bit width, allowing results to be independent usize.
---
 lib/std/rand.zig | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/lib/std/rand.zig b/lib/std/rand.zig
index 61ea6f7248..17f241adfb 100644
--- a/lib/std/rand.zig
+++ b/lib/std/rand.zig
@@ -323,14 +323,38 @@ 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 {
+        comptime std.debug.assert(@typeInfo(Index).Int.signedness == .unsigned);
+        const MinInt = std.meta.Int(.unsigned, @min(@typeInfo(Index).Int.bits, @typeInfo(usize).Int.bits));
         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]);
         }
     }