1#[derive(Debug, Clone)]
8pub struct SplitMix64 {
9 state: u64,
10}
11
12impl SplitMix64 {
13 pub fn new(seed: u64) -> Self {
15 Self { state: seed }
16 }
17
18 pub fn next_u64(&mut self) -> u64 {
20 self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
21 mix64(self.state)
22 }
23}
24
25#[derive(Debug, Clone)]
27pub struct Xoroshiro128Plus {
28 s0: u64,
29 s1: u64,
30}
31
32impl Xoroshiro128Plus {
33 pub fn new(seed: u64) -> Self {
35 let mut seeder = SplitMix64::new(seed);
36 let s0 = seeder.next_u64();
37 let mut s1 = seeder.next_u64();
38 if s0 == 0 && s1 == 0 {
39 s1 = 1;
40 }
41 Self { s0, s1 }
42 }
43
44 pub fn next_u64(&mut self) -> u64 {
46 let s0 = self.s0;
47 let mut s1 = self.s1;
48 let result = s0.wrapping_add(s1);
49
50 s1 ^= s0;
51 self.s0 = s0.rotate_left(24) ^ s1 ^ (s1 << 16);
52 self.s1 = s1.rotate_left(37);
53
54 result
55 }
56
57 pub fn next_f64(&mut self) -> f64 {
59 let bits = self.next_u64() >> 11;
60 bits as f64 * (1.0 / (1u64 << 53) as f64)
61 }
62}
63
64pub mod bucket {
66 pub fn assign_u64(key: &[u8], experiment: &[u8], salt: &[u8]) -> u64 {
68 let mut state = 0xCBF2_9CE4_8422_2325u64;
69 state = mix_bytes(state, key);
70 state = mix_bytes(state, &[0xFF]);
71 state = mix_bytes(state, experiment);
72 state = mix_bytes(state, &[0xFE]);
73 state = mix_bytes(state, salt);
74 super::mix64(state)
75 }
76
77 pub fn assign_ratio(key: &[u8], experiment: &[u8], salt: &[u8]) -> f64 {
79 let bits = assign_u64(key, experiment, salt) >> 11;
80 bits as f64 * (1.0 / (1u64 << 53) as f64)
81 }
82
83 pub fn assign_bucket(key: &[u8], experiment: &[u8], salt: &[u8], bucket_count: u64) -> u64 {
85 assert!(bucket_count > 0, "bucket_count must be positive");
86 assign_u64(key, experiment, salt) % bucket_count
87 }
88
89 fn mix_bytes(mut state: u64, bytes: &[u8]) -> u64 {
90 for &byte in bytes {
91 state ^= u64::from(byte);
92 state = state.wrapping_mul(0x0000_0100_0000_01B3);
93 }
94 state
95 }
96}
97
98#[inline]
99fn mix64(mut z: u64) -> u64 {
100 z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
101 z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
102 z ^ (z >> 31)
103}
104
105#[cfg(test)]
106mod tests {
107 use super::*;
108
109 #[test]
110 fn splitmix_reproducible() {
111 let mut a = SplitMix64::new(42);
112 let mut b = SplitMix64::new(42);
113 for _ in 0..128 {
114 assert_eq!(a.next_u64(), b.next_u64());
115 }
116 }
117
118 #[test]
119 fn xoroshiro_reproducible() {
120 let mut a = Xoroshiro128Plus::new(1234);
121 let mut b = Xoroshiro128Plus::new(1234);
122 for _ in 0..128 {
123 assert_eq!(a.next_u64(), b.next_u64());
124 }
125 }
126
127 #[test]
128 fn xoroshiro_next_f64_in_range() {
129 let mut rng = Xoroshiro128Plus::new(999);
130 for _ in 0..10_000 {
131 let value = rng.next_f64();
132 assert!((0.0..1.0).contains(&value), "value={value}");
133 }
134 }
135
136 #[test]
137 fn bucket_ratio_is_in_unit_interval() {
138 for seed in 0..256u64 {
139 let ratio = bucket::assign_ratio(&seed.to_le_bytes(), b"exp-a", b"salt");
140 assert!((0.0..1.0).contains(&ratio), "ratio={ratio}");
141 }
142 }
143
144 #[test]
145 fn bucket_assignment_is_stable() {
146 let a = bucket::assign_bucket(b"user-1", b"exp-a", b"salt", 100);
147 let b = bucket::assign_bucket(b"user-1", b"exp-a", b"salt", 100);
148 assert_eq!(a, b);
149 }
150
151 #[test]
152 fn bucket_assignment_respects_bucket_bounds() {
153 for bucket_count in [1u64, 2, 7, 100] {
154 let bucket = bucket::assign_bucket(b"user-1", b"exp-a", b"salt", bucket_count);
155 assert!(
156 bucket < bucket_count,
157 "bucket={bucket}, bucket_count={bucket_count}"
158 );
159 }
160 }
161
162 #[test]
163 fn bucket_assignment_changes_when_experiment_changes() {
164 let a = bucket::assign_u64(b"user-1", b"exp-a", b"salt");
165 let b = bucket::assign_u64(b"user-1", b"exp-b", b"salt");
166 assert_ne!(a, b);
167 }
168
169 #[test]
170 fn bucket_assignment_changes_when_salt_changes() {
171 let a = bucket::assign_u64(b"user-1", b"exp-a", b"salt-a");
172 let b = bucket::assign_u64(b"user-1", b"exp-a", b"salt-b");
173 assert_ne!(a, b);
174 }
175
176 #[test]
177 #[should_panic(expected = "bucket_count must be positive")]
178 fn bucket_assignment_rejects_zero_bucket_count() {
179 let _ = bucket::assign_bucket(b"user-1", b"exp-a", b"salt", 0);
180 }
181}