Skip to content

Commit

Permalink
Add faillible version of size_hint to properly handle recursive struc…
Browse files Browse the repository at this point in the history
…tures

The default implementation forwards to the current size_hint, so that
it can be relied upon also for structures that still use the default implementation.

The default implementation of `size_hint` is not changed, so that it does not break existing
implementations
  • Loading branch information
sosthene-nitrokey committed Sep 13, 2024
1 parent 17647a1 commit a8907c9
Show file tree
Hide file tree
Showing 16 changed files with 384 additions and 81 deletions.
36 changes: 22 additions & 14 deletions derive/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -364,27 +364,27 @@ fn gen_size_hint_method(input: &DeriveInput) -> Result<TokenStream> {
determine_field_constructor(f).map(|field_constructor| {
match field_constructor {
FieldConstructor::Default | FieldConstructor::Value(_) => {
quote!((0, Some(0)))
quote!(Ok((0, Some(0))))
}
FieldConstructor::Arbitrary => {
quote! { <#ty as arbitrary::Arbitrary>::size_hint(depth) }
quote! { <#ty as arbitrary::Arbitrary>::try_size_hint(depth) }
}

// Note that in this case it's hard to determine what size_hint must be, so size_of::<T>() is
// just an educated guess, although it's gonna be inaccurate for dynamically
// allocated types (Vec, HashMap, etc.).
FieldConstructor::With(_) => {
quote! { (::core::mem::size_of::<#ty>(), None) }
quote! { Ok((::core::mem::size_of::<#ty>(), None)) }
}
}
})
})
.collect::<Result<Vec<TokenStream>>>()
.map(|hints| {
quote! {
arbitrary::size_hint::and_all(&[
#( #hints ),*
])
Ok(arbitrary::size_hint::and_all(&[
#( #hints? ),*
]))
}
})
};
Expand All @@ -393,7 +393,12 @@ fn gen_size_hint_method(input: &DeriveInput) -> Result<TokenStream> {
quote! {
#[inline]
fn size_hint(depth: usize) -> (usize, ::core::option::Option<usize>) {
arbitrary::size_hint::recursion_guard(depth, |depth| #hint)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, ::core::option::Option<usize>), arbitrary::MaxRecursionReached> {
arbitrary::size_hint::try_recursion_guard(depth, |depth| #hint)
}
}
})
Expand All @@ -413,14 +418,17 @@ fn gen_size_hint_method(input: &DeriveInput) -> Result<TokenStream> {
.collect::<Result<Vec<TokenStream>>>()
.map(|variants| {
quote! {
#[inline]
fn size_hint(depth: usize) -> (usize, ::core::option::Option<usize>) {
arbitrary::size_hint::and(
<u32 as arbitrary::Arbitrary>::size_hint(depth),
arbitrary::size_hint::recursion_guard(depth, |depth| {
arbitrary::size_hint::or_all(&[ #( #variants ),* ])
}),
)
Self::try_size_hint(depth).unwrap_or_default()
}
#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, ::core::option::Option<usize>), arbitrary::MaxRecursionReached> {
Ok(arbitrary::size_hint::and(
<u32 as arbitrary::Arbitrary>::try_size_hint(depth)?,
arbitrary::size_hint::try_recursion_guard(depth, |depth| {
Ok(arbitrary::size_hint::or_all(&[ #( #variants? ),* ]))
})?,
))
}
}
}),
Expand Down
9 changes: 7 additions & 2 deletions src/foreign/alloc/borrow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,13 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::recursion_guard(depth, |depth| {
<<A as ToOwned>::Owned as Arbitrary>::size_hint(depth)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
size_hint::try_recursion_guard(depth, |depth| {
<<A as ToOwned>::Owned as Arbitrary>::try_size_hint(depth)
})
}
}
7 changes: 6 additions & 1 deletion src/foreign/alloc/boxed.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,12 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::recursion_guard(depth, <A as Arbitrary>::size_hint)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
size_hint::try_recursion_guard(depth, <A as Arbitrary>::try_size_hint)
}
}

Expand Down
7 changes: 6 additions & 1 deletion src/foreign/alloc/rc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,12 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::recursion_guard(depth, <A as Arbitrary>::size_hint)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
size_hint::try_recursion_guard(depth, <A as Arbitrary>::try_size_hint)
}
}

Expand Down
7 changes: 6 additions & 1 deletion src/foreign/alloc/sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,12 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::recursion_guard(depth, <A as Arbitrary>::size_hint)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
size_hint::try_recursion_guard(depth, <A as Arbitrary>::try_size_hint)
}
}

Expand Down
12 changes: 8 additions & 4 deletions src/foreign/core/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,13 @@ where
}

#[inline]
fn size_hint(d: usize) -> (usize, Option<usize>) {
size_hint::and_all(&array::from_fn::<_, N, _>(|_| {
<T as Arbitrary>::size_hint(d)
}))
fn size_hint(depth: usize) -> (usize, Option<usize>) {
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
let hint = <T as Arbitrary>::try_size_hint(depth)?;
Ok(size_hint::and_all(&array::from_fn::<_, N, _>(|_| hint)))
}
}
23 changes: 19 additions & 4 deletions src/foreign/core/cell.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use {
crate::{Arbitrary, Result, Unstructured},
crate::{Arbitrary, MaxRecursionReached, Result, Unstructured},
core::cell::{Cell, RefCell, UnsafeCell},
};

Expand All @@ -13,7 +13,12 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
<A as Arbitrary<'a>>::size_hint(depth)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
<A as Arbitrary<'a>>::try_size_hint(depth)
}
}

Expand All @@ -27,7 +32,12 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
<A as Arbitrary<'a>>::size_hint(depth)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
<A as Arbitrary<'a>>::try_size_hint(depth)
}
}

