mysten_common/
random.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::in_antithesis;
5
6/// Get a random number generator.
7///
8/// If we are running in antithesis mode, use the antithesis RNG.
9/// Otherwise, use the rand::thread_rng().
10#[inline(always)]
11pub fn get_rng() -> impl rand::Rng {
12    if in_antithesis() {
13        Either::Right(antithesis_sdk::random::AntithesisRng)
14    } else {
15        Either::Left(rand::thread_rng())
16    }
17}
18
19// Need our own Either implementation so we can impl rand::RngCore for it.
20pub enum Either<L, R> {
21    Left(L),
22    Right(R),
23}
24
25impl<L, R> rand::RngCore for Either<L, R>
26where
27    L: rand::RngCore,
28    R: rand::RngCore,
29{
30    fn next_u32(&mut self) -> u32 {
31        match self {
32            Either::Left(l) => l.next_u32(),
33            Either::Right(r) => r.next_u32(),
34        }
35    }
36
37    fn next_u64(&mut self) -> u64 {
38        match self {
39            Either::Left(l) => l.next_u64(),
40            Either::Right(r) => r.next_u64(),
41        }
42    }
43
44    fn fill_bytes(&mut self, dest: &mut [u8]) {
45        match self {
46            Either::Left(l) => l.fill_bytes(dest),
47            Either::Right(r) => r.fill_bytes(dest),
48        }
49    }
50
51    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
52        match self {
53            Either::Left(l) => l.try_fill_bytes(dest),
54            Either::Right(r) => r.try_fill_bytes(dest),
55        }
56    }
57}