Skip to main content

reifydb_runtime/context/rng/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::sync::{Arc, Mutex};
5
6use getrandom::fill as getrandom_fill;
7use rand::{Rng as RandRng, RngExt, SeedableRng, rngs::StdRng};
8
9#[derive(Clone, Default)]
10pub enum Rng {
11	#[default]
12	Os,
13
14	Seeded(SeededRng),
15}
16
17impl Rng {
18	pub fn seeded(seed: u64) -> Self {
19		Rng::Seeded(SeededRng::new(seed))
20	}
21
22	pub fn bytes_16(&self) -> [u8; 16] {
23		match self {
24			Rng::Os => {
25				let mut buf = [0u8; 16];
26				getrandom_fill(&mut buf).expect("getrandom failed");
27				buf
28			}
29			Rng::Seeded(seeded) => {
30				let mut buf = [0u8; 16];
31				let mut rng = seeded.inner.lock().unwrap();
32				rng.fill_bytes(&mut buf);
33				buf
34			}
35		}
36	}
37
38	pub fn bytes_32(&self) -> [u8; 32] {
39		match self {
40			Rng::Os => {
41				let mut buf = [0u8; 32];
42				getrandom_fill(&mut buf).expect("getrandom failed");
43				buf
44			}
45			Rng::Seeded(seeded) => {
46				let mut buf = [0u8; 32];
47				let mut rng = seeded.inner.lock().unwrap();
48				rng.fill_bytes(&mut buf);
49				buf
50			}
51		}
52	}
53
54	pub fn bytes_10(&self) -> [u8; 10] {
55		match self {
56			Rng::Os => {
57				let mut buf = [0u8; 10];
58				getrandom_fill(&mut buf).expect("getrandom failed");
59				buf
60			}
61			Rng::Seeded(seeded) => {
62				let mut buf = [0u8; 10];
63				let mut rng = seeded.inner.lock().unwrap();
64				rng.fill_bytes(&mut buf);
65				buf
66			}
67		}
68	}
69
70	pub fn infra_bytes_10(&self) -> [u8; 10] {
71		match self {
72			Rng::Os => {
73				let mut buf = [0u8; 10];
74				getrandom_fill(&mut buf).expect("getrandom failed");
75				buf
76			}
77			Rng::Seeded(seeded) => {
78				let mut buf = [0u8; 10];
79				let mut rng = seeded.infra.lock().unwrap();
80				rng.fill_bytes(&mut buf);
81				buf
82			}
83		}
84	}
85
86	pub fn infra_bytes_32(&self) -> [u8; 32] {
87		match self {
88			Rng::Os => {
89				let mut buf = [0u8; 32];
90				getrandom_fill(&mut buf).expect("getrandom failed");
91				buf
92			}
93			Rng::Seeded(seeded) => {
94				let mut buf = [0u8; 32];
95				let mut rng = seeded.infra.lock().unwrap();
96				rng.fill_bytes(&mut buf);
97				buf
98			}
99		}
100	}
101
102	pub fn infra_u64_inclusive(&self, max_inclusive: u64) -> u64 {
103		if max_inclusive == 0 {
104			return 0;
105		}
106		match self {
107			Rng::Os => {
108				let mut buf = [0u8; 8];
109				getrandom_fill(&mut buf).expect("getrandom failed");
110				let raw = u64::from_le_bytes(buf);
111				if max_inclusive == u64::MAX {
112					raw
113				} else {
114					raw % (max_inclusive + 1)
115				}
116			}
117			Rng::Seeded(seeded) => {
118				let mut rng = seeded.infra.lock().unwrap();
119				rng.random_range(0..=max_inclusive)
120			}
121		}
122	}
123}
124
125#[derive(Clone)]
126pub struct SeededRng {
127	inner: Arc<Mutex<StdRng>>,
128
129	infra: Arc<Mutex<StdRng>>,
130}
131
132impl SeededRng {
133	pub fn new(seed: u64) -> Self {
134		Self {
135			inner: Arc::new(Mutex::new(StdRng::seed_from_u64(seed))),
136			infra: Arc::new(Mutex::new(StdRng::seed_from_u64(seed ^ 0x5A5A5A5A5A5A5A5A))),
137		}
138	}
139}