| ... | @@ -1,6 +1,7 @@ | ... | @@ -1,6 +1,7 @@ |
| 1 | const assert = @import("debug.zig").assert; | 1 | const assert = @import("debug.zig").assert; |
| 2 | const rand_test = @import("rand_test.zig"); | 2 | const rand_test = @import("rand_test.zig"); |
| 3 | const mem = @import("mem.zig"); | 3 | const mem = @import("mem.zig"); |
| | 4 | const math = @import("math.zig"); |
| 4 | | 5 | |
| 5 | pub const MT19937_32 = MersenneTwister( | 6 | pub const MT19937_32 = MersenneTwister( |
| 6 | u32, 624, 397, 31, | 7 | u32, 624, 397, 31, |
| ... | @@ -63,18 +64,43 @@ pub const Rand = struct { | ... | @@ -63,18 +64,43 @@ pub const Rand = struct { |
| 63 | | 64 | |
| 64 | /// Get a random unsigned integer with even distribution between `start` | 65 | /// Get a random unsigned integer with even distribution between `start` |
| 65 | /// inclusive and `end` exclusive. | 66 | /// inclusive and `end` exclusive. |
| 66 | // TODO support signed integers and then rename to "range" | 67 | pub fn range(r: &Rand, comptime T: type, start: T, end: T) -> T { |
| 67 | pub fn rangeUnsigned(r: &Rand, comptime T: type, start: T, end: T) -> T { | 68 | assert(start <= end); |
| 68 | const range = end - start; | 69 | if (T.is_signed) { |
| 69 | const leftover = @maxValue(T) % range; | 70 | const uint = @IntType(false, T.bit_count); |
| 70 | const upper_bound = @maxValue(T) - leftover; | 71 | if (start >= 0 and end >= 0) { |
| 71 | var rand_val_array: [@sizeOf(T)]u8 = undefined; | 72 | return T(r.range(uint, uint(start), uint(end))); |
| 72 | | 73 | } else if (start < 0 and end < 0) { |
| 73 | while (true) { | 74 | // Can't overflow because the range is over signed ints |
| 74 | r.fillBytes(rand_val_array[0..]); | 75 | return %%math.negateCast(r.range(uint, math.absCast(end), math.absCast(start)) + 1); |
| 75 | const rand_val = mem.readInt(rand_val_array, T, false); | 76 | } else if (start < 0 and end >= 0) { |
| 76 | if (rand_val < upper_bound) { | 77 | const end_uint = uint(end); |
| 77 | return start + (rand_val % range); | 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,7 +120,7 @@ pub const Rand = struct { |
| 94 | } else { | 120 | } else { |
| 95 | @compileError("unknown floating point type") | 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,16 +201,38 @@ test "rand float 32" { |
| 175 | } | 201 | } |
| 176 | } | 202 | } |
| 177 | | 203 | |
| 178 | test "testMT19937_64" { | 204 | test "rand.MT19937_64" { |
| 179 | var rng = MT19937_64.init(rand_test.mt64_seed); | 205 | var rng = MT19937_64.init(rand_test.mt64_seed); |
| 180 | for (rand_test.mt64_data) |value| { | 206 | for (rand_test.mt64_data) |value| { |
| 181 | assert(value == rng.get()); | 207 | assert(value == rng.get()); |
| 182 | } | 208 | } |
| 183 | } | 209 | } |
| 184 | | 210 | |
| 185 | test "testMT19937_32" { | 211 | test "rand.MT19937_32" { |
| 186 | var rng = MT19937_32.init(rand_test.mt32_seed); | 212 | var rng = MT19937_32.init(rand_test.mt32_seed); |
| 187 | for (rand_test.mt32_data) |value| { | 213 | for (rand_test.mt32_data) |value| { |
| 188 | assert(value == rng.get()); | 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 | } |