random.ha (2387B)
1 // SPDX-License-Identifier: MPL-2.0 2 // (c) Hare authors <https://harelang.org> 3 4 // State for a pseudorandom number generator. 5 export type random = u64; 6 7 // Initializes a pseudorandom number generator with a given seed. This seed will 8 // yield the same sequence of psuedo-random numbers if used again. 9 export fn init(seed: u64) random = seed; 10 11 // Returns a psuedo-random 64-bit unsigned integer. 12 export fn next(r: *random) u64 = { 13 // SplitMix64 14 *r += 0x9e3779b97f4a7c15; 15 *r = (*r ^ *r >> 30) * 0xbf58476d1ce4e5b9; 16 *r = (*r ^ *r >> 27) * 0x94d049bb133111eb; 17 return *r ^ *r >> 31; 18 }; 19 20 // Returns a pseudo-random 32-bit unsigned integer in the half-open interval 21 // [0,n). n must be greater than zero. 22 export fn u32n(r: *random, n: u32) u32 = { 23 // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/ 24 // https://lemire.me/blog/2016/06/30/fast-random-shuffling/ 25 assert(n != 0); 26 let prod = next(r): u32: u64 * n: u64; 27 let leftover = prod: u32; 28 if (leftover < n) { 29 let thresh = -n % n; 30 for (leftover < thresh) { 31 prod = next(r): u32: u64 * n: u64; 32 leftover = prod: u32; 33 }; 34 }; 35 return (prod >> 32): u32; 36 }; 37 38 // Returns a pseudo-random 64-bit unsigned integer in the half-open interval 39 // [0,n). n must be greater than zero. 40 export fn u64n(r: *random, n: u64) u64 = { 41 assert(n != 0); 42 // Powers of 2 can be handled more efficiently 43 if (n & n - 1 == 0) return next(r) & n - 1; 44 // Equivalent to 2^64 - 1 - 2^64 % n 45 let max = -1 - -n % n; 46 let out = next(r); 47 for (out > max) out = next(r); 48 return out % n; 49 }; 50 51 // Returns a pseudo-random 64-bit floating-point number in the interval 52 // [0.0, 1.0) 53 export fn f64rand(r: *random) f64 = { 54 // 1.0 x 2^-53 55 const d: f64 = 1.1102230246251565e-16; 56 // Take the upper 53 bits 57 let n = next(r) >> 11; 58 59 return d * n: f64; 60 }; 61 62 @test fn rng() void = { 63 let r = init(0); 64 let expected: [_]u64 = [ 65 16294208416658607535, 66 15501543990041496116, 67 15737388954706874752, 68 15091258616627000950, 69 ]; 70 for (let i = 0z; i < len(expected); i += 1) { 71 assert(next(&r) == expected[i]); 72 }; 73 74 for (let i = 0z; i < 1000; i += 1) { 75 assert(u32n(&r, 3) < 3); 76 }; 77 78 for (let i = 0z; i < 1000; i += 1) { 79 assert(u64n(&r, 3) < 3); 80 }; 81 for (let i = 0z; i < 1000; i += 1) { 82 // Powers of 2 have a separate codepath 83 assert(u64n(&r, 2) < 2); 84 }; 85 for (let i = 0z; i < 1000; i += 1) { 86 assert(f64rand(&r) < 1.0); 87 }; 88 };