Expand All @@ -41,6 +51,11 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
<A as Arbitrary<'a>>::size_hint(depth)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
<A as Arbitrary<'a>>::try_size_hint(depth)
}
}
9 changes: 7 additions & 2 deletions src/foreign/core/num.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use {
crate::{Arbitrary, Error, Result, Unstructured},
crate::{Arbitrary, Error, MaxRecursionReached, Result, Unstructured},
core::{
mem,
num::{
Expand Down Expand Up @@ -131,6 +131,11 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
<A as Arbitrary<'a>>::size_hint(depth)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
<A as Arbitrary<'a>>::try_size_hint(depth)
}
}
41 changes: 25 additions & 16 deletions src/foreign/core/ops.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use {
crate::{size_hint, Arbitrary, Result, Unstructured},
crate::{size_hint, Arbitrary, MaxRecursionReached, Result, Unstructured},
core::{
mem,
ops::{Bound, Range, RangeBounds, RangeFrom, RangeInclusive, RangeTo, RangeToInclusive},
Expand All @@ -25,53 +25,57 @@ macro_rules! impl_range {

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
#[allow(clippy::redundant_closure_call)]
$size_hint_closure(depth)
}
}
};
}

impl_range!(
Range<A>,
|r: &Range<A>| (r.start.clone(), r.end.clone()),
(A, A),
bounded_range(|(a, b)| a..b),
|depth| size_hint::and(
<A as Arbitrary>::size_hint(depth),
<A as Arbitrary>::size_hint(depth)
)
|depth| Ok(crate::size_hint::and(
<A as Arbitrary>::try_size_hint(depth)?,
<A as Arbitrary>::try_size_hint(depth)?,
))
);
impl_range!(
RangeFrom<A>,
|r: &RangeFrom<A>| r.start.clone(),
A,
unbounded_range(|a| a..),
|depth| <A as Arbitrary>::size_hint(depth)
|depth| <A as Arbitrary>::try_size_hint(depth)
);
impl_range!(
RangeInclusive<A>,
|r: &RangeInclusive<A>| (r.start().clone(), r.end().clone()),
(A, A),
bounded_range(|(a, b)| a..=b),
|depth| size_hint::and(
<A as Arbitrary>::size_hint(depth),
<A as Arbitrary>::size_hint(depth)
)
|depth| Ok(crate::size_hint::and(
<A as Arbitrary>::try_size_hint(depth)?,
<A as Arbitrary>::try_size_hint(depth)?,
))
);
impl_range!(
RangeTo<A>,
|r: &RangeTo<A>| r.end.clone(),
A,
unbounded_range(|b| ..b),
|depth| <A as Arbitrary>::size_hint(depth)
|depth| <A as Arbitrary>::try_size_hint(depth)
);
impl_range!(
RangeToInclusive<A>,
|r: &RangeToInclusive<A>| r.end.clone(),
A,
unbounded_range(|b| ..=b),
|depth| <A as Arbitrary>::size_hint(depth)
|depth| <A as Arbitrary>::try_size_hint(depth)
);

