tiny.simd.random

1const std = @import("std");2const crypto = @import("crypto.zig");3const rotate = @import("rotate.zig");4const shift = @import("shift.zig");5const tag = @import("tag.zig");67const jump_constants = [4]u64{8    0x180e_c6d3_3cfd_0aba,9    0xd5a6_1266_f0c9_392c,10    0xa958_2618_e03f_c9aa,11    0x39ab_dc45_29b1_661c,12};1314const long_jump_constants = [4]u64{15    0x76e1_5d3e_fefd_cbbf,16    0xc500_4e44_1c52_2fb3,17    0x7771_0069_854e_e241,18    0x3910_9bb0_2acb_e635,19};2021const uniform_scale: f64 = 0x1.0p-53;2223pub const SplitMix64 = struct {24    state: u64,2526    pub fn init(state: u64) SplitMix64 {27        return .{ .state = state };28    }2930    pub fn next(self: *SplitMix64) u64 {31        self.state +%= 0x9e37_79b9_7f4a_7c15;32        var value = self.state;33        value = (value ^ (value >> 30)) *% 0xbf58_476d_1ce4_e5b9;34        value = (value ^ (value >> 27)) *% 0x94d0_49bb_1331_11eb;35        return value ^ (value >> 31);36    }37};3839pub const Xoshiro = struct {40    state: [4]u64,4142    pub fn init(seed: u64) Xoshiro {43        var split = SplitMix64.init(seed);44        var state: [4]u64 = undefined;45        for (&state) |*value| value.* = split.next();46        return .{ .state = state };47    }4849    pub fn initThread(seed: u64, thread_id: u64) Xoshiro {50        var result = init(seed);51        var stream: u64 = 0;52        while (stream < thread_id) : (stream += 1) result.jump();53        return result;54    }5556    pub fn next(self: *Xoshiro) u64 {57        const result = std.math.rotl(u64, self.state[0] +% self.state[3], 23) +%58            self.state[0];59        const temporary = self.state[1] << 17;60        self.state[2] ^= self.state[0];61        self.state[3] ^= self.state[1];62        self.state[1] ^= self.state[2];63        self.state[0] ^= self.state[3];64        self.state[2] ^= temporary;65        self.state[3] = std.math.rotl(u64, self.state[3], 45);66        return result;67    }6869    pub fn uniform(self: *Xoshiro) f64 {70        return @as(f64, @floatFromInt(self.next() >> 11)) * uniform_scale;71    }7273    pub fn getState(self: Xoshiro) [4]u64 {74        return self.state;75    }7677    pub fn setState(self: *Xoshiro, state: [4]u64) void {78        self.state = state;79    }8081    pub fn stateSize() usize {82        return 4;83    }8485    pub fn jump(self: *Xoshiro) void {86        self.jumpWith(jump_constants);87    }8889    pub fn longJump(self: *Xoshiro) void {90        self.jumpWith(long_jump_constants);91    }9293    fn jumpWith(self: *Xoshiro, constants: [4]u64) void {94        var state: [4]u64 = @splat(0);95        for (constants) |bits| {96            for (0..64) |bit| {97                if (bits & (@as(u64, 1) << @intCast(bit)) != 0) {98                    inline for (0..4) |index| state[index] ^= self.state[index];99                }100                _ = self.next();101            }102        }103        self.state = state;104    }105};106107pub fn VectorXoshiro(comptime D: type) type {108    requireU64(D);109    return struct {110        state: [4]D.Vector,111112        const Self = @This();113        const DF: type = D.rebind(f64);114115        pub fn init(seed: u64) Self {116            return initThread(seed, 0);117        }118119        pub fn initThread(seed: u64, thread_number: u64) Self {120            var scalar = Xoshiro.init(seed);121            var stream: u64 = 0;122            while (stream < thread_number) : (stream += 1) scalar.longJump();123            var state: [4]D.Vector = @splat(@splat(0));124            inline for (0..D.lane_count) |lane| {125                const scalar_state = scalar.getState();126                inline for (0..4) |index| state[index][lane] = scalar_state[index];127                scalar.jump();128            }129            return .{ .state = state };130        }131132        pub fn next(self: *Self) D.Vector {133            return update(&self.state);134        }135136        pub fn fill(self: *Self, output: []u64) void {137            var index: usize = 0;138            while (index + D.lane_count <= output.len) : (index += D.lane_count) {139                const value: [D.lane_count]u64 = self.next();140                @memcpy(output[index .. index + D.lane_count], &value);141            }142            if (index != output.len) {143                const value: [D.lane_count]u64 = self.next();144                @memcpy(output[index..], value[0 .. output.len - index]);145            }146        }147148        pub fn uniform(self: *Self) DF.Vector {149            return toUniform(self.next());150        }151152        pub fn fillUniform(self: *Self, output: []f64) void {153            var index: usize = 0;154            while (index + D.lane_count <= output.len) : (index += D.lane_count) {155                const value: [D.lane_count]f64 = self.uniform();156                @memcpy(output[index .. index + D.lane_count], &value);157            }158            if (index != output.len) {159                const value: [D.lane_count]f64 = self.uniform();160                @memcpy(output[index..], value[0 .. output.len - index]);161            }162        }163164        pub fn getState(self: Self) [4]D.Vector {165            return self.state;166        }167168        pub fn setState(self: *Self, state: [4]D.Vector) void {169            self.state = state;170        }171172        pub fn stateSize() usize {173            return 4 * D.lane_count;174        }175176        fn update(state: *[4]D.Vector) D.Vector {177            const result = rotate.rotateRightSame(D, state[0] +% state[3], 41) +% state[0];178            const temporary = shift.shiftLeft(D, 17, state[1]);179            state[2] ^= state[0];180            state[3] ^= state[1];181            state[1] ^= state[2];182            state[0] ^= state[3];183            state[2] ^= temporary;184            state[3] = rotate.rotateRightSame(D, state[3], 19);185            return result;186        }187188        fn toUniform(value: D.Vector) DF.Vector {189            const bits = shift.shiftRight(D, 11, value);190            var result: DF.Vector = undefined;191            inline for (0..D.lane_count) |index| {192                result[index] = @as(f64, @floatFromInt(bits[index])) * uniform_scale;193            }194            return result;195        }196    };197}198199pub fn CachedXoshiro(comptime D: type, comptime cache_size: usize) type {200    if (cache_size == 0 or !std.math.isPowerOfTwo(cache_size)) {201        @compileError("cached xoshiro size must be a nonzero power of two");202    }203    return struct {204        generator: Generator,205        cache: [cache_size]u64,206        index: usize,207208        const Self = @This();209        const Generator: type = VectorXoshiro(D);210211        pub fn init(seed: u64) Self {212            return initThread(seed, 0);213        }214215        pub fn initThread(seed: u64, thread_number: u64) Self {216            var result = Self{217                .generator = Generator.initThread(seed, thread_number),218                .cache = undefined,219                .index = 0,220            };221            result.generator.fill(&result.cache);222            return result;223        }224225        pub fn next(self: *Self) u64 {226            std.debug.assert(self.index <= cache_size);227            if (self.index == cache_size) {228                self.generator.fill(&self.cache);229                self.index = 0;230            }231            const result = self.cache[self.index];232            self.index += 1;233            return result;234        }235236        pub fn minimum() u64 {237            return 0;238        }239240        pub fn maximum() u64 {241            return std.math.maxInt(u64);242        }243    };244}245246pub fn DefaultCachedXoshiro(comptime D: type) type {247    return CachedXoshiro(D, 1024);248}249250pub const AesCtrEngine = struct {251    key: [2 * (1 + rounds)]u64,252253    const rounds: usize = 5;254    const D: type = tag.Full128(u8);255    const D64: type = tag.Full128(u64);256257    pub fn initDeterministic() AesCtrEngine {258        return initKey(.{ 0x243f_6a88_85a3_08d3, 0x1319_8a2e_0370_7344 });259    }260261    pub fn initKey(seed: [2]u64) AesCtrEngine {262        var result: AesCtrEngine = undefined;263        result.key[0] = seed[0];264        result.key[1] = seed[1];265        inline for (0..rounds) |index| {266            result.key[2 + 2 * index] = result.key[2 * index + 1] +%267                0xa409_3822_299f_31d0;268            result.key[2 + 2 * index + 1] = result.key[2 * index] +%269                0x082e_fa98_ec4e_6c89;270        }271        return result;272    }273274    pub fn initWithEntropy(entropy: anytype) !AesCtrEngine {275        var bytes: [16]u8 = undefined;276        try entropy.fill(&bytes);277        return initKey(@bitCast(bytes));278    }279280    pub fn generate(self: *const AesCtrEngine, stream: u64, counter: u64) u64 {281        var state: D.Vector = @bitCast(@as(D64.Vector, .{ counter, stream }));282        state ^= keyVector(self.key[0..2].*);283        inline for (0..rounds) |index| {284            state = crypto.aesRound(D, state, keyVector(self.key[2 + 2 * index ..][0..2].*));285        }286        const result: D64.Vector = @bitCast(state);287        return result[0];288    }289290    fn keyVector(pair: [2]u64) D.Vector {291        return @bitCast(@as(D64.Vector, pair));292    }293};294295pub const RngStream = struct {296    engine: *const AesCtrEngine,297    stream: u64,298    counter: u64 = 0,299300    pub fn init(engine: *const AesCtrEngine, stream: u64) RngStream {301        return .{ .engine = engine, .stream = stream };302    }303304    pub fn next(self: *RngStream) u64 {305        const result = self.engine.generate(self.stream, self.counter);306        self.counter +%= 1;307        return result;308    }309310    pub fn minimum() u64 {311        return 0;312    }313314    pub fn maximum() u64 {315        return std.math.maxInt(u64);316    }317};318319pub fn randomNormalizedFloat(stream: *RngStream) f32 {320    const exponent = @as(u32, @bitCast(@as(f32, 1)));321    const mantissa = @as(u32, @truncate(stream.next())) & 0x007f_ffff;322    const value: f32 = @bitCast(exponent | mantissa);323    std.debug.assert(value >= 1);324    std.debug.assert(value < 2);325    const result = (2 * (value - 1)) - 1;326    std.debug.assert(result >= -1);327    std.debug.assert(result < 1);328    return result;329}330331pub fn fillRandom(comptime T: type, engine: *const AesCtrEngine, stream_id: u64, output: []T) void {332    var stream = RngStream.init(engine, stream_id);333    for (output) |*value| value.* = castRandom(T, stream.next());334}335336fn castRandom(comptime T: type, value: u64) T {337    return switch (@typeInfo(T)) {338        .int => |info| if (info.signedness == .unsigned)339            @truncate(value)340        else blk: {341            const U = @Int(.unsigned, info.bits);342            break :blk @bitCast(@as(U, @truncate(value)));343        },344        .float => @floatFromInt(value),345        else => @compileError("fillRandom requires integer or floating-point outputs"),346    };347}348349fn requireU64(comptime D: type) void {350    if (comptime D.Lane != u64) @compileError("vector xoshiro requires u64 lanes");351}352353const TestEntropy = struct {354    offset: u8,355356    fn fill(self: *@This(), output: []u8) !void {357        for (output, 0..) |*byte, index| byte.* = @intCast(index * 7 + self.offset);358    }359};360361test "Highway scalar xoshiro state replay and stream jumps are deterministic" {362    var generator = Xoshiro.init(123);363    const state = generator.getState();364    const first = generator.next();365    const second = generator.next();366    try std.testing.expect(first != second);367    generator.setState(state);368    try std.testing.expectEqual(first, generator.next());369    try std.testing.expectEqual(second, generator.next());370    var jumped = Xoshiro.initThread(123, 1);371    var reference = Xoshiro.init(123);372    reference.jump();373    try std.testing.expectEqual(reference.getState(), jumped.getState());374    try std.testing.expectEqual(@as(usize, 4), Xoshiro.stateSize());375}376377test "Highway AVX2 scalar and vector xoshiro oracle matches exactly" {378    var scalar = Xoshiro.init(123);379    try std.testing.expectEqual(380        [4]u64{381            0xb4dc_9bd4_62de_412b,382            0xfa02_3ce9_f06f_b77c,383            0xdc12_d311_d371_cbe8,384            0xafd2_040c_9098_81ff,385        },386        scalar.getState(),387    );388    for ([_]u64{389        0xa556_5735_f810_987a,390        0xd691_4642_e58d_662e,391        0xaa75_21fe_b709_887f,392        0x863c_d15c_558d_6bfb,393    }) |expected| try std.testing.expectEqual(expected, scalar.next());394    var jumped = Xoshiro.init(123);395    jumped.jump();396    try std.testing.expectEqual(397        [4]u64{398            0xed6f_8f0b_3989_38af,399            0x17c5_be01_9e09_5507,400            0x05a0_a7d3_71bf_d778,401            0xbc77_948a_a522_4033,402        },403        jumped.getState(),404    );405    var long_jumped = Xoshiro.init(123);406    long_jumped.longJump();407    try std.testing.expectEqual(408        [4]u64{409            0xf98f_04cb_080b_8942,410            0xd884_97e0_e599_e6f1,411            0x3b4f_07f1_f6d2_6dd4,412            0x5d3e_2cd5_d040_18c1,413        },414        long_jumped.getState(),415    );416    const D = tag.FixedTag(u64, 4);417    const Generator = VectorXoshiro(D);418    var vector = Generator.init(456);419    try std.testing.expect(@reduce(.And, vector.next() == @as(D.Vector, .{420        0xa901_0f72_89ae_bac2,421        0x6990_d4cc_67a1_47d7,422        0x76bd_c3fa_db67_4b7c,423        0x6580_3fd8_b3a1_a35d,424    })));425    try std.testing.expect(@reduce(.And, vector.next() == @as(D.Vector, .{426        0x5fcc_d4ec_4e33_71f4,427        0x5321_fc38_a086_1200,428        0x0bce_898c_528d_1264,429        0xf716_df28_1d99_3979,430    })));431    var uniform = Generator.init(456);432    const real: [4]f64 = uniform.uniform();433    try std.testing.expectEqual(434        [4]u64{435            0x3fe5_2021_ee51_35d7,436            0x3fda_6435_3319_e850,437            0x3fdd_af70_feb6_d9d2,438            0x3fd9_600f_f62c_e868,439        },440        @as([4]u64, @bitCast(real)),441    );442}443444test "Highway vector xoshiro interleaves jump-separated scalar streams" {445    const D = tag.FixedTag(u64, 4);446    const Generator = VectorXoshiro(D);447    var generator = Generator.init(456);448    var streams: [D.lane_count]Xoshiro = undefined;449    streams[0] = Xoshiro.init(456);450    for (1..D.lane_count) |index| {451        streams[index] = streams[index - 1];452        streams[index].jump();453    }454    for (0..7) |_| {455        const actual: [D.lane_count]u64 = generator.next();456        for (&streams, 0..) |*stream, lane| {457            try std.testing.expectEqual(stream.next(), actual[lane]);458        }459    }460    try std.testing.expectEqual(@as(usize, 16), Generator.stateSize());461}462463test "Highway vector xoshiro thread streams begin at long-jump partitions" {464    const D = tag.FixedTag(u64, 4);465    const Generator = VectorXoshiro(D);466    var generator = Generator.initThread(654, 3);467    var reference = Xoshiro.init(654);468    for (0..3) |_| reference.longJump();469    const state = generator.getState();470    inline for (0..D.lane_count) |lane| {471        const scalar_state = reference.getState();472        inline for (0..4) |index| {473            try std.testing.expectEqual(scalar_state[index], state[index][lane]);474        }475        reference.jump();476    }477}478479test "Highway vector xoshiro fills awkward tails and uniform doubles" {480    const D = tag.FixedTag(u64, 4);481    const Generator = VectorXoshiro(D);482    var integers = Generator.init(789);483    var expected = Generator.init(789);484    var output: [13]u64 = undefined;485    integers.fill(&output);486    for (0..4) |block| {487        const values: [D.lane_count]u64 = expected.next();488        const begin = block * D.lane_count;489        const count = @min(D.lane_count, output.len - begin);490        for (0..count) |lane| try std.testing.expectEqual(values[lane], output[begin + lane]);491    }492    var uniforms = Generator.init(789);493    var real: [13]f64 = undefined;494    uniforms.fillUniform(&real);495    for (real) |value| try std.testing.expect(value >= 0 and value < 1);496}497498test "Highway cached xoshiro refills without changing stream order" {499    const D = tag.FixedTag(u64, 4);500    const Cached = CachedXoshiro(D, 8);501    const Generator = VectorXoshiro(D);502    var cached = Cached.init(987);503    var vector = Generator.init(987);504    var expected: [24]u64 = undefined;505    vector.fill(&expected);506    for (expected) |value| try std.testing.expectEqual(value, cached.next());507    try std.testing.expectEqual(@as(u64, 0), Cached.minimum());508    try std.testing.expectEqual(std.math.maxInt(u64), Cached.maximum());509}510511test "Highway AES counter streams are deterministic and independently addressed" {512    const engine = AesCtrEngine.initDeterministic();513    var first = RngStream.init(&engine, 7);514    var replay = RngStream.init(&engine, 7);515    var other = RngStream.init(&engine, 8);516    for (0..16) |_| {517        try std.testing.expectEqual(first.next(), replay.next());518        try std.testing.expect(first.counter == replay.counter);519    }520    try std.testing.expect(engine.generate(7, 0) != other.next());521}522523test "Highway AVX2 AES counter and normalized float oracle matches exactly" {524    const engine = AesCtrEngine.initDeterministic();525    try std.testing.expectEqual(@as(u64, 0xf0ea_ad2f_c4a2_c3e1), engine.generate(7, 0));526    try std.testing.expectEqual(@as(u64, 0xbf34_0afe_30f5_275d), engine.generate(7, 1));527    try std.testing.expectEqual(@as(u64, 0x0797_7286_13ce_fc88), engine.generate(8, 0));528    var stream = RngStream.init(&engine, 0);529    for ([_]u32{ 0xbed6_6698, 0x3e46_f1b0, 0x3f17_03d4 }) |bits| {530        try std.testing.expectEqual(bits, @as(u32, @bitCast(randomNormalizedFloat(&stream))));531    }532}533534test "Highway AES counter accepts caller entropy and fills typed buffers" {535    var entropy = TestEntropy{ .offset = 3 };536    const engine = try AesCtrEngine.initWithEntropy(&entropy);537    var other_entropy = TestEntropy{ .offset = 11 };538    const other_engine = try AesCtrEngine.initWithEntropy(&other_entropy);539    try std.testing.expect(engine.generate(0, 0) != other_engine.generate(0, 0));540    var actual: [17]u16 = undefined;541    fillRandom(u16, &engine, 91, &actual);542    var stream = RngStream.init(&engine, 91);543    for (actual) |value| try std.testing.expectEqual(@as(u16, @truncate(stream.next())), value);544}545546test "Highway normalized float spans the half-open signed unit interval" {547    const engine = AesCtrEngine.initDeterministic();548    var stream = RngStream.init(&engine, 0);549    var sum: f64 = 0;550    for (0..20_000) |_| {551        const value = randomNormalizedFloat(&stream);552        try std.testing.expect(value >= -1 and value < 1);553        sum += value;554    }555    try std.testing.expect(@abs(sum / 20_000) < 0.02);556}557558test "Highway AES counter bit and byte distributions remain balanced" {559    const engine = AesCtrEngine.initDeterministic();560    var stream = RngStream.init(&engine, 0);561    const count: usize = 20_000;562    var one_bits: u64 = 0;563    for (0..count) |_| one_bits += @popCount(stream.next());564    const bit_ratio = @as(f64, @floatFromInt(one_bits)) /565        @as(f64, @floatFromInt(count * 64));566    try std.testing.expect(bit_ratio >= 0.49 and bit_ratio <= 0.51);567    for (0..8) |byte_index| {568        var counts: [256]u32 = @splat(0);569        const amount: u6 = @intCast(byte_index * 8);570        for (0..count) |_| counts[@as(u8, @truncate(stream.next() >> amount))] += 1;571        const expected = @as(f64, @floatFromInt(count)) / 256;572        var chi_squared: f64 = 0;573        for (counts) |observed| {574            const difference = @as(f64, @floatFromInt(observed)) - expected;575            chi_squared += difference * difference / expected;576        }577        try std.testing.expect(chi_squared >= 170 and chi_squared <= 340);578    }579}

1pub const random = @import("random.zig");