1use std::fmt;
2
3pub const STREAM_VERSION: u32 = 1;
4pub const STATE_BYTES: usize = 44;
5
6const GAMMA: u64 = 0x9e37_79b9_7f4a_7c15;
7const FORK_TAG: u64 = 0xd1b5_4a32_d192_ed03;
8const F64_UNIT: f64 = 1.0 / 9007199254740992.0;
9const F32_UNIT: f32 = 1.0 / 16777216.0;
10
11pub fn mix64(value: u64) -> u64 {
12 let z = value.wrapping_add(GAMMA);
13 let z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
14 let z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
15 z ^ (z >> 31)
16}
17
18pub fn splitmix64(state: &mut u64) -> u64 {
19 let out = mix64(*state);
20 *state = state.wrapping_add(GAMMA);
21 out
22}
23
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
25pub enum SimError {
26 Length { found: usize },
27 Version { found: u32 },
28 ZeroState,
29}
30
31impl fmt::Display for SimError {
32 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33 match self {
34 Self::Length { found } => {
35 write!(
36 f,
37 "saved rng state is {found} bytes, expected {STATE_BYTES}"
38 )
39 }
40 Self::Version { found } => write!(
41 f,
42 "saved rng state is stream version {found}, this build reads {STREAM_VERSION}"
43 ),
44 Self::ZeroState => write!(f, "saved rng state has no bits set in its four words"),
45 }
46 }
47}
48
49impl std::error::Error for SimError {}
50
51#[derive(Clone, Copy, Debug, PartialEq, Eq)]
52pub struct RngState {
53 pub root: u64,
54 pub words: [u64; 4],
55}
56
57impl RngState {
58 pub fn to_bytes(&self) -> [u8; STATE_BYTES] {
59 let words = if self.words == [0; 4] {
60 Rng::new(self.root).s
61 } else {
62 self.words
63 };
64 let mut out = [0u8; STATE_BYTES];
65 out[..4].copy_from_slice(&STREAM_VERSION.to_le_bytes());
66 out[4..12].copy_from_slice(&self.root.to_le_bytes());
67 for (i, word) in words.iter().enumerate() {
68 out[12 + i * 8..20 + i * 8].copy_from_slice(&word.to_le_bytes());
69 }
70 out
71 }
72
73 pub fn from_bytes(bytes: &[u8]) -> Result<Self, SimError> {
74 if bytes.len() != STATE_BYTES {
75 return Err(SimError::Length { found: bytes.len() });
76 }
77 let version = u32::from_le_bytes(bytes[..4].try_into().unwrap());
78 if version != STREAM_VERSION {
79 return Err(SimError::Version { found: version });
80 }
81 let word_at = |at: usize| u64::from_le_bytes(bytes[at..at + 8].try_into().unwrap());
82 let words = [word_at(12), word_at(20), word_at(28), word_at(36)];
83 if words == [0; 4] {
84 return Err(SimError::ZeroState);
85 }
86 Ok(Self {
87 root: word_at(4),
88 words,
89 })
90 }
91}
92
93#[derive(Clone, Debug, PartialEq, Eq)]
94pub struct Rng {
95 root: u64,
96 s: [u64; 4],
97}
98
99impl Rng {
100 pub fn new(seed: u64) -> Self {
101 let mut sm = seed;
102 let s = [
103 splitmix64(&mut sm),
104 splitmix64(&mut sm),
105 splitmix64(&mut sm),
106 splitmix64(&mut sm),
107 ];
108 Self { root: seed, s }
109 }
110
111 pub fn fork(&self, label: u64) -> Self {
112 Self::new(mix64(
113 mix64(self.root ^ FORK_TAG).wrapping_add(mix64(label)),
114 ))
115 }
116
117 pub fn state(&self) -> RngState {
118 RngState {
119 root: self.root,
120 words: self.s,
121 }
122 }
123
124 pub fn from_state(state: RngState) -> Self {
125 if state.words == [0; 4] {
126 return Self::new(state.root);
127 }
128 Self {
129 root: state.root,
130 s: state.words,
131 }
132 }
133
134 pub fn save(&self) -> [u8; STATE_BYTES] {
135 self.state().to_bytes()
136 }
137
138 pub fn restore(bytes: &[u8]) -> Result<Self, SimError> {
139 RngState::from_bytes(bytes).map(Self::from_state)
140 }
141
142 pub fn next_u64(&mut self) -> u64 {
143 let result = self.s[1].wrapping_mul(5).rotate_left(7).wrapping_mul(9);
144 let t = self.s[1] << 17;
145 self.s[2] ^= self.s[0];
146 self.s[3] ^= self.s[1];
147 self.s[1] ^= self.s[2];
148 self.s[0] ^= self.s[3];
149 self.s[2] ^= t;
150 self.s[3] = self.s[3].rotate_left(45);
151 result
152 }
153
154 pub fn next_u32(&mut self) -> u32 {
155 (self.next_u64() >> 32) as u32
156 }
157
158 pub fn next_bool(&mut self) -> bool {
159 self.next_u64() >> 63 == 1
160 }
161
162 pub fn next_f64(&mut self) -> f64 {
163 (self.next_u64() >> 11) as f64 * F64_UNIT
164 }
165
166 pub fn next_f32(&mut self) -> f32 {
167 (self.next_u64() >> 40) as f32 * F32_UNIT
168 }
169
170 pub fn chance(&mut self, probability: f64) -> bool {
171 self.next_f64() < probability
172 }
173
174 pub fn below(&mut self, n: u64) -> u64 {
175 assert!(n > 0, "below(0)");
176 let mut m = (self.next_u64() as u128) * (n as u128);
177 let mut low = m as u64;
178 if low < n {
179 let threshold = n.wrapping_neg() % n;
180 while low < threshold {
181 m = (self.next_u64() as u128) * (n as u128);
182 low = m as u64;
183 }
184 }
185 (m >> 64) as u64
186 }
187
188 pub fn below_u32(&mut self, n: u32) -> u32 {
189 self.below(n as u64) as u32
190 }
191
192 pub fn index(&mut self, len: usize) -> usize {
193 self.below(len as u64) as usize
194 }
195
196 pub fn range_u64(&mut self, lo: u64, hi: u64) -> u64 {
197 assert!(lo < hi, "range_u64 needs lo < hi");
198 lo + self.below(hi - lo)
199 }
200
201 pub fn range_i64(&mut self, lo: i64, hi: i64) -> i64 {
202 assert!(lo < hi, "range_i64 needs lo < hi");
203 lo.wrapping_add(self.below(hi.wrapping_sub(lo) as u64) as i64)
204 }
205
206 pub fn range_i32(&mut self, lo: i32, hi: i32) -> i32 {
207 self.range_i64(lo as i64, hi as i64) as i32
208 }
209
210 pub fn range_f64(&mut self, lo: f64, hi: f64) -> f64 {
211 lo + (hi - lo) * self.next_f64()
212 }
213
214 pub fn range_f32(&mut self, lo: f32, hi: f32) -> f32 {
215 lo + (hi - lo) * self.next_f32()
216 }
217
218 pub fn shuffle<T>(&mut self, items: &mut [T]) {
219 for i in (1..items.len()).rev() {
220 let j = self.below(i as u64 + 1) as usize;
221 items.swap(i, j);
222 }
223 }
224
225 pub fn choose<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> {
226 if items.is_empty() {
227 None
228 } else {
229 Some(&items[self.index(items.len())])
230 }
231 }
232
233 pub fn weighted_index(&mut self, weights: &[f64]) -> Option<usize> {
234 let mut total = 0.0;
235 for &w in weights {
236 if w.is_nan() || w < 0.0 || !w.is_finite() {
237 return None;
238 }
239 total += w;
240 }
241 if total.is_nan() || total <= 0.0 || !total.is_finite() {
242 return None;
243 }
244 let target = self.next_f64() * total;
245 let mut acc = 0.0;
246 let mut last = None;
247 for (i, &w) in weights.iter().enumerate() {
248 if w > 0.0 {
249 last = Some(i);
250 acc += w;
251 if target < acc {
252 return Some(i);
253 }
254 }
255 }
256 last
257 }
258
259 pub fn weighted_index_u32(&mut self, weights: &[u32]) -> Option<usize> {
260 let total: u64 = weights.iter().map(|&w| w as u64).sum();
261 if total == 0 {
262 return None;
263 }
264 let mut target = self.below(total);
265 for (i, &w) in weights.iter().enumerate() {
266 if target < w as u64 {
267 return Some(i);
268 }
269 target -= w as u64;
270 }
271 None
272 }
273}