pub(crate) fn bounded_range<CB, I, R>(bounds: (I, I), cb: CB) -> R
Expand Down Expand Up @@ -110,9 +114,14 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::or(
size_hint::and((1, Some(1)), A::size_hint(depth)),
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
Ok(size_hint::or(
size_hint::and((1, Some(1)), A::try_size_hint(depth)?),
(1, Some(1)),
)
))
}
}
15 changes: 10 additions & 5 deletions src/foreign/core/option.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::{size_hint, Arbitrary, Result, Unstructured};
use crate::{size_hint, Arbitrary, MaxRecursionReached, Result, Unstructured};

impl<'a, A> Arbitrary<'a> for Option<A>
where
Expand All @@ -14,9 +14,14 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::and(
<bool as Arbitrary>::size_hint(depth),
size_hint::or((0, Some(0)), <A as Arbitrary>::size_hint(depth)),
)
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
Ok(size_hint::and(
<bool as Arbitrary>::try_size_hint(depth)?,
size_hint::or((0, Some(0)), <A as Arbitrary>::try_size_hint(depth)?),
))
}
}
15 changes: 10 additions & 5 deletions src/foreign/core/result.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::{size_hint, Arbitrary, Error, Unstructured};
use crate::{size_hint, Arbitrary, Error, MaxRecursionReached, Unstructured};

impl<'a, T, E> Arbitrary<'a> for Result<T, E>
where
Expand All @@ -15,12 +15,17 @@ where

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::and(
Self::try_size_hint(depth).unwrap_or_default()
}

#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
Ok(size_hint::and(
<bool as Arbitrary>::size_hint(depth),
size_hint::or(
<T as Arbitrary>::size_hint(depth),
<E as Arbitrary>::size_hint(depth),
<T as Arbitrary>::try_size_hint(depth)?,
<E as Arbitrary>::try_size_hint(depth)?,
),
)
))
}
}
14 changes: 9 additions & 5 deletions src/foreign/core/tuple.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::{size_hint, Arbitrary, Result, Unstructured};
use crate::{size_hint, Arbitrary, MaxRecursionReached, Result, Unstructured};

macro_rules! arbitrary_tuple {
() => {};
Expand All @@ -23,10 +23,14 @@ macro_rules! arbitrary_tuple {

#[inline]
fn size_hint(depth: usize) -> (usize, Option<usize>) {
size_hint::and_all(&[
<$last as Arbitrary>::size_hint(depth),
$( <$xs as Arbitrary>::size_hint(depth) ),*
])
Self::try_size_hint(depth).unwrap_or_default()
}
#[inline]
fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), MaxRecursionReached> {
Ok(size_hint::and_all(&[
<$last as Arbitrary>::try_size_hint(depth)?,
$( <$xs as Arbitrary>::try_size_hint(depth)?),*
]))
}
}
};
Expand Down
Loading

0 comments on commit a8907c9

Please sign in to comment.