Skip to main content

reifydb_runtime/context/rng/
mod.rs

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