runmat_runtime/builtins/common/
random.rs

1use std::collections::{HashMap, VecDeque};
2use std::f64::consts::PI;
3use std::sync::{Mutex, OnceLock};
4
5use crate::{build_runtime_error, BuiltinResult, RuntimeError};
6
7pub(crate) const DEFAULT_RNG_SEED: u64 = 0x9e3779b97f4a7c15;
8pub(crate) const DEFAULT_USER_SEED: u64 = 0;
9const RNG_MULTIPLIER: u64 = 6364136223846793005;
10const RNG_INCREMENT: u64 = 1;
11const RNG_SHIFT: u32 = 11;
12const RNG_SCALE: f64 = 1.0 / ((1u64 << 53) as f64);
13const MIN_UNIFORM: f64 = f64::MIN_POSITIVE;
14
15fn random_error(label: &str, message: impl Into<String>) -> RuntimeError {
16    build_runtime_error(message).with_builtin(label).build()
17}
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub(crate) enum RngAlgorithm {
21    RunMatLcg,
22}
23
24impl RngAlgorithm {
25    pub(crate) fn as_str(&self) -> &'static str {
26        match self {
27            RngAlgorithm::RunMatLcg => "twister",
28        }
29    }
30}
31
32#[derive(Clone, Copy, Debug)]
33pub(crate) struct RngSnapshot {
34    pub state: u64,
35    pub seed: Option<u64>,
36    pub algorithm: RngAlgorithm,
37}
38
39impl RngSnapshot {
40    pub(crate) fn new(state: u64, seed: Option<u64>, algorithm: RngAlgorithm) -> Self {
41        Self {
42            state,
43            seed,
44            algorithm,
45        }
46    }
47}
48
49#[derive(Clone, Copy)]
50struct GlobalRng {
51    state: u64,
52    seed: Option<u64>,
53    algorithm: RngAlgorithm,
54}
55
56impl GlobalRng {
57    fn new() -> Self {
58        Self {
59            state: DEFAULT_RNG_SEED,
60            seed: Some(DEFAULT_USER_SEED),
61            algorithm: RngAlgorithm::RunMatLcg,
62        }
63    }
64
65    fn snapshot(&self) -> RngSnapshot {
66        RngSnapshot {
67            state: self.state,
68            seed: self.seed,
69            algorithm: self.algorithm,
70        }
71    }
72}
73
74impl From<RngSnapshot> for GlobalRng {
75    fn from(snapshot: RngSnapshot) -> Self {
76        Self {
77            state: snapshot.state,
78            seed: snapshot.seed,
79            algorithm: snapshot.algorithm,
80        }
81    }
82}
83
84static RNG_STATE: OnceLock<Mutex<GlobalRng>> = OnceLock::new();
85const LEGACY_TOKEN_BASE: u64 = (1_u64 << 52) + 0x5eed;
86const LEGACY_TOKEN_CAPACITY: usize = 1024;
87
88struct LegacyTokenState {
89    next: u64,
90    entries: HashMap<u64, RngSnapshot>,
91    order: VecDeque<u64>,
92}
93
94impl LegacyTokenState {
95    fn new() -> Self {
96        Self {
97            next: 0,
98            entries: HashMap::new(),
99            order: VecDeque::new(),
100        }
101    }
102
103    fn insert(&mut self, snapshot: RngSnapshot) -> u64 {
104        let token = LEGACY_TOKEN_BASE + self.next;
105        self.next = self.next.wrapping_add(1);
106        if self.entries.insert(token, snapshot).is_none() {
107            self.order.push_back(token);
108        }
109        while self.order.len() > LEGACY_TOKEN_CAPACITY {
110            if let Some(oldest) = self.order.pop_front() {
111                self.entries.remove(&oldest);
112            }
113        }
114        token
115    }
116}
117
118static LEGACY_TOKENS: OnceLock<Mutex<LegacyTokenState>> = OnceLock::new();
119
120fn rng_state() -> &'static Mutex<GlobalRng> {
121    RNG_STATE.get_or_init(|| Mutex::new(GlobalRng::new()))
122}
123
124fn legacy_tokens() -> &'static Mutex<LegacyTokenState> {
125    LEGACY_TOKENS.get_or_init(|| Mutex::new(LegacyTokenState::new()))
126}
127
128fn mix_seed(seed: u64) -> u64 {
129    if seed == 0 {
130        return DEFAULT_RNG_SEED;
131    }
132    let mut z = seed.wrapping_add(0x9e3779b97f4a7c15);
133    z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
134    z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
135    let mixed = z ^ (z >> 31);
136    if mixed == 0 {
137        DEFAULT_RNG_SEED
138    } else {
139        mixed
140    }
141}
142
143pub(crate) fn snapshot() -> BuiltinResult<RngSnapshot> {
144    rng_state()
145        .lock()
146        .map(|guard| guard.snapshot())
147        .map_err(|_| random_error("rng", "rng: failed to acquire RNG lock"))
148}
149
150pub(crate) fn apply_snapshot(snapshot: RngSnapshot) -> BuiltinResult<RngSnapshot> {
151    let mut guard = rng_state()
152        .lock()
153        .map_err(|_| random_error("rng", "rng: failed to acquire RNG lock"))?;
154    let previous = guard.snapshot();
155    guard.state = snapshot.state;
156    guard.seed = snapshot.seed;
157    guard.algorithm = snapshot.algorithm;
158    Ok(previous)
159}
160
161pub(crate) fn set_seed(seed: u64) -> BuiltinResult<RngSnapshot> {
162    let state = mix_seed(seed);
163    apply_snapshot(RngSnapshot::new(state, Some(seed), RngAlgorithm::RunMatLcg))
164}
165
166pub(crate) fn legacy_seed_value() -> BuiltinResult<u64> {
167    let snapshot = snapshot()?;
168    if let Some(seed) = snapshot.seed {
169        if snapshot.state == mix_seed(seed) {
170            return Ok(seed);
171        }
172    }
173    let mut tokens = legacy_tokens()
174        .lock()
175        .map_err(|_| random_error("rand", "rand: failed to acquire legacy RNG token lock"))?;
176    Ok(tokens.insert(snapshot))
177}
178
179pub(crate) fn set_legacy_seed(seed_or_token: u64) -> BuiltinResult<RngSnapshot> {
180    if let Some(snapshot) = legacy_tokens()
181        .lock()
182        .map_err(|_| random_error("rand", "rand: failed to acquire legacy RNG token lock"))?
183        .entries
184        .get(&seed_or_token)
185        .copied()
186    {
187        apply_snapshot(snapshot)
188    } else {
189        set_seed(seed_or_token)
190    }
191}
192
193pub(crate) fn set_default() -> BuiltinResult<RngSnapshot> {
194    apply_snapshot(default_snapshot())
195}
196
197pub(crate) fn default_snapshot() -> RngSnapshot {
198    RngSnapshot::new(
199        DEFAULT_RNG_SEED,
200        Some(DEFAULT_USER_SEED),
201        RngAlgorithm::RunMatLcg,
202    )
203}
204
205pub(crate) fn generate_uniform(len: usize, label: &str) -> BuiltinResult<Vec<f64>> {
206    let mut guard = rng_state()
207        .lock()
208        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
209    let mut out = Vec::with_capacity(len);
210    for _ in 0..len {
211        out.push(next_uniform_state(&mut guard.state));
212    }
213    Ok(out)
214}
215
216pub(crate) fn generate_uniform_single(len: usize, label: &str) -> BuiltinResult<Vec<f64>> {
217    generate_uniform(len, label).map(|data| {
218        data.into_iter()
219            .map(|v| {
220                let value = v as f32;
221                value as f64
222            })
223            .collect()
224    })
225}
226
227pub(crate) fn skip_uniform(len: usize, label: &str) -> BuiltinResult<()> {
228    if len == 0 {
229        return Ok(());
230    }
231    let mut guard = rng_state()
232        .lock()
233        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
234    guard.state = advance_state(guard.state, len as u64);
235    Ok(())
236}
237
238fn advance_state(state: u64, mut delta: u64) -> u64 {
239    if delta == 0 {
240        return state;
241    }
242    let mut cur_mult = RNG_MULTIPLIER;
243    let mut cur_plus = RNG_INCREMENT;
244    let mut acc_mult = 1u64;
245    let mut acc_plus = 0u64;
246    while delta > 0 {
247        if (delta & 1) != 0 {
248            acc_mult = acc_mult.wrapping_mul(cur_mult);
249            acc_plus = acc_plus.wrapping_mul(cur_mult).wrapping_add(cur_plus);
250        }
251        cur_plus = cur_plus.wrapping_mul(cur_mult.wrapping_add(1));
252        cur_mult = cur_mult.wrapping_mul(cur_mult);
253        delta >>= 1;
254    }
255    acc_mult.wrapping_mul(state).wrapping_add(acc_plus)
256}
257
258pub(crate) fn generate_complex(len: usize, label: &str) -> BuiltinResult<Vec<(f64, f64)>> {
259    let mut guard = rng_state()
260        .lock()
261        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
262    let mut out = Vec::with_capacity(len);
263    for _ in 0..len {
264        let re = next_uniform_state(&mut guard.state);
265        let im = next_uniform_state(&mut guard.state);
266        out.push((re, im));
267    }
268    Ok(out)
269}
270
271pub(crate) fn next_uniform_state(state: &mut u64) -> f64 {
272    *state = state
273        .wrapping_mul(RNG_MULTIPLIER)
274        .wrapping_add(RNG_INCREMENT);
275    let bits = *state >> RNG_SHIFT;
276    (bits as f64) * RNG_SCALE
277}
278
279fn next_normal_pair(state: &mut u64) -> (f64, f64) {
280    let mut u1 = next_uniform_state(state);
281    if u1 <= 0.0 {
282        u1 = MIN_UNIFORM;
283    }
284    let u2 = next_uniform_state(state);
285    let radius = (-2.0 * u1.ln()).sqrt();
286    let angle = 2.0 * PI * u2;
287    (radius * angle.cos(), radius * angle.sin())
288}
289
290pub(crate) fn generate_exponential(mu: f64, len: usize, label: &str) -> BuiltinResult<Vec<f64>> {
291    let mut guard = rng_state()
292        .lock()
293        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
294    let mut out = Vec::with_capacity(len);
295    for _ in 0..len {
296        let u = next_uniform_state(&mut guard.state).max(MIN_UNIFORM);
297        out.push(-mu * u.ln());
298    }
299    Ok(out)
300}
301
302pub(crate) fn generate_normal_scaled(
303    mu: f64,
304    sigma: f64,
305    len: usize,
306    label: &str,
307) -> BuiltinResult<Vec<f64>> {
308    let mut guard = rng_state()
309        .lock()
310        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
311    let mut out = Vec::with_capacity(len);
312    while out.len() < len {
313        let (z0, z1) = next_normal_pair(&mut guard.state);
314        out.push(mu + sigma * z0);
315        if out.len() < len {
316            out.push(mu + sigma * z1);
317        }
318    }
319    Ok(out)
320}
321
322pub(crate) fn generate_student_t(nu: &[f64], len: usize, label: &str) -> BuiltinResult<Vec<f64>> {
323    let mut guard = rng_state()
324        .lock()
325        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
326    let mut out = Vec::with_capacity(len);
327    let mut spare_normal = None;
328    for idx in 0..len {
329        let df = if nu.len() == 1 { nu[0] } else { nu[idx] };
330        let z = if let Some(value) = spare_normal.take() {
331            value
332        } else {
333            let (z0, z1) = next_normal_pair(&mut guard.state);
334            spare_normal = Some(z1);
335            z0
336        };
337        if df == f64::INFINITY {
338            out.push(z);
339        } else {
340            let chi_square = sample_gamma_shape_scale(&mut guard.state, df / 2.0, 2.0);
341            out.push(z / (chi_square / df).sqrt());
342        }
343    }
344    Ok(out)
345}
346
347pub(crate) fn generate_gamma_shape_scale(
348    shape: &[f64],
349    scale: &[f64],
350    len: usize,
351    label: &str,
352) -> BuiltinResult<Vec<f64>> {
353    let mut guard = rng_state()
354        .lock()
355        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
356    let mut out = Vec::with_capacity(len);
357    for idx in 0..len {
358        let a = if shape.len() == 1 {
359            shape[0]
360        } else {
361            shape[idx]
362        };
363        let b = if scale.len() == 1 {
364            scale[0]
365        } else {
366            scale[idx]
367        };
368        out.push(if a == 0.0 {
369            0.0
370        } else {
371            sample_gamma_shape_scale(&mut guard.state, a, b)
372        });
373    }
374    Ok(out)
375}
376
377pub(crate) fn generate_binomial(
378    trials: &[f64],
379    probability: &[f64],
380    len: usize,
381    label: &str,
382) -> BuiltinResult<Vec<f64>> {
383    let mut guard = rng_state()
384        .lock()
385        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
386    let mut out = Vec::with_capacity(len);
387    for idx in 0..len {
388        let n = if trials.len() == 1 {
389            trials[0]
390        } else {
391            trials[idx]
392        };
393        let p = if probability.len() == 1 {
394            probability[0]
395        } else {
396            probability[idx]
397        };
398        out.push(sample_binomial(&mut guard.state, n, p));
399    }
400    Ok(out)
401}
402
403pub(crate) fn generate_weibull(
404    scale: &[f64],
405    shape: &[f64],
406    len: usize,
407    label: &str,
408) -> BuiltinResult<Vec<f64>> {
409    let mut guard = rng_state()
410        .lock()
411        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
412    let mut out = Vec::with_capacity(len);
413    for idx in 0..len {
414        let a = if scale.len() == 1 {
415            scale[0]
416        } else {
417            scale[idx]
418        };
419        let b = if shape.len() == 1 {
420            shape[0]
421        } else {
422            shape[idx]
423        };
424        let u = next_uniform_state(&mut guard.state).max(MIN_UNIFORM);
425        out.push(a * (-u.ln()).powf(1.0 / b));
426    }
427    Ok(out)
428}
429
430fn sample_gamma_shape_scale(state: &mut u64, shape: f64, scale: f64) -> f64 {
431    if shape < 1.0 {
432        let u = next_uniform_state(state).max(MIN_UNIFORM);
433        return sample_gamma_shape_scale(state, shape + 1.0, scale) * u.powf(1.0 / shape);
434    }
435
436    let d = shape - 1.0 / 3.0;
437    let c = (1.0 / (9.0 * d)).sqrt();
438    loop {
439        let (x, _) = next_normal_pair(state);
440        let v_base = 1.0 + c * x;
441        if v_base <= 0.0 {
442            continue;
443        }
444        let v = v_base * v_base * v_base;
445        let u = next_uniform_state(state);
446        let x2 = x * x;
447        if u < 1.0 - 0.0331 * x2 * x2 || u.ln() < 0.5 * x2 + d * (1.0 - v + v.ln()) {
448            return scale * d * v;
449        }
450    }
451}
452
453fn sample_binomial(state: &mut u64, trials: f64, probability: f64) -> f64 {
454    if trials == 0.0 || probability == 0.0 {
455        return 0.0;
456    }
457    if probability == 1.0 {
458        return trials;
459    }
460
461    let effective_p = probability.min(1.0 - probability);
462    let mean = trials * effective_p;
463    let sampled = if trials <= 10_000.0 {
464        sample_binomial_direct(state, trials as u64, effective_p)
465    } else if mean < 30.0 {
466        sample_poisson_knuth(state, mean).min(trials)
467    } else {
468        sample_binomial_normal_approx(state, trials, effective_p)
469    };
470
471    if probability <= 0.5 {
472        sampled
473    } else {
474        trials - sampled
475    }
476}
477
478fn sample_binomial_direct(state: &mut u64, trials: u64, probability: f64) -> f64 {
479    let mut successes = 0_u64;
480    for _ in 0..trials {
481        if next_uniform_state(state) < probability {
482            successes += 1;
483        }
484    }
485    successes as f64
486}
487
488fn sample_poisson_knuth(state: &mut u64, lambda: f64) -> f64 {
489    if lambda <= 0.0 {
490        return 0.0;
491    }
492    let threshold = (-lambda).exp();
493    let mut product = 1.0;
494    let mut count = 0_u64;
495    loop {
496        count += 1;
497        product *= next_uniform_state(state).max(MIN_UNIFORM);
498        if product <= threshold {
499            return (count - 1) as f64;
500        }
501    }
502}
503
504fn sample_binomial_normal_approx(state: &mut u64, trials: f64, probability: f64) -> f64 {
505    let mean = trials * probability;
506    let variance = mean * (1.0 - probability);
507    if variance <= 0.0 {
508        return mean.round().clamp(0.0, trials);
509    }
510    let (z, _) = next_normal_pair(state);
511    (mean + variance.sqrt() * z).round().clamp(0.0, trials)
512}
513
514pub(crate) fn generate_uniform_scaled(
515    a: f64,
516    b: f64,
517    len: usize,
518    label: &str,
519) -> BuiltinResult<Vec<f64>> {
520    let mut guard = rng_state()
521        .lock()
522        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
523    let mut out = Vec::with_capacity(len);
524    for _ in 0..len {
525        out.push(a + (b - a) * next_uniform_state(&mut guard.state));
526    }
527    Ok(out)
528}
529
530pub(crate) fn generate_normal(len: usize, label: &str) -> BuiltinResult<Vec<f64>> {
531    let mut guard = rng_state()
532        .lock()
533        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
534    let mut out = Vec::with_capacity(len);
535    while out.len() < len {
536        let (z0, z1) = next_normal_pair(&mut guard.state);
537        out.push(z0);
538        if out.len() < len {
539            out.push(z1);
540        }
541    }
542    Ok(out)
543}
544
545pub(crate) fn generate_normal_complex(len: usize, label: &str) -> BuiltinResult<Vec<(f64, f64)>> {
546    let mut guard = rng_state()
547        .lock()
548        .map_err(|_| random_error(label, format!("{label}: failed to acquire RNG lock")))?;
549    let mut out = Vec::with_capacity(len);
550    for _ in 0..len {
551        let (re, im) = next_normal_pair(&mut guard.state);
552        out.push((re, im));
553    }
554    Ok(out)
555}
556
557#[cfg(test)]
558pub(crate) fn reset_rng() {
559    if let Some(mutex) = RNG_STATE.get() {
560        if let Ok(mut guard) = mutex.lock() {
561            *guard = GlobalRng::from(default_snapshot());
562        }
563    } else {
564        let _ = RNG_STATE.set(Mutex::new(GlobalRng::new()));
565    }
566    if let Some(mutex) = LEGACY_TOKENS.get() {
567        if let Ok(mut guard) = mutex.lock() {
568            *guard = LegacyTokenState::new();
569        }
570    }
571}
572
573#[cfg(test)]
574pub(crate) fn expected_exponential_sequence(mu: f64, count: usize) -> Vec<f64> {
575    let mut seed = DEFAULT_RNG_SEED;
576    let mut seq = Vec::with_capacity(count);
577    for _ in 0..count {
578        let u = next_uniform_state(&mut seed).max(MIN_UNIFORM);
579        seq.push(-mu * u.ln());
580    }
581    seq
582}
583
584#[cfg(test)]
585pub(crate) fn expected_normal_scaled_sequence(mu: f64, sigma: f64, count: usize) -> Vec<f64> {
586    let mut seed = DEFAULT_RNG_SEED;
587    let mut seq = Vec::with_capacity(count);
588    while seq.len() < count {
589        let (z0, z1) = next_normal_pair(&mut seed);
590        seq.push(mu + sigma * z0);
591        if seq.len() < count {
592            seq.push(mu + sigma * z1);
593        }
594    }
595    seq
596}
597
598#[cfg(test)]
599pub(crate) fn expected_uniform_scaled_sequence(a: f64, b: f64, count: usize) -> Vec<f64> {
600    let mut seed = DEFAULT_RNG_SEED;
601    let mut seq = Vec::with_capacity(count);
602    for _ in 0..count {
603        seq.push(a + (b - a) * next_uniform_state(&mut seed));
604    }
605    seq
606}
607
608#[cfg(test)]
609pub(crate) fn expected_uniform_sequence(count: usize) -> Vec<f64> {
610    let mut seed = DEFAULT_RNG_SEED;
611    let mut seq = Vec::with_capacity(count);
612    for _ in 0..count {
613        seq.push(next_uniform_state(&mut seed));
614    }
615    seq
616}
617
618#[cfg(test)]
619pub(crate) fn expected_complex_sequence(count: usize) -> Vec<(f64, f64)> {
620    let mut seed = DEFAULT_RNG_SEED;
621    let mut seq = Vec::with_capacity(count);
622    for _ in 0..count {
623        let re = next_uniform_state(&mut seed);
624        let im = next_uniform_state(&mut seed);
625        seq.push((re, im));
626    }
627    seq
628}
629
630#[cfg(test)]
631pub(crate) fn expected_normal_sequence(count: usize) -> Vec<f64> {
632    let mut seed = DEFAULT_RNG_SEED;
633    let mut seq = Vec::with_capacity(count);
634    while seq.len() < count {
635        let (z0, z1) = next_normal_pair(&mut seed);
636        seq.push(z0);
637        if seq.len() < count {
638            seq.push(z1);
639        }
640    }
641    seq
642}
643
644#[cfg(test)]
645pub(crate) fn expected_complex_normal_sequence(count: usize) -> Vec<(f64, f64)> {
646    let mut seed = DEFAULT_RNG_SEED;
647    let mut seq = Vec::with_capacity(count);
648    for _ in 0..count {
649        seq.push(next_normal_pair(&mut seed));
650    }
651    seq
652}
653
654#[cfg(test)]
655pub(crate) fn test_guard() -> super::test_support::GlobalStateTestGuard {
656    super::test_support::global_state_test_guard()
657}
runmat_runtime/builtins/common/random.rs

runmat_runtime/builtins/common/
random.rs
