add boxed union feature to circumvent stack overflow

For enums with a lot of variants the `Strategy` synthesized by the
`#[derive(Arbitrary)]` macro is a right-deep `TupleUnion`.

This commit adds a `boxed_union` feature, which when enabled will change
the `enum` expansion to use `Union<BoxedStrategy<Self>>`.

see github issues:
#249
#152

code included from #249 (comment)

Co-authored-by: Alexander Alexandrov <[email protected]>
Co-authored-by: Sophie Dawson <[email protected]>

.github/workflows/rust.yml

@@ -70,3 +70,6 @@ jobs:
       - name: Clean and Run tests for derive
         if: ${{ matrix.build == 'nightly' }}
         run: cd proptest-derive && cargo clean && cargo test
+      - name: Clean and Run tests for derive with features enabled
+        if: ${{ matrix.build == 'nightly' }}
+        run: cd proptest-derive && cargo clean && cargo test --features boxed_union

proptest-derive/Cargo.toml

@@ -37,6 +37,10 @@ proc-macro2 = "1.0"
 syn = { version = "1.0.0", features = ["visit", "extra-traits", "full"] }
 quote = "1.0"
 
+[features]
+# Don't generate TupleUnion structs in #[derive(Arbitrary)] code
+boxed_union = []
+
 [[bench]]
 name = "large_enum"
 harness = false

proptest-derive/src/ast.rs

@@ -29,6 +29,7 @@ use crate::util::self_ty;
 /// Increase this if the behaviour is changed in `proptest`.
 /// Keeping this lower than what `proptest` supports will also work
 /// but for optimality this should follow what `proptest` supports.
+#[cfg(not(feature = "boxed_union"))]
 const UNION_CHUNK_SIZE: usize = 9;
 
 /// The `MAX - 1` tuple length `Arbitrary` is implemented for. After this number,
@@ -380,7 +381,10 @@ impl ToTokens for Strategy {
                     >
                 )
             }
+            #[cfg(not(feature = "boxed_union"))]
             Union(strats) => union_strat_to_tokens(tokens, strats),
+            #[cfg(feature = "boxed_union")]
+            Union(strats) => union_strat_to_tokens_boxed(tokens, strats),
             Filter(strat, ty) => quote_append!(tokens,
                 _proptest::strategy::Filter<#strat, fn(&#ty) -> bool>
             ),
@@ -508,14 +512,17 @@ impl ToTokens for Ctor {
             ),
             Existential(expr) => quote_append!(tokens,
                 _proptest::strategy::Strategy::boxed( #expr ) ),
-            Value(expr) => quote_append!(tokens, || #expr ),
+            Value(expr) => quote_append!(tokens, (|| #expr) as fn() -> _),
             ValueExistential(expr) => quote_append!(tokens,
                 _proptest::strategy::Strategy::boxed(
                     _proptest::strategy::LazyJust::new(move || #expr)
                 )
             ),
             Map(ctors, closure) => map_ctor_to_tokens(tokens, &ctors, closure),
+            #[cfg(not(feature = "boxed_union"))]
             Union(ctors) => union_ctor_to_tokens(tokens, ctors),
+            #[cfg(feature = "boxed_union")]
+            Union(ctors) => union_ctor_to_tokens_boxed(tokens, ctors),
         }
     }
 }
@@ -607,6 +614,7 @@ fn map_ctor_to_tokens(
 ///     (10, 11, 12, 13, 14, 15, 16, 17, 18,
 ///         (19, ..)))
 /// ```
+#[cfg(not(feature = "boxed_union"))]
 fn union_ctor_to_tokens(tokens: &mut TokenStream, ctors: &[(u32, Ctor)]) {
     if ctors.is_empty() {
         return;
@@ -665,6 +673,7 @@ fn union_ctor_to_tokens(tokens: &mut TokenStream, ctors: &[(u32, Ctor)]) {
 
 /// Tokenizes a weighted list of `Strategy`.
 /// For details, see `union_ctor_to_tokens`.
+#[cfg(not(feature = "boxed_union"))]
 fn union_strat_to_tokens(tokens: &mut TokenStream, strats: &[Strategy]) {
     if strats.is_empty() {
         return;
@@ -714,6 +723,55 @@ fn union_strat_to_tokens(tokens: &mut TokenStream, strats: &[Strategy]) {
     }
 }
 
+/// Tokenizes a weighted list of `Ctor`.
+///
+/// This can be used instead of `union_ctor_to_tokens` to generate a boxing
+/// macro.
+#[cfg(feature = "boxed_union")]
+fn union_ctor_to_tokens_boxed(tokens: &mut TokenStream, ctors: &[(u32, Ctor)]) {
+    if ctors.is_empty() {
+        return;
+    }
+
+    if let [(_, ctor)] = ctors {
+        // This is not a union at all - user provided an enum with one variant.
+        ctor.to_tokens(tokens);
+        return;
+    }
+
+    let ctors_boxed = ctors.iter().map(wrap_boxed);
+
+    quote_append!(
+        tokens,
+        _proptest::strategy::Union::new_weighted(vec![ #(#ctors_boxed,)* ])
+    );
+
+    fn wrap_boxed(arg: &(u32, Ctor)) -> TokenStream {
+        let (w, c) = arg;
+        quote!( (#w, _proptest::strategy::Strategy::boxed(#c)) )
+    }
+}
+
+/// Tokenizes a weighted list of `Strategy`.
+/// For details, see `union_ctor_to_tokens_boxed`.
+#[cfg(feature = "boxed_union")]
+fn union_strat_to_tokens_boxed(tokens: &mut TokenStream, strats: &[Strategy]) {
+    if strats.is_empty() {
+        return;
+    }
+
+    if let [strat] = strats {
+        // This is not a union at all - user provided an enum with one variant.
+        strat.to_tokens(tokens);
+        return;
+    }
+
+    quote_append!(
+        tokens,
+        _proptest::strategy::Union<_proptest::strategy::BoxedStrategy<Self>>
+    );
+}
+
 /// Wraps a `Ctor` that expects the `to` "register" to be filled with
 /// contents of the `from` register. The correctness of this wrt. the
 /// generated Rust code has to be verified externally by checking the

proptest-derive/src/tests.rs

@@ -87,7 +87,7 @@ test! {
             type Strategy = fn() -> Self;
 
             fn arbitrary_with(_top: Self::Parameters) -> Self::Strategy {
-                || MyUnitStruct {}
+                (|| MyUnitStruct {}) as fn() -> _
             }
         }
         };
@@ -108,7 +108,7 @@ test! {
             type Strategy = fn() -> Self;
 
             fn arbitrary_with(_top: Self::Parameters) -> Self::Strategy {
-                || MyTupleUnitStruct {}
+                (|| MyTupleUnitStruct {}) as fn() -> _
             }
         }
         };
@@ -129,7 +129,7 @@ test! {
             type Strategy = fn() -> Self;
 
             fn arbitrary_with(_top: Self::Parameters) -> Self::Strategy {
-                || MyNamedUnitStruct {}
+                (|| MyNamedUnitStruct {}) as fn () -> _
             }
         }
         };