feat(bolero-generator): add support for calling generators from anywhere (#258)

Author: camshaft

lib/bolero-engine/src/lib.rs

@@ -4,6 +4,9 @@ pub use bolero_generator::{
     TypeGenerator, ValueGenerator,
 };
 
+#[cfg(kani)]
+pub use bolero_generator::kani;
+
 pub type Seed = u128;
 
 #[cfg(not(kani))]

lib/bolero-generator/Cargo.toml

@@ -12,17 +12,21 @@ readme = "README.md"
 rust-version = "1.66.0"
 
 [features]
-default = ["either", "std"]
+default = ["any", "either", "std"]
+any = ["getrandom", "rand_xoshiro", "std"]
 std = ["alloc", "either/use_std"]
 alloc = ["rand_core/alloc"]
 
 [dependencies]
 arbitrary = { version = "1.0", optional = true }
 bolero-generator-derive = { version = "0.11", path = "../bolero-generator-derive" }
 either = { version = "1.5", default-features = false, optional = true }
+getrandom = { version = "0.2", optional = true }
 rand_core = { version = "0.6", default-features = false }
+rand_xoshiro = { version = "0.6", optional = true }
 
 [dev-dependencies]
+insta = "1"
 rand = "0.8"
 
 [lints.rust.unexpected_cfgs]

lib/bolero-generator/src/any.rs

@@ -0,0 +1,116 @@
+use crate::{gen, TypeGenerator, ValueGenerator};
+
+#[cfg(not(kani))]
+mod default;
+
+#[cfg(any(kani, test))]
+#[cfg_attr(not(kani), allow(dead_code, unused_imports))]
+mod kani;
+
+#[cfg(test)]
+mod tests;
+
+pub mod scope {
+    #[cfg(not(kani))]
+    pub use super::default::*;
+    #[cfg(kani)]
+    pub use super::kani::*;
+}
+
+pub use scope::{assume, fill_bytes, Error};
+
+pub trait Any: ValueGenerator {
+    fn any(&self) -> Self::Output;
+}
+
+impl<G: 'static + ValueGenerator> Any for G {
+    #[track_caller]
+    fn any(&self) -> Self::Output {
+        scope::any(self)
+    }
+}
+
+#[inline]
+pub fn any<T: TypeGenerator>() -> T {
+    gen().any()
+}
+
+pub trait AnySliceExt<T> {
+    fn pick(&self) -> &T;
+}
+
+impl<T> AnySliceExt<T> for [T] {
+    #[inline]
+    fn pick(&self) -> &T {
+        let index = (0..self.len()).any();
+        &self[index]
+    }
+}
+
+pub trait AnySliceMutExt<T> {
+    fn shuffle(&mut self);
+    fn fill_any(&mut self)
+    where
+        T: TypeGenerator;
+}
+
+impl<T> AnySliceMutExt<T> for [T] {
+    #[inline]
+    fn shuffle(&mut self) {
+        let max_dst = self.len().saturating_sub(1);
+        for src in 0..max_dst {
+            let dst = (src..=max_dst).any();
+            self.swap(src, dst);
+        }
+    }
+
+    #[inline]
+    fn fill_any(&mut self)
+    where
+        T: TypeGenerator,
+    {
+        for value in self {
+            *value = any();
+        }
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<T> AnySliceMutExt<T> for alloc::collections::VecDeque<T> {
+    #[inline]
+    fn shuffle(&mut self) {
+        let max_dst = self.len().saturating_sub(1);
+        for src in 0..max_dst {
+            let dst = (src..=max_dst).any();
+            self.swap(src, dst);
+        }
+    }
+
+    #[inline]
+    fn fill_any(&mut self)
+    where
+        T: TypeGenerator,
+    {
+        for value in self {
+            *value = any();
+        }
+    }
+}
+
+#[inline]
+pub fn fill<T>(values: &mut [T])
+where
+    T: TypeGenerator,
+{
+    values.fill_any()
+}
+
+#[inline]
+pub fn shuffle<T>(items: &mut [T]) {
+    items.shuffle()
+}
+
+#[inline]
+pub fn pick<T>(items: &[T]) -> &T {
+    items.pick()
+}

lib/bolero-generator/src/any/default.rs

@@ -0,0 +1,162 @@
+use crate::driver::object::{self, DynDriver, Object};
+use core::fmt;
+use std::cell::RefCell;
+
+pub trait Scope: 'static + DynDriver + core::any::Any {
+    fn borrowed(&mut self) -> object::Borrowed;
+}
+
+impl<T> Scope for T
+where
+    T: 'static + DynDriver + core::any::Any,
+{
+    fn borrowed(&mut self) -> object::Borrowed {
+        object::Borrowed(self)
+    }
+}
+
+type Type = Box<dyn Scope>;
+
+thread_local! {
+    static SCOPE: RefCell<Type> = RefCell::new(Box::new(Object(default())));
+}
+
+fn default() -> impl crate::Driver {
+    use rand_core::SeedableRng;
+    use rand_xoshiro::Xoshiro128PlusPlus;
+
+    let mut seed = [42; 16];
+    // make a best effort to get random seeds
+    let _ = getrandom::getrandom(&mut seed);
+    let rng = Xoshiro128PlusPlus::from_seed(seed);
+    crate::driver::Rng::new(rng, &Default::default())
+}
+
+fn set(value: Type) -> Type {
+    SCOPE.with(|r| core::mem::replace(&mut *r.borrow_mut(), value))
+}
+
+// protect against panics in the `with` function
+struct Prev(Option<Type>);
+
+impl Prev {
+    fn reset(mut self) -> Type {
+        set(self.0.take().unwrap())
+    }
+}
+
+impl Drop for Prev {
+    fn drop(&mut self) {
+        if let Some(prev) = self.0.take() {
+            let _ = set(prev);
+        }
+    }
+}
+
+pub fn with<D, F, R>(driver: Box<D>, f: F) -> (Box<D>, R)
+where
+    D: Scope,
+    F: FnOnce() -> R,
+{
+    let prev = Prev(Some(set(driver)));
+    let res = f();
+    let driver = prev.reset();
+    let driver = if driver.type_id() == core::any::TypeId::of::<D>() {
+        unsafe {
+            let raw = Box::into_raw(driver);
+            Box::from_raw(raw as *mut D)
+        }
+    } else {
+        panic!(
+            "invalid scope state; expected {}",
+            core::any::type_name::<D>()
+        )
+    };
+    (driver, res)
+}
+
+fn borrow_with<F: FnOnce(&mut object::Borrowed) -> R, R>(f: F) -> R {
+    SCOPE.with(|r| {
+        let mut driver = r.borrow_mut();
+        let mut driver = driver.borrowed();
+        f(&mut driver)
+    })
+}
+
+#[track_caller]
+pub fn any<G: crate::ValueGenerator>(g: &G) -> G::Output {
+    borrow_with(|driver| {
+        g.generate(driver).unwrap_or_else(|| {
+            std::panic::panic_any(Error {
+                location: core::panic::Location::caller(),
+                generator: core::any::type_name::<G>(),
+                output: core::any::type_name::<G::Output>(),
+            })
+        })
+    })
+}
+
+#[track_caller]
+pub fn assume(condition: bool, message: &'static str) {
+    if !condition {
+        std::panic::panic_any(Error {
+            location: core::panic::Location::caller(),
+            generator: "<assume>",
+            output: message,
+        });
+    }
+}
+
+#[track_caller]
+pub fn fill_bytes(bytes: &mut [u8]) {
+    borrow_with(|driver| {
+        let len = bytes.len();
+        let mut hint = || (len, Some(len));
+        driver
+            .0
+            .gen_from_bytes(&mut hint, &mut |src: &[u8]| {
+                if src.len() == len {
+                    bytes.copy_from_slice(src);
+                    Some(len)
+                } else {
+                    None
+                }
+            })
+            .unwrap_or_else(|| {
+                std::panic::panic_any(Error {
+                    location: core::panic::Location::caller(),
+                    generator: "<fill_bytes>",
+                    output: "could not generate enough bytes",
+                });
+            })
+    })
+}
+
+#[derive(Clone)]
+pub struct Error {
+    location: &'static core::panic::Location<'static>,
+    generator: &'static str,
+    output: &'static str,
+}
+
+impl fmt::Debug for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Error")
+            .field("location", &self.location)
+            .field("generator", &self.generator)
+            .field("output", &self.output)
+            .finish()
+    }
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Could not generate value of type {} at {}",
+            self.output, self.location,
+        )
+    }
+}
+
+impl std::error::Error for Error {}

