| ... | ... | @@ -1,6 +1,7 @@ |
| 1 | 1 | const assert = @import("debug.zig").assert; |
| 2 | 2 | const rand_test = @import("rand_test.zig"); |
| 3 | 3 | const mem = @import("mem.zig"); |
| 4 | const math = @import("math.zig"); |
| 4 | 5 | |
| 5 | 6 | pub const MT19937_32 = MersenneTwister( |
| 6 | 7 | u32, 624, 397, 31, |
| ... | ... | @@ -63,18 +64,43 @@ pub const Rand = struct { |
| 63 | 64 | |
| 64 | 65 | /// Get a random unsigned integer with even distribution between `start` |
| 65 | 66 | /// inclusive and `end` exclusive. |
| 66 | | // TODO support signed integers and then rename to "range" |
| 67 | | pub fn rangeUnsigned(r: &Rand, comptime T: type, start: T, end: T) -> T { |
| 68 | | const range = end - start; |
| 69 | | const leftover = @maxValue(T) % range; |
| 70 | | const upper_bound = @maxValue(T) - leftover; |
| 71 | | var rand_val_array: [@sizeOf(T)]u8 = undefined; |
| 72 | | |
| 73 | | while (true) { |
| 74 | | r.fillBytes(rand_val_array[0..]); |
| 75 | | const rand_val = mem.readInt(rand_val_array, T, false); |
| 76 | | if (rand_val < upper_bound) { |
| 77 | | return start + (rand_val % range); |
| 67 | pub fn range(r: &Rand, comptime T: type, start: T, end: T) -> T { |
| 68 | assert(start <= end); |
| 69 | if (T.is_signed) { |
| 70 | const uint = @IntType(false, T.bit_count); |
| 71 | if (start >= 0 and end >= 0) { |
| 72 | return T(r.range(uint, uint(start), uint(end))); |
| 73 | } else if (start < 0 and end < 0) { |
| 74 | // Can't overflow because the range is over signed ints |
| 75 | return %%math.negateCast(r.range(uint, math.absCast(end), math.absCast(start)) + 1); |
| 76 | } else if (start < 0 and end >= 0) { |
| 77 | const end_uint = uint(end); |
| 78 | const total_range = math.absCast(start) + end_uint; |
| 79 | const value = r.range(uint, 0, total_range); |
| 80 | const result = if (value < end_uint) { |
| 81 | T(value) |
| 82 | } else if (value == end_uint) { |
| 83 | start |
| 84 | } else { |
| 85 | // Can't overflow because the range is over signed ints |
| 86 | %%math.negateCast(value - end_uint) |
| 87 | }; |
| 88 | return result; |
| 89 | } else { |
| 90 | unreachable; |
| 91 | } |
| 92 | } else { |
| 93 | const total_range = end - start; |
| 94 | const leftover = @maxValue(T) % total_range; |
| 95 | const upper_bound = @maxValue(T) - leftover; |
| 96 | var rand_val_array: [@sizeOf(T)]u8 = undefined; |
| 97 | |
| 98 | while (true) { |
| 99 | r.fillBytes(rand_val_array[0..]); |
| 100 | const rand_val = mem.readInt(rand_val_array, T, false); |
| 101 | if (rand_val < upper_bound) { |
| 102 | return start + (rand_val % total_range); |
| 103 | } |
| 78 | 104 | } |
| 79 | 105 | } |
| 80 | 106 | } |
| ... | ... | @@ -94,7 +120,7 @@ pub const Rand = struct { |
| 94 | 120 | } else { |
| 95 | 121 | @compileError("unknown floating point type") |
| 96 | 122 | }; |
| 97 | | return T(r.rangeUnsigned(int_type, 0, precision)) / T(precision); |
| 123 | return T(r.range(int_type, 0, precision)) / T(precision); |
| 98 | 124 | } |
| 99 | 125 | }; |
| 100 | 126 | |
| ... | ... | @@ -175,16 +201,38 @@ test "rand float 32" { |
| 175 | 201 | } |
| 176 | 202 | } |
| 177 | 203 | |
| 178 | | test "testMT19937_64" { |
| 204 | test "rand.MT19937_64" { |
| 179 | 205 | var rng = MT19937_64.init(rand_test.mt64_seed); |
| 180 | 206 | for (rand_test.mt64_data) |value| { |
| 181 | 207 | assert(value == rng.get()); |
| 182 | 208 | } |
| 183 | 209 | } |
| 184 | 210 | |
| 185 | | test "testMT19937_32" { |
| 211 | test "rand.MT19937_32" { |
| 186 | 212 | var rng = MT19937_32.init(rand_test.mt32_seed); |
| 187 | 213 | for (rand_test.mt32_data) |value| { |
| 188 | 214 | assert(value == rng.get()); |
| 189 | 215 | } |
| 190 | 216 | } |
| 217 | |
| 218 | test "rand.Rand.range" { |
| 219 | var r = Rand.init(42); |
| 220 | testRange(&r, -4, 3); |
| 221 | testRange(&r, -4, -1); |
| 222 | testRange(&r, 10, 14); |
| 223 | } |
| 224 | |
| 225 | fn testRange(r: &Rand, start: i32, end: i32) { |
| 226 | const count = usize(end - start); |
| 227 | var values_buffer = []bool{false} ** 20; |
| 228 | const values = values_buffer[0..count]; |
| 229 | var i: usize = 0; |
| 230 | while (i < count) { |
| 231 | const value = r.range(i32, start, end); |
| 232 | const index = usize(value - start); |
| 233 | if (!values[index]) { |
| 234 | i += 1; |
| 235 | values[index] = true; |
| 236 | } |
| 237 | } |
| 238 | } |