Skip to content

Commit

Permalink
Make test_rng randomized by default in std
Browse files Browse the repository at this point in the history
  • Loading branch information
Pratyush committed Jul 26, 2021
1 parent 7c752c5 commit 7850e7e
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 8 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

### Breaking changes

- [\#32](https://github.com/arkworks-rs/utils/pull/32) Change `test_rng` to return `impl Rng`, and make the output randomized by default when the `std` feature is set. Introduces a `DETERMINISTIC_TEST_RNG` environment variable that forces the old deterministic behavior when `DETERMINISTIC_TEST_RNG=1` is set.

### Features

### Improvements
Expand Down
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ num-traits = { version = "0.2", default-features = false }

[features]
default = [ "std" ]
std = []
std = [ "rand/std" ]
parallel = [ "rayon", "std" ]
print-trace = [ "std", "colored" ]

Expand Down
94 changes: 87 additions & 7 deletions src/rand_helper.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,4 @@
use rand::{
distributions::{Distribution, Standard},
rngs::StdRng,
Rng,
};
use rand::{Rng, RngCore, distributions::{Distribution, Standard}, prelude::{StdRng, ThreadRng}};

pub use rand;

Expand All @@ -20,13 +16,97 @@ where
}
}

/// Should be used only for tests, not for any real world usage.
pub fn test_rng() -> StdRng {
fn test_rng_helper() -> StdRng {
use rand::SeedableRng;
// arbitrary seed
let seed = [
1, 0, 0, 0, 23, 0, 0, 0, 200, 1, 0, 0, 210, 30, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0,
];
rand::rngs::StdRng::from_seed(seed)

}

/// Should be used only for tests, not for any real world usage.
#[cfg(not(feature = "std"))]
pub fn test_rng() -> impl rand::Rng {
test_rng_helper()
}

/// Should be used only for tests, not for any real world usage.
#[cfg(feature = "std")]
pub fn test_rng() -> impl rand::Rng {
let is_deterministic = std::env::vars().any(|(var, value)| var == "DETERMINISTIC_TEST_RNG" && value == "1");
if is_deterministic {
RngWrapper::Deterministic(test_rng_helper())
} else {
RngWrapper::Randomized(rand::thread_rng())
}
}

/// Helper wrapper to enable `test_rng` to return `impl::Rng`.
#[cfg(feature = "std")]
enum RngWrapper {
Deterministic(StdRng),
Randomized(ThreadRng),
}

impl RngCore for RngWrapper {
#[inline(always)]
fn next_u32(&mut self) -> u32 {
match self {
Self::Deterministic(rng) => rng.next_u32(),
Self::Randomized(rng) => rng.next_u32(),
}
}

#[inline(always)]
fn next_u64(&mut self) -> u64 {
match self {
Self::Deterministic(rng) => rng.next_u64(),
Self::Randomized(rng) => rng.next_u64(),
}
}

#[inline(always)]
fn fill_bytes(&mut self, dest: &mut [u8]) {
match self {
Self::Deterministic(rng) => rng.fill_bytes(dest),
Self::Randomized(rng) => rng.fill_bytes(dest),
}
}

#[inline(always)]
fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
match self {
Self::Deterministic(rng) => rng.try_fill_bytes(dest),
Self::Randomized(rng) => rng.try_fill_bytes(dest),
}
}
}

#[cfg(all(test, feature = "std"))]
mod test {
#[test]
fn test_deterministic_rng() {
use super::*;

let mut rng = super::test_rng();
let a = u128::rand(&mut rng);

// Reset the rng by sampling a new one.
let mut rng = super::test_rng();
let b = u128::rand(&mut rng);
assert_ne!(a, b); // should be unequal with high probability.

// Let's make the rng deterministic.
std::env::set_var("DETERMINISTIC_TEST_RNG", "1");
let mut rng = super::test_rng();
let a = u128::rand(&mut rng);

// Reset the rng by sampling a new one.
let mut rng = super::test_rng();
let b = u128::rand(&mut rng);
assert_eq!(a, b); // should be unequal with high probability.
}
}

0 comments on commit 7850e7e

Please sign in to comment.