From 826148af5693fe284fa307f90d496fc6dd8c956a Mon Sep 17 00:00:00 2001 From: Andrew Gazelka Date: Wed, 30 Oct 2024 09:53:21 -0700 Subject: [PATCH 1/2] feat: implement `rand::RngCore` for `AntithesisRng` Adds `AntithesisRng` struct implementing the `rand` crate's `RngCore` trait, allowing integration with the broader Rust random number generation ecosystem. The implementation uses the existing `get_random()` function and adds methods like `next_u32()`, `next_u64()`, and `fill_bytes()`. Includes comprehensive tests for the new functionality. Key changes: - Add `use rand::{Error, RngCore}` - Create `AntithesisRng` struct with `RngCore` implementation - Add tests using `rand::Rng` and `SliceRandom` traits --- lib/src/random.rs | 114 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) diff --git a/lib/src/random.rs b/lib/src/random.rs index bb087c4..d2567b5 100644 --- a/lib/src/random.rs +++ b/lib/src/random.rs @@ -1,3 +1,4 @@ +use rand::{Error, RngCore}; use crate::internal; /// Returns a u64 value chosen by Antithesis. You should not @@ -45,10 +46,71 @@ pub fn random_choice(slice: &[T]) -> Option<&T> { } } +/// A random number generator that uses Antithesis's random number generation. +/// +/// This implements the `RngCore` trait from the `rand` crate, allowing it to be used +/// with any code that expects a random number generator from that ecosystem. +/// +/// # Example +/// +/// ``` +/// use antithesis_sdk::random::AntithesisRng; +/// use rand::RngCore; +/// +/// let mut rng = AntithesisRng; +/// let random_u32 = rng.next_u32(); +/// let random_u64 = rng.next_u64(); +/// +/// let mut bytes = [0u8; 16]; +/// rng.fill_bytes(&mut bytes); +/// ``` +pub struct AntithesisRng; + +impl RngCore for AntithesisRng { + fn next_u32(&mut self) -> u32 { + get_random() as u32 + } + + fn next_u64(&mut self) -> u64 { + get_random() + } + + fn fill_bytes(&mut self, dest: &mut [u8]) { + // Split the destination buffer into chunks of 8 bytes each + // (since we'll fill each chunk with a u64/8 bytes of random data) + let mut chunks = dest.chunks_exact_mut(8); + + // Fill each complete 8-byte chunk with random bytes + for chunk in chunks.by_ref() { + // Generate 8 random bytes from a u64 in native endian order + let random_bytes = self.next_u64().to_ne_bytes(); + // Copy those random bytes into this chunk + chunk.copy_from_slice(&random_bytes); + } + + // Get any remaining bytes that didn't fit in a complete 8-byte chunk + let remainder = chunks.into_remainder(); + + if !remainder.is_empty() { + // Generate 8 more random bytes + let random_bytes = self.next_u64().to_ne_bytes(); + // Copy just enough random bytes to fill the remainder + remainder.copy_from_slice(&random_bytes[..remainder.len()]); + } + } + + fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> { + self.fill_bytes(dest); + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; use std::collections::{HashMap, HashSet}; + use rand::Rng; + use rand::seq::SliceRandom; #[test] fn random_choice_no_choices() { @@ -97,4 +159,56 @@ mod tests { random_numbers.insert(rn); } } + + #[test] + fn rng_no_choices() { + let mut rng = AntithesisRng; + let array = [""; 0]; + assert_eq!(0, array.len()); + assert_eq!(None, array.choose(&mut rng)); + } + + #[test] + fn rng_one_choice() { + let mut rng = AntithesisRng; + let array = ["ABc"; 1]; + assert_eq!(1, array.len()); + assert_eq!(Some(&"ABc"), array.choose(&mut rng)); + } + + #[test] + fn rng_few_choices() { + let mut rng = AntithesisRng; + // For each map key, the value is the count of the number of + // random_choice responses received matching that key + let mut counted_items: HashMap<&str, i64> = HashMap::new(); + counted_items.insert("a", 0); + counted_items.insert("b", 0); + counted_items.insert("c", 0); + + let all_keys: Vec<&str> = counted_items.keys().cloned().collect(); + assert_eq!(counted_items.len(), all_keys.len()); + for _i in 0..15 { + let rc = all_keys.choose(&mut rng); + if let Some(choice) = rc { + if let Some(x) = counted_items.get_mut(choice) { + *x += 1; + } + } + } + for (key, val) in counted_items.iter() { + assert_ne!(*val, 0, "Did not produce the choice: {}", key); + } + } + + #[test] + fn rng_100k() { + let mut rng = AntithesisRng; + let mut random_numbers: HashSet = HashSet::new(); + for _i in 0..100000 { + let rn: u64 = rng.gen(); + assert!(!random_numbers.contains(&rn)); + random_numbers.insert(rn); + } + } } From 50cc08bcf8cd056fff073ab140c9f5e182082219 Mon Sep 17 00:00:00 2001 From: Andrew Gazelka Date: Wed, 30 Oct 2024 10:18:56 -0700 Subject: [PATCH 2/2] add more examples in doc --- lib/src/random.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/lib/src/random.rs b/lib/src/random.rs index d2567b5..cd8186f 100644 --- a/lib/src/random.rs +++ b/lib/src/random.rs @@ -55,11 +55,12 @@ pub fn random_choice(slice: &[T]) -> Option<&T> { /// /// ``` /// use antithesis_sdk::random::AntithesisRng; -/// use rand::RngCore; +/// use rand::{Rng, RngCore}; /// /// let mut rng = AntithesisRng; -/// let random_u32 = rng.next_u32(); -/// let random_u64 = rng.next_u64(); +/// let random_u32: u32 = rng.gen(); +/// let random_u64: u64 = rng.gen(); +/// let random_char: char = rng.gen(); /// /// let mut bytes = [0u8; 16]; /// rng.fill_bytes(&mut bytes);