use crate::{Arbitrary, Error, Result};
use std::marker::PhantomData;
use std::ops::ControlFlow;
use std::{mem, ops};
#[derive(Debug)]
pub struct Unstructured<'a> {
data: &'a [u8],
}
impl<'a> Unstructured<'a> {
pub fn new(data: &'a [u8]) -> Self {
Unstructured { data }
}
#[inline]
pub fn len(&self) -> usize {
self.data.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn arbitrary<A>(&mut self) -> Result<A>
where
A: Arbitrary<'a>,
{
<A as Arbitrary<'a>>::arbitrary(self)
}
pub fn arbitrary_len<ElementType>(&mut self) -> Result<usize>
where
ElementType: Arbitrary<'a>,
{
let byte_size = self.arbitrary_byte_size()?;
let (lower, upper) = <ElementType as Arbitrary>::size_hint(0);
let elem_size = upper.unwrap_or(lower * 2);
let elem_size = std::cmp::max(1, elem_size);
Ok(byte_size / elem_size)
}
fn arbitrary_byte_size(&mut self) -> Result<usize> {
if self.data.is_empty() {
Ok(0)
} else if self.data.len() == 1 {
self.data = &[];
Ok(0)
} else {
let len = if self.data.len() as u64 <= u8::MAX as u64 + 1 {
let bytes = 1;
let max_size = self.data.len() - bytes;
let (rest, for_size) = self.data.split_at(max_size);
self.data = rest;
Self::int_in_range_impl(0..=max_size as u8, for_size.iter().copied())?.0 as usize
} else if self.data.len() as u64 <= u16::MAX as u64 + 2 {
let bytes = 2;
let max_size = self.data.len() - bytes;
let (rest, for_size) = self.data.split_at(max_size);
self.data = rest;
Self::int_in_range_impl(0..=max_size as u16, for_size.iter().copied())?.0 as usize
} else if self.data.len() as u64 <= u32::MAX as u64 + 4 {
let bytes = 4;
let max_size = self.data.len() - bytes;
let (rest, for_size) = self.data.split_at(max_size);
self.data = rest;
Self::int_in_range_impl(0..=max_size as u32, for_size.iter().copied())?.0 as usize
} else {
let bytes = 8;
let max_size = self.data.len() - bytes;
let (rest, for_size) = self.data.split_at(max_size);
self.data = rest;
Self::int_in_range_impl(0..=max_size as u64, for_size.iter().copied())?.0 as usize
};
Ok(len)
}
}
pub fn int_in_range<T>(&mut self, range: ops::RangeInclusive<T>) -> Result<T>
where
T: Int,
{
let (result, bytes_consumed) = Self::int_in_range_impl(range, self.data.iter().cloned())?;
self.data = &self.data[bytes_consumed..];
Ok(result)
}
fn int_in_range_impl<T>(
range: ops::RangeInclusive<T>,
mut bytes: impl Iterator<Item = u8>,
) -> Result<(T, usize)>
where
T: Int,
{
let start = *range.start();
let end = *range.end();
assert!(
start <= end,
"`arbitrary::Unstructured::int_in_range` requires a non-empty range"
);
if start == end {
return Ok((start, 0));
}
let start = start.to_unsigned();
let end = end.to_unsigned();
let delta = end.wrapping_sub(start);
debug_assert_ne!(delta, T::Unsigned::ZERO);
let mut arbitrary_int = T::Unsigned::ZERO;
let mut bytes_consumed: usize = 0;
while (bytes_consumed < mem::size_of::<T>())
&& (delta >> T::Unsigned::from_usize(bytes_consumed * 8)) > T::Unsigned::ZERO
{
let byte = match bytes.next() {
None => break,
Some(b) => b,
};
bytes_consumed += 1;
arbitrary_int = if mem::size_of::<T>() == 1 {
T::Unsigned::from_u8(byte)
} else {
(arbitrary_int << 8) | T::Unsigned::from_u8(byte)
};
}
let offset = if delta == T::Unsigned::MAX {
arbitrary_int
} else {
arbitrary_int % (delta.checked_add(T::Unsigned::ONE).unwrap())
};
let result = start.wrapping_add(offset);
let result = T::from_unsigned(result);
debug_assert!(*range.start() <= result);
debug_assert!(result <= *range.end());
Ok((result, bytes_consumed))
}
pub fn choose<'b, T>(&mut self, choices: &'b [T]) -> Result<&'b T> {
let idx = self.choose_index(choices.len())?;
Ok(&choices[idx])
}
pub fn choose_iter<T, I>(&mut self, choices: I) -> Result<T>
where
I: IntoIterator<Item = T>,
I::IntoIter: ExactSizeIterator,
{
let mut choices = choices.into_iter();
let idx = self.choose_index(choices.len())?;
let choice = choices
.nth(idx)
.expect("ExactSizeIterator should have correct len");
Ok(choice)
}
pub fn choose_index(&mut self, len: usize) -> Result<usize> {
if len == 0 {
return Err(Error::EmptyChoose);
}
let idx = self.int_in_range(0..=len - 1)?;
Ok(idx)
}
pub fn ratio<T>(&mut self, numerator: T, denominator: T) -> Result<bool>
where
T: Int,
{
assert!(T::ZERO < numerator);
assert!(numerator <= denominator);
let x = self.int_in_range(T::ONE..=denominator)?;
Ok(x <= numerator)
}
pub fn fill_buffer(&mut self, buffer: &mut [u8]) -> Result<()> {
let n = std::cmp::min(buffer.len(), self.data.len());
buffer[..n].copy_from_slice(&self.data[..n]);
for byte in buffer[n..].iter_mut() {
*byte = 0;
}
self.data = &self.data[n..];
Ok(())
}
pub fn bytes(&mut self, size: usize) -> Result<&'a [u8]> {
if self.data.len() < size {
return Err(Error::NotEnoughData);
}
let (for_buf, rest) = self.data.split_at(size);
self.data = rest;
Ok(for_buf)
}
pub fn peek_bytes(&self, size: usize) -> Option<&'a [u8]> {
self.data.get(..size)
}
pub fn take_rest(mut self) -> &'a [u8] {
mem::take(&mut self.data)
}
pub fn arbitrary_iter<'b, ElementType: Arbitrary<'a>>(
&'b mut self,
) -> Result<ArbitraryIter<'a, 'b, ElementType>> {
Ok(ArbitraryIter {
u: &mut *self,
_marker: PhantomData,
})
}
pub fn arbitrary_take_rest_iter<ElementType: Arbitrary<'a>>(
self,
) -> Result<ArbitraryTakeRestIter<'a, ElementType>> {
Ok(ArbitraryTakeRestIter {
u: self,
_marker: PhantomData,
})
}
pub fn arbitrary_loop(
&mut self,
min: Option<u32>,
max: Option<u32>,
mut f: impl FnMut(&mut Self) -> Result<ControlFlow<(), ()>>,
) -> Result<()> {
let min = min.unwrap_or(0);
let max = max.unwrap_or(u32::MAX);
for _ in 0..self.int_in_range(min..=max)? {
match f(self)? {
ControlFlow::Continue(_) => continue,
ControlFlow::Break(_) => break,
}
}
Ok(())
}
}
pub struct ArbitraryIter<'a, 'b, ElementType> {
u: &'b mut Unstructured<'a>,
_marker: PhantomData<ElementType>,
}
impl<'a, 'b, ElementType: Arbitrary<'a>> Iterator for ArbitraryIter<'a, 'b, ElementType> {
type Item = Result<ElementType>;
fn next(&mut self) -> Option<Result<ElementType>> {
let keep_going = self.u.arbitrary().unwrap_or(false);
if keep_going {
Some(Arbitrary::arbitrary(self.u))
} else {
None
}
}
}
pub struct ArbitraryTakeRestIter<'a, ElementType> {
u: Unstructured<'a>,
_marker: PhantomData<ElementType>,
}
impl<'a, ElementType: Arbitrary<'a>> Iterator for ArbitraryTakeRestIter<'a, ElementType> {
type Item = Result<ElementType>;
fn next(&mut self) -> Option<Result<ElementType>> {
let keep_going = self.u.arbitrary().unwrap_or(false);
if keep_going {
Some(Arbitrary::arbitrary(&mut self.u))
} else {
None
}
}
}
pub trait Int:
Copy
+ std::fmt::Debug
+ PartialOrd
+ Ord
+ ops::Sub<Self, Output = Self>
+ ops::Rem<Self, Output = Self>
+ ops::Shr<Self, Output = Self>
+ ops::Shl<usize, Output = Self>
+ ops::BitOr<Self, Output = Self>
{
#[doc(hidden)]
type Unsigned: Int;
#[doc(hidden)]
const ZERO: Self;
#[doc(hidden)]
const ONE: Self;
#[doc(hidden)]
const MAX: Self;
#[doc(hidden)]
fn from_u8(b: u8) -> Self;
#[doc(hidden)]
fn from_usize(u: usize) -> Self;
#[doc(hidden)]
fn checked_add(self, rhs: Self) -> Option<Self>;
#[doc(hidden)]
fn wrapping_add(self, rhs: Self) -> Self;
#[doc(hidden)]
fn wrapping_sub(self, rhs: Self) -> Self;
#[doc(hidden)]
fn to_unsigned(self) -> Self::Unsigned;
#[doc(hidden)]
fn from_unsigned(unsigned: Self::Unsigned) -> Self;
}
macro_rules! impl_int {
( $( $ty:ty : $unsigned_ty: ty ; )* ) => {
$(
impl Int for $ty {
type Unsigned = $unsigned_ty;
const ZERO: Self = 0;
const ONE: Self = 1;
const MAX: Self = Self::MAX;
fn from_u8(b: u8) -> Self {
b as Self
}
fn from_usize(u: usize) -> Self {
u as Self
}
fn checked_add(self, rhs: Self) -> Option<Self> {
<$ty>::checked_add(self, rhs)
}
fn wrapping_add(self, rhs: Self) -> Self {
<$ty>::wrapping_add(self, rhs)
}
fn wrapping_sub(self, rhs: Self) -> Self {
<$ty>::wrapping_sub(self, rhs)
}
fn to_unsigned(self) -> Self::Unsigned {
self as $unsigned_ty
}
fn from_unsigned(unsigned: $unsigned_ty) -> Self {
unsigned as Self
}
}
)*
}
}
impl_int! {
u8: u8;
u16: u16;
u32: u32;
u64: u64;
u128: u128;
usize: usize;
i8: u8;
i16: u16;
i32: u32;
i64: u64;
i128: u128;
isize: usize;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_byte_size() {
let mut u = Unstructured::new(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 6]);
assert_eq!(u.arbitrary_byte_size().unwrap(), 6);
assert_eq!(u.len(), 9);
let mut v = vec![0; 260];
v.push(1);
v.push(4);
let mut u = Unstructured::new(&v);
assert_eq!(u.arbitrary_byte_size().unwrap(), 0x104);
assert_eq!(u.len(), 260);
}
#[test]
fn int_in_range_of_one() {
let mut u = Unstructured::new(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 6]);
let x = u.int_in_range(0..=0).unwrap();
assert_eq!(x, 0);
let choice = *u.choose(&[42]).unwrap();
assert_eq!(choice, 42)
}
#[test]
fn int_in_range_uses_minimal_amount_of_bytes() {
let mut u = Unstructured::new(&[1, 2]);
assert_eq!(1, u.int_in_range::<u8>(0..=u8::MAX).unwrap());
assert_eq!(u.len(), 1);
let mut u = Unstructured::new(&[1, 2]);
assert_eq!(1, u.int_in_range::<u32>(0..=u8::MAX as u32).unwrap());
assert_eq!(u.len(), 1);
let mut u = Unstructured::new(&[1]);
assert_eq!(1, u.int_in_range::<u32>(0..=u8::MAX as u32 + 1).unwrap());
assert!(u.is_empty());
}
#[test]
fn int_in_range_in_bounds() {
for input in u8::MIN..=u8::MAX {
let input = [input];
let mut u = Unstructured::new(&input);
let x = u.int_in_range(1..=u8::MAX).unwrap();
assert_ne!(x, 0);
let mut u = Unstructured::new(&input);
let x = u.int_in_range(0..=u8::MAX - 1).unwrap();
assert_ne!(x, u8::MAX);
}
}
#[test]
fn int_in_range_covers_unsigned_range() {
let mut full = [false; u8::MAX as usize + 1];
let mut no_zero = [false; u8::MAX as usize];
let mut no_max = [false; u8::MAX as usize];
let mut narrow = [false; 10];
for input in u8::MIN..=u8::MAX {
let input = [input];
let mut u = Unstructured::new(&input);
let x = u.int_in_range(0..=u8::MAX).unwrap();
full[x as usize] = true;
let mut u = Unstructured::new(&input);
let x = u.int_in_range(1..=u8::MAX).unwrap();
no_zero[x as usize - 1] = true;
let mut u = Unstructured::new(&input);
let x = u.int_in_range(0..=u8::MAX - 1).unwrap();
no_max[x as usize] = true;
let mut u = Unstructured::new(&input);
let x = u.int_in_range(100..=109).unwrap();
narrow[x as usize - 100] = true;
}
for (i, covered) in full.iter().enumerate() {
assert!(covered, "full[{}] should have been generated", i);
}
for (i, covered) in no_zero.iter().enumerate() {
assert!(covered, "no_zero[{}] should have been generated", i);
}
for (i, covered) in no_max.iter().enumerate() {
assert!(covered, "no_max[{}] should have been generated", i);
}
for (i, covered) in narrow.iter().enumerate() {
assert!(covered, "narrow[{}] should have been generated", i);
}
}
#[test]
fn int_in_range_covers_signed_range() {
let mut full = [false; u8::MAX as usize + 1];
let mut no_min = [false; u8::MAX as usize];
let mut no_max = [false; u8::MAX as usize];
let mut narrow = [false; 21];
let abs_i8_min: isize = 128;
for input in 0..=u8::MAX {
let input = [input];
let mut u = Unstructured::new(&input);
let x = u.int_in_range(i8::MIN..=i8::MAX).unwrap();
full[(x as isize + abs_i8_min) as usize] = true;
let mut u = Unstructured::new(&input);
let x = u.int_in_range(i8::MIN + 1..=i8::MAX).unwrap();
no_min[(x as isize + abs_i8_min - 1) as usize] = true;
let mut u = Unstructured::new(&input);
let x = u.int_in_range(i8::MIN..=i8::MAX - 1).unwrap();
no_max[(x as isize + abs_i8_min) as usize] = true;
let mut u = Unstructured::new(&input);
let x = u.int_in_range(-10..=10).unwrap();
narrow[(x as isize + 10) as usize] = true;
}
for (i, covered) in full.iter().enumerate() {
assert!(covered, "full[{}] should have been generated", i);
}
for (i, covered) in no_min.iter().enumerate() {
assert!(covered, "no_min[{}] should have been generated", i);
}
for (i, covered) in no_max.iter().enumerate() {
assert!(covered, "no_max[{}] should have been generated", i);
}
for (i, covered) in narrow.iter().enumerate() {
assert!(covered, "narrow[{}] should have been generated", i);
}
}
}