Skip to content

Commit 40fe82d

Browse files
committed
Shrink derived code size with with_recursive_count.
The amount of code generated by `derive(Arbitrary)` is large, mostly due to the recursion count guard. This commit factors out that code with a new hidden function `with_recursive_count`. Because `with_recursive_count` is mark with `inline`, the generated code should end up being much the same. But compile times are reduced by 30-40% because rustc's frontend has much less code to chew through.
1 parent 3fbc2d7 commit 40fe82d

File tree

2 files changed

+70
-29
lines changed

2 files changed

+70
-29
lines changed

derive/src/lib.rs

Lines changed: 15 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -158,31 +158,6 @@ fn add_trait_bounds(mut generics: Generics, lifetime: LifetimeParam) -> Generics
158158
generics
159159
}
160160

161-
fn with_recursive_count_guard(recursive_count: &syn::Ident, expr: TokenStream) -> TokenStream {
162-
quote! {
163-
let guard_against_recursion = u.is_empty();
164-
if guard_against_recursion {
165-
#recursive_count.with(|count| {
166-
if count.get() > 0 {
167-
return Err(arbitrary::Error::NotEnoughData);
168-
}
169-
count.set(count.get() + 1);
170-
Ok(())
171-
})?;
172-
}
173-
174-
let result = (|| { #expr })();
175-
176-
if guard_against_recursion {
177-
#recursive_count.with(|count| {
178-
count.set(count.get() - 1);
179-
});
180-
}
181-
182-
result
183-
}
184-
}
185-
186161
fn gen_arbitrary_method(
187162
input: &DeriveInput,
188163
lifetime: LifetimeParam,
@@ -195,11 +170,18 @@ fn gen_arbitrary_method(
195170
recursive_count: &syn::Ident,
196171
) -> Result<TokenStream> {
197172
let arbitrary = construct(fields, |_idx, field| gen_constructor_for_field(field))?;
198-
let body = with_recursive_count_guard(recursive_count, quote! { Ok(#ident #arbitrary) });
173+
let body = quote! {
174+
arbitrary::details::with_recursive_count(u, &#recursive_count, |mut u| {
175+
Ok(#ident #arbitrary)
176+
})
177+
};
199178

200179
let arbitrary_take_rest = construct_take_rest(fields)?;
201-
let take_rest_body =
202-
with_recursive_count_guard(recursive_count, quote! { Ok(#ident #arbitrary_take_rest) });
180+
let take_rest_body = quote! {
181+
arbitrary::details::with_recursive_count(u, &#recursive_count, |mut u| {
182+
Ok(#ident #arbitrary_take_rest)
183+
})
184+
};
203185

204186
Ok(quote! {
205187
fn arbitrary(u: &mut arbitrary::Unstructured<#lifetime>) -> arbitrary::Result<Self> {
@@ -243,7 +225,11 @@ fn gen_arbitrary_method(
243225
};
244226

245227
if needs_recursive_count {
246-
with_recursive_count_guard(recursive_count, do_variants)
228+
quote! {
229+
arbitrary::details::with_recursive_count(u, &#recursive_count, |mut u| {
230+
#do_variants
231+
})
232+
}
247233
} else {
248234
do_variants
249235
}

src/lib.rs

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -551,3 +551,58 @@ mod test {
551551
/// ```
552552
#[cfg(all(doctest, feature = "derive"))]
553553
pub struct CompileFailTests;
554+
555+
// Support for `#[derive(Arbitrary)]`.
556+
#[doc(hidden)]
557+
#[cfg(feature = "derive")]
558+
pub mod details {
559+
use super::*;
560+
561+
// Hidden trait that papers over the difference between `&mut Unstructured` and
562+
// `Unstructured` arguments so that `with_recursive_count` can be used for both
563+
// `arbitrary` and `arbitrary_take_rest`.
564+
pub trait IsEmpty {
565+
fn is_empty(&self) -> bool;
566+
}
567+
568+
impl IsEmpty for Unstructured<'_> {
569+
fn is_empty(&self) -> bool {
570+
Unstructured::is_empty(self)
571+
}
572+
}
573+
574+
impl IsEmpty for &mut Unstructured<'_> {
575+
fn is_empty(&self) -> bool {
576+
Unstructured::is_empty(self)
577+
}
578+
}
579+
580+
// Calls `f` with a recursive count guard.
581+
#[inline]
582+
pub fn with_recursive_count<U: IsEmpty, R>(
583+
u: U,
584+
recursive_count: &'static std::thread::LocalKey<std::cell::Cell<u32>>,
585+
f: impl FnOnce(U) -> Result<R>,
586+
) -> Result<R> {
587+
let guard_against_recursion = u.is_empty();
588+
if guard_against_recursion {
589+
recursive_count.with(|count| {
590+
if count.get() > 0 {
591+
return Err(Error::NotEnoughData);
592+
}
593+
count.set(count.get() + 1);
594+
Ok(())
595+
})?;
596+
}
597+
598+
let result = f(u);
599+
600+
if guard_against_recursion {
601+
recursive_count.with(|count| {
602+
count.set(count.get() - 1);
603+
});
604+
}
605+
606+
result
607+
}
608+
}

0 commit comments

Comments
 (0)