Skip to content

Commit

Permalink
feat(bolero-generator): add support for calling generators from anywh…
Browse files Browse the repository at this point in the history
…ere (#258)
  • Loading branch information
camshaft authored Nov 26, 2024
1 parent 665f1c7 commit e01216c
Show file tree
Hide file tree
Showing 17 changed files with 793 additions and 246 deletions.
3 changes: 3 additions & 0 deletions lib/bolero-engine/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@ pub use bolero_generator::{
TypeGenerator, ValueGenerator,
};

#[cfg(kani)]
pub use bolero_generator::kani;

pub type Seed = u128;

#[cfg(not(kani))]
Expand Down
6 changes: 5 additions & 1 deletion lib/bolero-generator/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
116 changes: 116 additions & 0 deletions lib/bolero-generator/src/any.rs
Original file line number Diff line number Diff line change
@@ -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()
}
162 changes: 162 additions & 0 deletions lib/bolero-generator/src/any/default.rs
Original file line number Diff line number Diff line change
@@ -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 {}
82 changes: 82 additions & 0 deletions lib/bolero-generator/src/any/kani.rs
Original file line number Diff line number Diff line change
@@ -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;
}
Loading

0 comments on commit e01216c

Please sign in to comment.