Skip to main content

icydb_utils/
rand.rs

1use rand_chacha::{
2    ChaCha20Rng,
3    rand_core::{Rng, SeedableRng},
4};
5use std::cell::RefCell;
6use thiserror::Error;
7
8thread_local! {
9    static RNG: RefCell<Option<ChaCha20Rng>> = const { RefCell::new(None) };
10}
11
12///
13/// RngError
14///
15/// Errors raised when the shared deterministic RNG has not been seeded yet.
16///
17
18#[derive(Debug, Error)]
19pub enum RngError {
20    #[error("Randomness is not initialized. Please try again later")]
21    NotInitialized,
22}
23
24/// Seed the shared RNG from one 32-byte seed.
25pub fn seed_from(seed: [u8; 32]) {
26    RNG.with_borrow_mut(|rng| {
27        *rng = Some(ChaCha20Rng::from_seed(seed));
28    });
29}
30
31/// Return whether the shared RNG is currently seeded.
32#[must_use]
33pub fn is_seeded() -> bool {
34    RNG.with_borrow(Option::is_some)
35}
36
37fn with_rng<T>(f: impl FnOnce(&mut ChaCha20Rng) -> T) -> Result<T, RngError> {
38    RNG.with_borrow_mut(|rng| match rng.as_mut() {
39        Some(rand) => Ok(f(rand)),
40        None => Err(RngError::NotInitialized),
41    })
42}
43
44/// Fill one caller-provided buffer using the shared RNG.
45pub fn fill_bytes(dest: &mut [u8]) -> Result<(), RngError> {
46    with_rng(|rand| rand.fill_bytes(dest))
47}
48
49/// Return one owned buffer of random bytes from the shared RNG.
50pub fn random_bytes(size: usize) -> Result<Vec<u8>, RngError> {
51    let mut buf = vec![0u8; size];
52    fill_bytes(&mut buf)?;
53    Ok(buf)
54}
55
56/// Return one random `u8` from the shared RNG.
57pub fn next_u8() -> Result<u8, RngError> {
58    Ok((next_u16()? & 0xFF) as u8)
59}
60
61/// Return one random `u16` from the shared RNG.
62#[expect(clippy::cast_possible_truncation)]
63pub fn next_u16() -> Result<u16, RngError> {
64    with_rng(|rand| rand.next_u32() as u16)
65}
66
67/// Return one random `u32` from the shared RNG.
68pub fn next_u32() -> Result<u32, RngError> {
69    with_rng(Rng::next_u32)
70}
71
72/// Return one random `u64` from the shared RNG.
73pub fn next_u64() -> Result<u64, RngError> {
74    with_rng(Rng::next_u64)
75}
76
77/// Return one random `u128` from the shared RNG.
78pub fn next_u128() -> Result<u128, RngError> {
79    with_rng(|rand| {
80        let hi = u128::from(rand.next_u64());
81        let lo = u128::from(rand.next_u64());
82        (hi << 64) | lo
83    })
84}