doge_runtime/stdlib/
roll.rs1use std::cell::RefCell;
12use std::rc::Rc;
13use std::sync::atomic::{AtomicU64, Ordering};
14use std::time::{SystemTime, UNIX_EPOCH};
15
16use crate::error::{DogeError, DogeResult};
17use crate::stdlib::int_arg;
18use crate::value::Value;
19
20const FLOAT_DIVISOR: f64 = (1u64 << 53) as f64;
23
24static SEED_COUNTER: AtomicU64 = AtomicU64::new(0);
27
28thread_local! {
29 static RNG: RefCell<Xoshiro256ss> = RefCell::new(Xoshiro256ss::from_entropy());
32}
33
34struct SplitMix64(u64);
37
38impl SplitMix64 {
39 fn next(&mut self) -> u64 {
40 self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
41 let mut z = self.0;
42 z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
43 z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
44 z ^ (z >> 31)
45 }
46}
47
48struct Xoshiro256ss {
50 s: [u64; 4],
51}
52
53impl Xoshiro256ss {
54 fn from_seed(seed: u64) -> Self {
56 let mut sm = SplitMix64(seed);
57 Xoshiro256ss {
58 s: [sm.next(), sm.next(), sm.next(), sm.next()],
59 }
60 }
61
62 fn from_entropy() -> Self {
65 let nanos = SystemTime::now()
66 .duration_since(UNIX_EPOCH)
67 .map(|d| d.as_nanos() as u64)
68 .unwrap_or(0);
69 let count = SEED_COUNTER.fetch_add(1, Ordering::Relaxed);
70 Xoshiro256ss::from_seed(nanos ^ count.wrapping_mul(0x9E37_79B9_7F4A_7C15))
71 }
72
73 fn next_u64(&mut self) -> u64 {
74 let result = self.s[1].wrapping_mul(5).rotate_left(7).wrapping_mul(9);
75 let t = self.s[1] << 17;
76 self.s[2] ^= self.s[0];
77 self.s[3] ^= self.s[1];
78 self.s[1] ^= self.s[2];
79 self.s[0] ^= self.s[3];
80 self.s[2] ^= t;
81 self.s[3] = self.s[3].rotate_left(45);
82 result
83 }
84
85 fn below(&mut self, bound: u64) -> u64 {
89 let mut x = self.next_u64();
90 let mut m = (x as u128).wrapping_mul(bound as u128);
91 let mut low = m as u64;
92 if low < bound {
93 let threshold = bound.wrapping_neg() % bound;
94 while low < threshold {
95 x = self.next_u64();
96 m = (x as u128).wrapping_mul(bound as u128);
97 low = m as u64;
98 }
99 }
100 (m >> 64) as u64
101 }
102
103 fn next_float(&mut self) -> f64 {
105 (self.next_u64() >> 11) as f64 / FLOAT_DIVISOR
106 }
107}
108
109fn with_rng<T>(f: impl FnOnce(&mut Xoshiro256ss) -> T) -> T {
111 RNG.with(|rng| f(&mut rng.borrow_mut()))
112}
113
114fn list_arg<'a>(fname: &str, v: &'a Value) -> DogeResult<&'a Rc<RefCell<Vec<Value>>>> {
117 match v {
118 Value::List(items) => Ok(items),
119 _ => Err(DogeError::type_error(format!(
120 "roll.{fname} needs a List, got {}",
121 v.describe()
122 ))),
123 }
124}
125
126pub fn roll_seed(n: &Value) -> DogeResult {
129 let seed = int_arg("roll", "seed", n)?;
130 with_rng(|rng| *rng = Xoshiro256ss::from_seed(seed as u64));
131 Ok(Value::None)
132}
133
134pub fn roll_int(low: &Value, high: &Value) -> DogeResult {
137 let low = int_arg("roll", "int", low)?;
138 let high = int_arg("roll", "int", high)?;
139 if low > high {
140 return Err(DogeError::value_error(format!(
141 "roll.int needs low <= high, but {low} > {high} — swap the bounds"
142 )));
143 }
144 let span = (high as i128 - low as i128 + 1) as u128;
145 let draw = if span > u64::MAX as u128 {
146 with_rng(|rng| rng.next_u64())
147 } else {
148 with_rng(|rng| rng.below(span as u64))
149 };
150 Ok(Value::int((low as i128 + draw as i128) as i64))
151}
152
153pub fn roll_float() -> DogeResult {
155 Ok(Value::Float(with_rng(|rng| rng.next_float())))
156}
157
158pub fn roll_choice(list: &Value) -> DogeResult {
161 let items = list_arg("choice", list)?.borrow();
162 if items.is_empty() {
163 return Err(DogeError::value_error(
164 "roll.choice needs a non-empty List — there is nothing to choose from",
165 ));
166 }
167 let index = with_rng(|rng| rng.below(items.len() as u64)) as usize;
168 Ok(items[index].clone())
169}
170
171pub fn roll_shuffle(list: &Value) -> DogeResult {
175 let mut items: Vec<Value> = list_arg("shuffle", list)?.borrow().clone();
176 let len = items.len();
177 with_rng(|rng| fisher_yates(rng, &mut items, len));
178 Ok(Value::list(items))
179}
180
181pub fn roll_sample(list: &Value, k: &Value) -> DogeResult {
185 let items: Vec<Value> = list_arg("sample", list)?.borrow().clone();
186 let k = int_arg("roll", "sample", k)?;
187 if k < 0 || k as u128 > items.len() as u128 {
188 return Err(DogeError::value_error(format!(
189 "roll.sample needs 0 <= k <= {}, got {k}",
190 items.len()
191 )));
192 }
193 let mut items = items;
194 let k = k as usize;
195 with_rng(|rng| fisher_yates(rng, &mut items, k));
196 items.truncate(k);
197 Ok(Value::list(items))
198}
199
200fn fisher_yates(rng: &mut Xoshiro256ss, items: &mut [Value], count: usize) {
204 let len = items.len();
205 for i in 0..count.min(len) {
206 let j = i + rng.below((len - i) as u64) as usize;
207 items.swap(i, j);
208 }
209}
210
211#[cfg(test)]
212mod tests {
213 use super::*;
214 use crate::error::ErrorKind;
215 use bigdecimal::ToPrimitive;
216
217 fn seed(n: i64) {
218 roll_seed(&Value::int(n)).unwrap();
219 }
220
221 fn as_int(v: Value) -> i64 {
222 match v {
223 Value::Int(n) => n.to_i64().unwrap(),
224 other => panic!("expected an Int, got {other:?}"),
225 }
226 }
227
228 fn as_list(v: Value) -> Vec<Value> {
229 match v {
230 Value::List(items) => items.borrow().clone(),
231 other => panic!("expected a List, got {other:?}"),
232 }
233 }
234
235 #[test]
236 fn same_seed_reproduces_the_sequence() {
237 seed(1234);
238 let first: Vec<i64> = (0..10)
239 .map(|_| as_int(roll_int(&Value::int(0), &Value::int(1_000_000)).unwrap()))
240 .collect();
241 seed(1234);
242 let second: Vec<i64> = (0..10)
243 .map(|_| as_int(roll_int(&Value::int(0), &Value::int(1_000_000)).unwrap()))
244 .collect();
245 assert_eq!(first, second);
246 }
247
248 #[test]
249 fn int_stays_within_the_inclusive_range() {
250 seed(7);
251 for _ in 0..1000 {
252 let n = as_int(roll_int(&Value::int(-3), &Value::int(3)).unwrap());
253 assert!((-3..=3).contains(&n), "{n} out of range");
254 }
255 }
256
257 #[test]
258 fn int_with_equal_bounds_is_that_bound() {
259 seed(7);
260 assert_eq!(as_int(roll_int(&Value::int(5), &Value::int(5)).unwrap()), 5);
261 }
262
263 #[test]
264 fn int_low_above_high_is_value_error() {
265 assert_eq!(
266 roll_int(&Value::int(5), &Value::int(1)).unwrap_err().kind,
267 ErrorKind::ValueError
268 );
269 }
270
271 #[test]
272 fn float_stays_in_the_unit_interval() {
273 seed(99);
274 for _ in 0..1000 {
275 let f = match roll_float().unwrap() {
276 Value::Float(f) => f,
277 other => panic!("expected a Float, got {other:?}"),
278 };
279 assert!((0.0..1.0).contains(&f), "{f} out of range");
280 }
281 }
282
283 #[test]
284 fn choice_returns_a_member_and_rejects_empty() {
285 seed(42);
286 let list = Value::list(vec![Value::int(10), Value::int(20), Value::int(30)]);
287 for _ in 0..100 {
288 let n = as_int(roll_choice(&list).unwrap());
289 assert!([10, 20, 30].contains(&n));
290 }
291 assert_eq!(
292 roll_choice(&Value::list(vec![])).unwrap_err().kind,
293 ErrorKind::ValueError
294 );
295 }
296
297 #[test]
298 fn shuffle_is_a_permutation_that_leaves_the_input_untouched() {
299 seed(2024);
300 let list = Value::list((0..8).map(Value::int).collect());
301 let shuffled = as_list(roll_shuffle(&list).unwrap());
302 let after: Vec<i64> = as_list(list).into_iter().map(as_int).collect();
304 assert_eq!(after, (0..8).collect::<Vec<_>>());
305 let mut got: Vec<i64> = shuffled.into_iter().map(as_int).collect();
307 got.sort_unstable();
308 assert_eq!(got, (0..8).collect::<Vec<_>>());
309 }
310
311 #[test]
312 fn sample_draws_k_distinct_positions() {
313 seed(2024);
314 let list = Value::list((0..10).map(Value::int).collect());
315 let picked = as_list(roll_sample(&list, &Value::int(4)).unwrap());
316 assert_eq!(picked.len(), 4);
317 let mut values: Vec<i64> = picked.into_iter().map(as_int).collect();
318 values.sort_unstable();
319 values.dedup();
320 assert_eq!(values.len(), 4, "sampled positions should be distinct");
321 }
322
323 #[test]
324 fn sample_bounds_are_value_errors() {
325 let list = Value::list((0..3).map(Value::int).collect());
326 assert_eq!(
327 roll_sample(&list, &Value::int(4)).unwrap_err().kind,
328 ErrorKind::ValueError
329 );
330 assert_eq!(
331 roll_sample(&list, &Value::int(-1)).unwrap_err().kind,
332 ErrorKind::ValueError
333 );
334 }
335
336 #[test]
337 fn sample_of_zero_is_empty_and_of_all_is_a_permutation() {
338 seed(1);
339 let list = Value::list((0..5).map(Value::int).collect());
340 assert!(as_list(roll_sample(&list, &Value::int(0)).unwrap()).is_empty());
341 let all = as_list(roll_sample(&list, &Value::int(5)).unwrap());
342 let mut values: Vec<i64> = all.into_iter().map(as_int).collect();
343 values.sort_unstable();
344 assert_eq!(values, (0..5).collect::<Vec<_>>());
345 }
346
347 #[test]
348 fn wrong_arg_types_are_type_errors() {
349 assert_eq!(
350 roll_seed(&Value::str("x")).unwrap_err().kind,
351 ErrorKind::TypeError
352 );
353 assert_eq!(
354 roll_int(&Value::str("x"), &Value::int(1)).unwrap_err().kind,
355 ErrorKind::TypeError
356 );
357 assert_eq!(
358 roll_choice(&Value::int(1)).unwrap_err().kind,
359 ErrorKind::TypeError
360 );
361 assert_eq!(
362 roll_shuffle(&Value::int(1)).unwrap_err().kind,
363 ErrorKind::TypeError
364 );
365 assert_eq!(
366 roll_sample(&Value::int(1), &Value::int(1))
367 .unwrap_err()
368 .kind,
369 ErrorKind::TypeError
370 );
371 }
372}