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 {
113113 // TODO: endian portability is pointless if the underlying prng isn't endian portable.
114114 // TODO: document the endian portability of this library.
115115 const byte_aligned_result = mem.readIntSliceLittle(ByteAlignedT, &rand_bytes);
116 const unsigned_result = @as(UnsignedT, @truncate(byte_aligned_result));
117 return @as(T, @bitCast(unsigned_result));
116 const unsigned_result: UnsignedT = @truncate(byte_aligned_result);
117 return @bitCast(unsigned_result);
118118 }
119119
120120 /// Constant-time implementation off `uintLessThan`.
121121 /// The results of this function may be biased.
122122 pub fn uintLessThanBiased(r: Random, comptime T: type, less_than: T) T {
123123 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!
126124 assert(0 < less_than);
127 if (bits <= 32) {
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 }
125 return limitRangeBiased(T, r.int(T), less_than);
132126 }
133127
134128 /// Returns an evenly distributed random unsigned integer `0 <= i < less_than`.
......@@ -142,22 +136,16 @@ pub const Random = struct {
142136 pub fn uintLessThan(r: Random, comptime T: type, less_than: T) T {
143137 comptime assert(@typeInfo(T).Int.signedness == .unsigned);
144138 const bits = @typeInfo(T).Int.bits;
145 comptime assert(bits <= 64); // TODO: workaround: LLVM ERROR: Unsupported library call operation!
146139 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
153141 // adapted from:
154142 // http://www.pcg-random.org/posts/bounded-rands.html
155143 // "Lemire's (with an extra tweak from me)"
156 var x: Small = r.int(Small);
157 var m: Large = @as(Large, x) * @as(Large, less_than);
158 var l: Small = @as(Small, @truncate(m));
144 var x = r.int(T);
145 var m = math.mulWide(T, x, less_than);
146 var l: T = @truncate(m);
159147 if (l < less_than) {
160 var t: Small = -%less_than;
148 var t = -%less_than;
161149
162150 if (t >= less_than) {
163151 t -= less_than;
......@@ -166,12 +154,12 @@ pub const Random = struct {
166154 }
167155 }
168156 while (l < t) {
169 x = r.int(Small);
170 m = @as(Large, x) * @as(Large, less_than);
171 l = @as(Small, @truncate(m));
157 x = r.int(T);
158 m = math.mulWide(T, x, less_than);
159 l = @truncate(m);
172160 }
173161 }
174 return @as(T, @intCast(m >> small_bits));
162 return @intCast(m >> bits);
175163 }
176164
177165 /// Constant-time implementation off `uintAtMost`.
......@@ -205,10 +193,10 @@ pub const Random = struct {
205193 if (info.signedness == .signed) {
206194 // Two's complement makes this math pretty easy.
207195 const UnsignedT = std.meta.Int(.unsigned, info.bits);
208 const lo = @as(UnsignedT, @bitCast(at_least));
209 const hi = @as(UnsignedT, @bitCast(less_than));
196 const lo: UnsignedT = @bitCast(at_least);
197 const hi: UnsignedT = @bitCast(less_than);
210198 const result = lo +% r.uintLessThanBiased(UnsignedT, hi -% lo);
211 return @as(T, @bitCast(result));
199 return @bitCast(result);
212200 } else {
213201 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
214202 return at_least + r.uintLessThanBiased(T, less_than - at_least);
......@@ -224,10 +212,10 @@ pub const Random = struct {
224212 if (info.signedness == .signed) {
225213 // Two's complement makes this math pretty easy.
226214 const UnsignedT = std.meta.Int(.unsigned, info.bits);
227 const lo = @as(UnsignedT, @bitCast(at_least));
228 const hi = @as(UnsignedT, @bitCast(less_than));
215 const lo: UnsignedT = @bitCast(at_least);
216 const hi: UnsignedT = @bitCast(less_than);
229217 const result = lo +% r.uintLessThan(UnsignedT, hi -% lo);
230 return @as(T, @bitCast(result));
218 return @bitCast(result);
231219 } else {
232220 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
233221 return at_least + r.uintLessThan(T, less_than - at_least);
......@@ -242,10 +230,10 @@ pub const Random = struct {
242230 if (info.signedness == .signed) {
243231 // Two's complement makes this math pretty easy.
244232 const UnsignedT = std.meta.Int(.unsigned, info.bits);
245 const lo = @as(UnsignedT, @bitCast(at_least));
246 const hi = @as(UnsignedT, @bitCast(at_most));
233 const lo: UnsignedT = @bitCast(at_least);
234 const hi: UnsignedT = @bitCast(at_most);
247235 const result = lo +% r.uintAtMostBiased(UnsignedT, hi -% lo);
248 return @as(T, @bitCast(result));
236 return @bitCast(result);
249237 } else {
250238 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
251239 return at_least + r.uintAtMostBiased(T, at_most - at_least);
......@@ -261,10 +249,10 @@ pub const Random = struct {
261249 if (info.signedness == .signed) {
262250 // Two's complement makes this math pretty easy.
263251 const UnsignedT = std.meta.Int(.unsigned, info.bits);
264 const lo = @as(UnsignedT, @bitCast(at_least));
265 const hi = @as(UnsignedT, @bitCast(at_most));
252 const lo: UnsignedT = @bitCast(at_least);
253 const hi: UnsignedT = @bitCast(at_most);
266254 const result = lo +% r.uintAtMost(UnsignedT, hi -% lo);
267 return @as(T, @bitCast(result));
255 return @bitCast(result);
268256 } else {
269257 // The signed implementation would work fine, but we can use stricter arithmetic operators here.
270258 return at_least + r.uintAtMost(T, at_most - at_least);
......@@ -293,9 +281,9 @@ pub const Random = struct {
293281 rand_lz += @clz(r.int(u32) | 0x7FF);
294282 }
295283 }
296 const mantissa = @as(u23, @truncate(rand));
284 const mantissa: u23 = @truncate(rand);
297285 const exponent = @as(u32, 126 - rand_lz) << 23;
298 return @as(f32, @bitCast(exponent | mantissa));
286 return @bitCast(exponent | mantissa);
299287 },
300288 f64 => {
301289 // Use 52 random bits for the mantissa, and the rest for the exponent.
......@@ -320,7 +308,7 @@ pub const Random = struct {
320308 }
321309 const mantissa = rand & 0xFFFFFFFFFFFFF;
322310 const exponent = (1022 - rand_lz) << 52;
323 return @as(f64, @bitCast(exponent | mantissa));
311 return @bitCast(exponent | mantissa);
324312 },
325313 else => @compileError("unknown floating point type"),
326314 }
......@@ -332,7 +320,7 @@ pub const Random = struct {
332320 pub fn floatNorm(r: Random, comptime T: type) T {
333321 const value = ziggurat.next_f64(r, ziggurat.NormDist);
334322 switch (T) {
335 f32 => return @as(f32, @floatCast(value)),
323 f32 => return @floatCast(value),
336324 f64 => return value,
337325 else => @compileError("unknown floating point type"),
338326 }
......@@ -344,7 +332,7 @@ pub const Random = struct {
344332 pub fn floatExp(r: Random, comptime T: type) T {
345333 const value = ziggurat.next_f64(r, ziggurat.ExpDist);
346334 switch (T) {
347 f32 => return @as(f32, @floatCast(value)),
335 f32 => return @floatCast(value),
348336 f64 => return value,
349337 else => @compileError("unknown floating point type"),
350338 }
......@@ -378,10 +366,10 @@ pub const Random = struct {
378366 }
379367
380368 // `i <= j < max <= maxInt(MinInt)`
381 const max = @as(MinInt, @intCast(buf.len));
369 const max: MinInt = @intCast(buf.len);
382370 var i: MinInt = 0;
383371 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));
385373 mem.swap(T, &buf[i], &buf[j]);
386374 }
387375 }
......@@ -438,13 +426,12 @@ pub const Random = struct {
438426pub fn limitRangeBiased(comptime T: type, random_int: T, less_than: T) T {
439427 comptime assert(@typeInfo(T).Int.signedness == .unsigned);
440428 const bits = @typeInfo(T).Int.bits;
441 const T2 = std.meta.Int(.unsigned, bits * 2);
442429
443430 // adapted from:
444431 // http://www.pcg-random.org/posts/bounded-rands.html
445432 // "Integer Multiplication (Biased)"
446 var m: T2 = @as(T2, random_int) * @as(T2, less_than);
447 return @as(T, @intCast(m >> bits));
433 const m = math.mulWide(T, random_int, less_than);
434 return @intCast(m >> bits);
448435}
449436
450437// Generator to extend 64-bit seed values into longer sequences.