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 + rand::CryptoRng {
12    if in_antithesis() {
13        Either::Antithesis(antithesis_sdk::random::AntithesisRng)
14    } else {
15        Either::Prod(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    Prod(L),
22    Antithesis(R),
23}
24
25// We don't require that AntithesisRng is a CryptoRng, but it is only
26// used in the antithesis test environment.
27impl<Prod: rand::CryptoRng, Antithesis> rand::CryptoRng for Either<Prod, Antithesis> {}
28
29impl<L, R> rand::RngCore for Either<L, R>
30where
31    L: rand::RngCore,
32    R: rand::RngCore,
33{
34    fn next_u32(&mut self) -> u32 {
35        match self {
36            Either::Prod(l) => l.next_u32(),
37            Either::Antithesis(r) => r.next_u32(),
38        }
39    }
40
41    fn next_u64(&mut self) -> u64 {
42        match self {
43            Either::Prod(l) => l.next_u64(),
44            Either::Antithesis(r) => r.next_u64(),
45        }
46    }
47
48    fn fill_bytes(&mut self, dest: &mut [u8]) {
49        match self {
50            Either::Prod(l) => l.fill_bytes(dest),
51            Either::Antithesis(r) => r.fill_bytes(dest),
52        }
53    }
54
55    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
56        match self {
57            Either::Prod(l) => l.try_fill_bytes(dest),
58            Either::Antithesis(r) => r.try_fill_bytes(dest),
59        }
60    }
61}
mysten_common/random.rs

mysten_common/
random.rs
