authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-10-21 05:26:19-04:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-10-21 05:26:19-04:00
log3d6e63337164b42fec4df55687071be38d33dce9
treeb8a49b76cf5cd576964b00bde6f18851da81c51f
parentb3aaf854cdd814998a09195ba05f5e2bad4c60cb
parent2b4c5d990cc7ec22da037b37e68b5ca48da958de
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #16840 from e4m2/rand-int

std.rand: Support integers with >64 bits in more functions

1 files changed, 32 insertions(+), 45 deletions(-)

lib/std/rand.zig+32-45
...@@ -113,22 +113,16 @@ pub const Random = struct {...@@ -113,22 +113,16 @@ pub const Random = struct {
113 // TODO: endian portability is pointless if the underlying prng isn't endian portable.113 // TODO: endian portability is pointless if the underlying prng isn't endian portable.
114 // TODO: document the endian portability of this library.114 // TODO: document the endian portability of this library.
115 const byte_aligned_result = mem.readIntSliceLittle(ByteAlignedT, &rand_bytes);115 const byte_aligned_result = mem.readIntSliceLittle(ByteAlignedT, &rand_bytes);
116 const unsigned_result = @as(UnsignedT, @truncate(byte_aligned_result));116 const unsigned_result: UnsignedT = @truncate(byte_aligned_result);
117 return @as(T, @bitCast(unsigned_result));117 return @bitCast(unsigned_result);
118 }118 }
119119
120 /// Constant-time implementation off `uintLessThan`.120 /// Constant-time implementation off `uintLessThan`.
121 /// The results of this function may be biased.121 /// The results of this function may be biased.
122 pub fn uintLessThanBiased(r: Random, comptime T: type, less_than: T) T {122 pub fn uintLessThanBiased(r: Random, comptime T: type, less_than: T) T {
123 comptime assert(@typeInfo(T).Int.signedness == .unsigned);123 comptime assert(@typeInfo(T).Int.signedness == .unsigned);
124 const bits = @typeInfo(T).Int.bits;
125 comptime assert(bits <= 64); // TODO: workaround: LLVM ERROR: Unsupported library call operation!
126 assert(0 < less_than);124 assert(0 < less_than);
127 if (bits <= 32) {125 return limitRangeBiased(T, r.int(T), less_than);
128 return @as(T, @intCast(limitRangeBiased(u32, r.int(u32), less_than)));
129 } else {
130 return @as(T, @intCast(limitRangeBiased(u64, r.int(u64), less_than)));
131 }
132 }126 }
133127
134 /// Returns an evenly distributed random unsigned integer `0 <= i < less_than`.128 /// Returns an evenly distributed random unsigned integer `0 <= i < less_than`.
...@@ -142,22 +136,16 @@ pub const Random = struct {...@@ -142,22 +136,16 @@ pub const Random = struct {
142 pub fn uintLessThan(r: Random, comptime T: type, less_than: T) T {136 pub fn uintLessThan(r: Random, comptime T: type, less_than: T) T {
143 comptime assert(@typeInfo(T).Int.signedness == .unsigned);137 comptime assert(@typeInfo(T).Int.signedness == .unsigned);
144 const bits = @typeInfo(T).Int.bits;138 const bits = @typeInfo(T).Int.bits;
145 comptime assert(bits <= 64); // TODO: workaround: LLVM ERROR: Unsupported library call operation!
146 assert(0 < less_than);139 assert(0 < less_than);
147 // Small is typically u32
148 const small_bits = @divTrunc(bits + 31, 32) * 32;
149 const Small = std.meta.Int(.unsigned, small_bits);
150 // Large is typically u64
151 const Large = std.meta.Int(.unsigned, small_bits * 2);
152140
153 // adapted from:141 // adapted from:
154 // http://www.pcg-random.org/posts/bounded-rands.html142 // http://www.pcg-random.org/posts/bounded-rands.html
155 // "Lemire's (with an extra tweak from me)"143 // "Lemire's (with an extra tweak from me)"
156 var x: Small = r.int(Small);144 var x = r.int(T);
157 var m: Large = @as(Large, x) * @as(Large, less_than);145 var m = math.mulWide(T, x, less_than);
158 var l: Small = @as(Small, @truncate(m));146 var l: T = @truncate(m);
159 if (l < less_than) {147 if (l < less_than) {
160 var t: Small = -%less_than;148 var t = -%less_than;
161149
162 if (t >= less_than) {150 if (t >= less_than) {
163 t -= less_than;151 t -= less_than;
...@@ -166,12 +154,12 @@ pub const Random = struct {...@@ -166,12 +154,12 @@ pub const Random = struct {
166 }154 }
167 }155 }
168 while (l < t) {156 while (l < t) {
169 x = r.int(Small);157 x = r.int(T);
170 m = @as(Large, x) * @as(Large, less_than);158 m = math.mulWide(T, x, less_than);
171 l = @as(Small, @truncate(m));159 l = @truncate(m);
172 }160 }
173 }161 }
174 return @as(T, @intCast(m >> small_bits));162 return @intCast(m >> bits);
175 }163 }
176164
177 /// Constant-time implementation off `uintAtMost`.165 /// Constant-time implementation off `uintAtMost`.
...@@ -205,10 +193,10 @@ pub const Random = struct {...@@ -205,10 +193,10 @@ pub const Random = struct {
205 if (info.signedness == .signed) {193 if (info.signedness == .signed) {
206 // Two's complement makes this math pretty easy.194 // Two's complement makes this math pretty easy.
207 const UnsignedT = std.meta.Int(.unsigned, info.bits);195 const UnsignedT = std.meta.Int(.unsigned, info.bits);
208 const lo = @as(UnsignedT, @bitCast(at_least));196 const lo: UnsignedT = @bitCast(at_least);
209 const hi = @as(UnsignedT, @bitCast(less_than));197 const hi: UnsignedT = @bitCast(less_than);
210 const result = lo +% r.uintLessThanBiased(UnsignedT, hi -% lo);198 const result = lo +% r.uintLessThanBiased(UnsignedT, hi -% lo);
211 return @as(T, @bitCast(result));199 return @bitCast(result);
212 } else {200 } else {
213 // The signed implementation would work fine, but we can use stricter arithmetic operators here.201 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
214 return at_least + r.uintLessThanBiased(T, less_than - at_least);202 return at_least + r.uintLessThanBiased(T, less_than - at_least);
...@@ -224,10 +212,10 @@ pub const Random = struct {...@@ -224,10 +212,10 @@ pub const Random = struct {
224 if (info.signedness == .signed) {212 if (info.signedness == .signed) {
225 // Two's complement makes this math pretty easy.213 // Two's complement makes this math pretty easy.
226 const UnsignedT = std.meta.Int(.unsigned, info.bits);214 const UnsignedT = std.meta.Int(.unsigned, info.bits);
227 const lo = @as(UnsignedT, @bitCast(at_least));215 const lo: UnsignedT = @bitCast(at_least);
228 const hi = @as(UnsignedT, @bitCast(less_than));216 const hi: UnsignedT = @bitCast(less_than);
229 const result = lo +% r.uintLessThan(UnsignedT, hi -% lo);217 const result = lo +% r.uintLessThan(UnsignedT, hi -% lo);
230 return @as(T, @bitCast(result));218 return @bitCast(result);
231 } else {219 } else {
232 // The signed implementation would work fine, but we can use stricter arithmetic operators here.220 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
233 return at_least + r.uintLessThan(T, less_than - at_least);221 return at_least + r.uintLessThan(T, less_than - at_least);
...@@ -242,10 +230,10 @@ pub const Random = struct {...@@ -242,10 +230,10 @@ pub const Random = struct {
242 if (info.signedness == .signed) {230 if (info.signedness == .signed) {
243 // Two's complement makes this math pretty easy.231 // Two's complement makes this math pretty easy.
244 const UnsignedT = std.meta.Int(.unsigned, info.bits);232 const UnsignedT = std.meta.Int(.unsigned, info.bits);
245 const lo = @as(UnsignedT, @bitCast(at_least));233 const lo: UnsignedT = @bitCast(at_least);
246 const hi = @as(UnsignedT, @bitCast(at_most));234 const hi: UnsignedT = @bitCast(at_most);
247 const result = lo +% r.uintAtMostBiased(UnsignedT, hi -% lo);235 const result = lo +% r.uintAtMostBiased(UnsignedT, hi -% lo);
248 return @as(T, @bitCast(result));236 return @bitCast(result);
249 } else {237 } else {
250 // The signed implementation would work fine, but we can use stricter arithmetic operators here.238 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
251 return at_least + r.uintAtMostBiased(T, at_most - at_least);239 return at_least + r.uintAtMostBiased(T, at_most - at_least);
...@@ -261,10 +249,10 @@ pub const Random = struct {...@@ -261,10 +249,10 @@ pub const Random = struct {
261 if (info.signedness == .signed) {249 if (info.signedness == .signed) {
262 // Two's complement makes this math pretty easy.250 // Two's complement makes this math pretty easy.
263 const UnsignedT = std.meta.Int(.unsigned, info.bits);251 const UnsignedT = std.meta.Int(.unsigned, info.bits);
264 const lo = @as(UnsignedT, @bitCast(at_least));252 const lo: UnsignedT = @bitCast(at_least);
265 const hi = @as(UnsignedT, @bitCast(at_most));253 const hi: UnsignedT = @bitCast(at_most);
266 const result = lo +% r.uintAtMost(UnsignedT, hi -% lo);254 const result = lo +% r.uintAtMost(UnsignedT, hi -% lo);
267 return @as(T, @bitCast(result));255 return @bitCast(result);
268 } else {256 } else {
269 // The signed implementation would work fine, but we can use stricter arithmetic operators here.257 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
270 return at_least + r.uintAtMost(T, at_most - at_least);258 return at_least + r.uintAtMost(T, at_most - at_least);
...@@ -293,9 +281,9 @@ pub const Random = struct {...@@ -293,9 +281,9 @@ pub const Random = struct {
293 rand_lz += @clz(r.int(u32) | 0x7FF);281 rand_lz += @clz(r.int(u32) | 0x7FF);
294 }282 }
295 }283 }
296 const mantissa = @as(u23, @truncate(rand));284 const mantissa: u23 = @truncate(rand);
297 const exponent = @as(u32, 126 - rand_lz) << 23;285 const exponent = @as(u32, 126 - rand_lz) << 23;
298 return @as(f32, @bitCast(exponent | mantissa));286 return @bitCast(exponent | mantissa);
299 },287 },
300 f64 => {288 f64 => {
301 // Use 52 random bits for the mantissa, and the rest for the exponent.289 // Use 52 random bits for the mantissa, and the rest for the exponent.
...@@ -320,7 +308,7 @@ pub const Random = struct {...@@ -320,7 +308,7 @@ pub const Random = struct {
320 }308 }
321 const mantissa = rand & 0xFFFFFFFFFFFFF;309 const mantissa = rand & 0xFFFFFFFFFFFFF;
322 const exponent = (1022 - rand_lz) << 52;310 const exponent = (1022 - rand_lz) << 52;
323 return @as(f64, @bitCast(exponent | mantissa));311 return @bitCast(exponent | mantissa);
324 },312 },
325 else => @compileError("unknown floating point type"),313 else => @compileError("unknown floating point type"),
326 }314 }
...@@ -332,7 +320,7 @@ pub const Random = struct {...@@ -332,7 +320,7 @@ pub const Random = struct {
332 pub fn floatNorm(r: Random, comptime T: type) T {320 pub fn floatNorm(r: Random, comptime T: type) T {
333 const value = ziggurat.next_f64(r, ziggurat.NormDist);321 const value = ziggurat.next_f64(r, ziggurat.NormDist);
334 switch (T) {322 switch (T) {
335 f32 => return @as(f32, @floatCast(value)),323 f32 => return @floatCast(value),
336 f64 => return value,324 f64 => return value,
337 else => @compileError("unknown floating point type"),325 else => @compileError("unknown floating point type"),
338 }326 }
...@@ -344,7 +332,7 @@ pub const Random = struct {...@@ -344,7 +332,7 @@ pub const Random = struct {
344 pub fn floatExp(r: Random, comptime T: type) T {332 pub fn floatExp(r: Random, comptime T: type) T {
345 const value = ziggurat.next_f64(r, ziggurat.ExpDist);333 const value = ziggurat.next_f64(r, ziggurat.ExpDist);
346 switch (T) {334 switch (T) {
347 f32 => return @as(f32, @floatCast(value)),335 f32 => return @floatCast(value),
348 f64 => return value,336 f64 => return value,
349 else => @compileError("unknown floating point type"),337 else => @compileError("unknown floating point type"),
350 }338 }
...@@ -378,10 +366,10 @@ pub const Random = struct {...@@ -378,10 +366,10 @@ pub const Random = struct {
378 }366 }
379367
380 // `i <= j < max <= maxInt(MinInt)`368 // `i <= j < max <= maxInt(MinInt)`
381 const max = @as(MinInt, @intCast(buf.len));369 const max: MinInt = @intCast(buf.len);
382 var i: MinInt = 0;370 var i: MinInt = 0;
383 while (i < max - 1) : (i += 1) {371 while (i < max - 1) : (i += 1) {
384 const j = @as(MinInt, @intCast(r.intRangeLessThan(Index, i, max)));372 const j: MinInt = @intCast(r.intRangeLessThan(Index, i, max));
385 mem.swap(T, &buf[i], &buf[j]);373 mem.swap(T, &buf[i], &buf[j]);
386 }374 }
387 }375 }
...@@ -438,13 +426,12 @@ pub const Random = struct {...@@ -438,13 +426,12 @@ pub const Random = struct {
438pub fn limitRangeBiased(comptime T: type, random_int: T, less_than: T) T {426pub fn limitRangeBiased(comptime T: type, random_int: T, less_than: T) T {
439 comptime assert(@typeInfo(T).Int.signedness == .unsigned);427 comptime assert(@typeInfo(T).Int.signedness == .unsigned);
440 const bits = @typeInfo(T).Int.bits;428 const bits = @typeInfo(T).Int.bits;
441 const T2 = std.meta.Int(.unsigned, bits * 2);
442429
443 // adapted from:430 // adapted from:
444 // http://www.pcg-random.org/posts/bounded-rands.html431 // http://www.pcg-random.org/posts/bounded-rands.html
445 // "Integer Multiplication (Biased)"432 // "Integer Multiplication (Biased)"
446 var m: T2 = @as(T2, random_int) * @as(T2, less_than);433 const m = math.mulWide(T, random_int, less_than);
447 return @as(T, @intCast(m >> bits));434 return @intCast(m >> bits);
448}435}
449436
450// Generator to extend 64-bit seed values into longer sequences.437// Generator to extend 64-bit seed values into longer sequences.