lib/bolero-generator/src/any/kani.rs

@@ -0,0 +1,82 @@
+use std::cell::RefCell;
+
+type Type = crate::kani::Driver;
+
+pub use core::convert::Infallible as Error;
+
+/// Kani doesn't support thread_locals so use a static global instead
+static mut CURRENT: Type = Type {
+    depth: 0,
+    max_depth: crate::driver::Options::DEFAULT_MAX_DEPTH,
+};
+
+fn get() -> Type {
+    unsafe {
+        let depth = CURRENT.depth;
+        let max_depth = CURRENT.max_depth;
+
+        Type { depth, max_depth }
+    }
+}
+
+fn set(value: Type) -> Type {
+    let prev = get();
+
+    unsafe {
+        CURRENT.depth = value.depth;
+        CURRENT.max_depth = value.max_depth;
+    }
+
+    prev
+}
+
+pub fn with<F, R>(driver: Type, f: F) -> (Type, R)
+where
+    F: FnOnce() -> R,
+{
+    let prev = set(driver);
+    let res = f();
+    let driver = set(prev);
+    (driver, res)
+}
+
+fn borrow_with<F: FnOnce(&mut Type) -> R, R>(f: F) -> R {
+    let mut driver = unsafe {
+        let depth = CURRENT.depth;
+        let max_depth = CURRENT.max_depth;
+
+        Type { depth, max_depth }
+    };
+    let result = f(&mut driver);
+    set(driver);
+    result
+}
+
+pub fn any<G: crate::ValueGenerator>(g: &G) -> G::Output {
+    borrow_with(|driver| {
+        let v = g.generate(driver);
+        assume(v.is_some(), "generator should return at least one value");
+        v.unwrap()
+    })
+}
+
+pub fn fill_bytes(bytes: &mut [u8]) {
+    for dst in bytes {
+        #[cfg(kani)]
+        let src = ::kani::any();
+
+        #[cfg(not(kani))]
+        let src = 0;
+
+        *dst = src;
+    }
+}
+
+#[inline]
+pub fn assume(condition: bool, message: &'static str) {
+    #[cfg(kani)]
+    ::kani::assume(condition);
+
+    let _ = condition;
+    let _ = message;
+}