proptest_derive/
derive.rs

1// Copyright 2018 The proptest developers
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9//! Provides actual deriving logic for the crate.
10
11use proc_macro2::{Span, TokenStream};
12use syn::spanned::Spanned;
13use syn::{DeriveInput, Expr, Field, Ident, Path, Type, Variant};
14
15use crate::ast::*;
16use crate::attr::{self, ParamsMode, ParsedAttributes, StratMode};
17use crate::error::{self, Context, Ctx, DeriveResult};
18use crate::use_tracking::{UseMarkable, UseTracker};
19use crate::util::{fields_to_vec, is_unit_type, self_ty};
20use crate::void::IsUninhabited;
21
22//==============================================================================
23// API
24//==============================================================================
25
26pub fn impl_proptest_arbitrary(ast: DeriveInput) -> TokenStream {
27    let mut ctx = Context::default();
28    let result = derive_proptest_arbitrary(&mut ctx, ast);
29    match (result, ctx.check()) {
30        (Ok(derive), Ok(())) => derive,
31        (_, Err(err)) => err,
32        (Err(result), Ok(())) => panic!(
33            "[proptest_derive]: internal error, this is a bug! \
34             result: {:?}",
35            result
36        ),
37    }
38}
39
40/// Simplified version of `DeriveInput` from syn letting us be generic over
41/// the body.
42struct DeriveData<B> {
43    ident: Ident,
44    attrs: ParsedAttributes,
45    tracker: UseTracker,
46    body: B,
47}
48
49/// Entry point for deriving `Arbitrary`.
50fn derive_proptest_arbitrary(
51    ctx: Ctx,
52    ast: DeriveInput,
53) -> DeriveResult<TokenStream> {
54    use syn::Data::*;
55
56    // Deny lifetimes on type.
57    error::if_has_lifetimes(ctx, &ast);
58
59    // Parse top level attributes:
60    let attrs = attr::parse_top_attributes(ctx, &ast.attrs)?;
61
62    // Initialize tracker:
63    let mut tracker = UseTracker::new(ast.generics);
64    if attrs.no_bound {
65        tracker.no_track();
66    }
67
68    // Compile into our own high level IR for the impl:
69    let the_impl = match ast.data {
70        // Deal with structs:
71        Struct(data) => derive_struct(
72            ctx,
73            DeriveData {
74                tracker,
75                attrs,
76                ident: ast.ident,
77                body: fields_to_vec(data.fields),
78            },
79        ),
80        // Deal with enums:
81        Enum(data) => derive_enum(
82            ctx,
83            DeriveData {
84                tracker,
85                attrs,
86                ident: ast.ident,
87                body: data.variants.into_iter().collect(),
88            },
89        ),
90        // Unions are not supported:
91        _ => error::not_struct_or_enum(ctx)?,
92    }?;
93
94    // Linearise the IR into Rust code:
95    let q = the_impl.into_tokens(ctx)?;
96
97    // We're done!
98    Ok(q)
99}
100
101//==============================================================================
102// Struct
103//==============================================================================
104
105/// Entry point for deriving `Arbitrary` for `struct`s.
106fn derive_struct(
107    ctx: Ctx,
108    mut ast: DeriveData<Vec<Field>>,
109) -> DeriveResult<Impl> {
110    // Deny attributes that are only for enum variants:
111    error::if_enum_attrs_present(ctx, &ast.attrs, error::STRUCT);
112
113    // Deny an explicit strategy directly on the struct.
114    error::if_strategy_present(ctx, &ast.attrs, error::STRUCT);
115
116    let v_path = ast.ident.clone().into();
117    let parts = if ast.body.is_empty() {
118        // Deriving for a unit struct.
119        error::if_present_on_unit_struct(ctx, &ast.attrs);
120        let (strat, ctor) = pair_unit_self(&v_path);
121        (Params::empty(), strat, ctor)
122    } else {
123        // Not a unit struct.
124
125        // Ensures that the fields of the given struct has fields which are all
126        // inhabited. If one field is uninhabited, the entire product type is
127        // uninhabited.
128        //
129        // A unit struct in the other branch is by definition always inhabited.
130        if (&*ast.body).is_uninhabited() {
131            error::uninhabited_struct(ctx);
132        }
133
134        // Construct the closure for `.prop_map`:
135        let closure = map_closure(v_path, &ast.body);
136
137        // The complexity of the logic depends mostly now on whether
138        // parameters were set directly on the type or not.
139        let parts = if let Some(param_ty) = ast.attrs.params.into_option() {
140            // Parameters was set on the struct itself, the logic is simpler.
141            add_top_params(
142                param_ty,
143                derive_product_has_params(
144                    ctx,
145                    &mut ast.tracker,
146                    error::STRUCT_FIELD,
147                    closure,
148                    ast.body,
149                )?,
150            )
151        } else {
152            // We need considerably more complex logic.
153            derive_product_no_params(
154                ctx,
155                &mut ast.tracker,
156                ast.body,
157                error::STRUCT_FIELD,
158            )?
159            .finish(closure)
160        };
161
162        // Possibly apply filter:
163        add_top_filter(ast.attrs.filter, parts)
164    };
165
166    // We're done!
167    Ok(Impl::new(ast.ident, ast.tracker, parts))
168}
169
170/// Apply the filter at the top level if provided.
171fn add_top_filter(filter: Vec<Expr>, parts: ImplParts) -> ImplParts {
172    let (params, strat, ctor) = parts;
173    let (strat, ctor) = add_filter_self(filter, (strat, ctor));
174    (params, strat, ctor)
175}
176
177/// Apply a filter with `Self` as the input type to the predicate.
178fn add_filter_self(filter: Vec<Expr>, pair: StratPair) -> StratPair {
179    pair_filter(filter, self_ty(), pair)
180}
181
182/// Determine the `Parameters` part. We've already handled everything else.
183/// After this, we have all parts needed for an impl. If `None` is given,
184/// then the unit type `()` will be used for `Parameters`.
185fn add_top_params(
186    param_ty: Option<Type>,
187    (strat, ctor): StratPair,
188) -> ImplParts {
189    let params = Params::empty();
190    if let Some(params_ty) = param_ty {
191        // We need to add `let params = _top;`.
192        (params + params_ty, strat, extract_api(ctor, FromReg::Top))
193    } else {
194        (params, strat, ctor)
195    }
196}
197
198/// Deriving for a list of fields (product type) on
199/// which `params` or `no_params` was set directly.
200fn derive_product_has_params(
201    ctx: Ctx,
202    ut: &mut UseTracker,
203    item: &str,
204    closure: MapClosure,
205    fields: Vec<Field>,
206) -> DeriveResult<StratPair> {
207    // Fold into an accumulator of the strategy types and the expressions
208    // that produces the strategy. Finally turn the accumulator into
209    // a `.prop_map(..)` that produces the composite strategy.
210    let len = fields.len();
211    fields
212        .into_iter()
213        .try_fold(StratAcc::new(len), |acc, field| {
214            let attrs = attr::parse_attributes(ctx, &field.attrs)?;
215
216            // Deny attributes that are only for enum variants:
217            error::if_enum_attrs_present(ctx, &attrs, item);
218
219            // Deny setting parameters on the field since it has been set on parent:
220            error::if_specified_params(ctx, &attrs, item);
221
222            // Determine the strategy for this field and add it to acc.
223            let span = field.span();
224            let ty = field.ty.clone();
225            let pair =
226                product_handle_default_params(ut, ty, span, attrs.strategy);
227            let pair = pair_filter(attrs.filter, field.ty, pair);
228            Ok(acc.add(pair))
229        })
230        .map(|acc| acc.finish(closure))
231}
232
233/// Determine strategy using "Default" semantics for a product.
234fn product_handle_default_params(
235    ut: &mut UseTracker,
236    ty: Type,
237    span: Span,
238    strategy: StratMode,
239) -> StratPair {
240    match strategy {
241        // Specific strategy - use the given expr and erase the type
242        // (since we don't know about it):
243        StratMode::Strategy(strat) => pair_existential(ty, strat),
244        // Specific value - use the given expr:
245        StratMode::Value(value) => pair_value(ty, value),
246        // Specific regex - dispatch to `_regex` function based on `ty`:
247        StratMode::Regex(regex) => pair_regex(ty, regex),
248        // Use Arbitrary for the given type and mark the type as used:
249        StratMode::Arbitrary => {
250            ty.mark_uses(ut);
251            pair_any(ty, span)
252        }
253    }
254}
255
256/// Deriving for a list of fields (product type) on
257/// which `params` or `no_params` was NOT set directly.
258fn derive_product_no_params(
259    ctx: Ctx,
260    ut: &mut UseTracker,
261    fields: Vec<Field>,
262    item: &str,
263) -> DeriveResult<PartsAcc<Ctor>> {
264    // Fold into an accumulator of the strategy types and the expressions
265    // that produces the strategy. We then just return that accumulator
266    // and let the caller of this function determine what to do with it.
267    let acc = PartsAcc::new(fields.len());
268    fields.into_iter().try_fold(acc, |mut acc, field| {
269        let attrs = attr::parse_attributes(ctx, &field.attrs)?;
270
271        // Deny attributes that are only for enum variants:
272        error::if_enum_attrs_present(ctx, &attrs, item);
273
274        let span = field.span();
275        let ty = field.ty;
276
277        let strat = pair_filter(
278            attrs.filter,
279            ty.clone(),
280            match attrs.params {
281                // Parameters were not set on the field:
282                ParamsMode::Passthrough => match attrs.strategy {
283                    // Specific strategy - use the given expr and erase the type:
284                    StratMode::Strategy(strat) => pair_existential(ty, strat),
285                    // Specific value - use the given expr:
286                    StratMode::Value(value) => pair_value(ty, value),
287                    // Specific regex - dispatch to `_regex` function:
288                    StratMode::Regex(regex) => pair_regex(ty, regex),
289                    // Use Arbitrary for the given type and mark the type as used:
290                    StratMode::Arbitrary => {
291                        ty.mark_uses(ut);
292
293                        // We use the Parameters type of the field's type.
294                        let pref = acc.add_param(arbitrary_param(&ty));
295                        pair_any_with(ty, pref, span)
296                    }
297                },
298                // no_params set on the field:
299                ParamsMode::Default => {
300                    product_handle_default_params(ut, ty, span, attrs.strategy)
301                }
302                // params(<type>) set on the field:
303                ParamsMode::Specified(params_ty) =>
304                // We need to extract the param as the binding `params`:
305                {
306                    extract_nparam(
307                        &mut acc,
308                        params_ty,
309                        match attrs.strategy {
310                            // Specific strategy - use the given expr and erase the type:
311                            StratMode::Strategy(strat) => {
312                                pair_existential(ty, strat)
313                            }
314                            // Specific value - use the given expr in a closure and erase:
315                            StratMode::Value(value) => {
316                                pair_value_exist(ty, value)
317                            }
318                            // Logic error by user; Pointless to specify params and
319                            // regex because the params can never be used in the regex.
320                            StratMode::Regex(regex) => {
321                                error::cant_set_param_and_regex(ctx, item);
322                                pair_regex(ty, regex)
323                            }
324                            // Logic error by user.
325                            // Pointless to specify params and not the strategy. Bail!
326                            StratMode::Arbitrary => {
327                                error::cant_set_param_but_not_strat(
328                                    ctx, &ty, item,
329                                )?
330                            }
331                        },
332                    )
333                }
334            },
335        );
336        Ok(acc.add_strat(strat))
337    })
338}
339
340/// Wrap the given constructor with a let binding
341/// moving `param_<x>` into `params`.
342fn extract_nparam<C>(
343    acc: &mut PartsAcc<C>,
344    params_ty: Type,
345    (strat, ctor): StratPair,
346) -> StratPair {
347    (
348        strat,
349        extract_api(ctor, FromReg::Num(acc.add_param(params_ty))),
350    )
351}
352
353//==============================================================================
354// Enum
355//==============================================================================
356
357/// Entry point for deriving `Arbitrary` for `enum`s.
358fn derive_enum(
359    ctx: Ctx,
360    mut ast: DeriveData<Vec<Variant>>,
361) -> DeriveResult<Impl> {
362    // An enum can't be skipped, ensure it hasn't been:
363    error::if_skip_present(ctx, &ast.attrs, error::ENUM);
364
365    // We don't allow a strategy on the enum directly:
366    error::if_strategy_present(ctx, &ast.attrs, error::ENUM);
367
368    // We don't allow weight on enums directly:
369    error::if_weight_present(ctx, &ast.attrs, error::ENUM);
370
371    // Bail if there are no variants:
372    if ast.body.is_empty() {
373        error::uninhabited_enum_with_no_variants(ctx)?;
374    }
375
376    // Bail if all variants are uninhabited:
377    if (&*ast.body).is_uninhabited() {
378        error::uninhabited_enum_variants_uninhabited(ctx)?;
379    }
380
381    // The complexity of the logic depends mostly now on whether
382    // parameters were set directly on the type or not.
383    let parts = if let Some(sty) = ast.attrs.params.into_option() {
384        // The logic is much simpler in this branch.
385        derive_enum_has_params(ctx, &mut ast.tracker, &ast.ident, ast.body, sty)
386    } else {
387        // And considerably more complex here.
388        derive_enum_no_params(ctx, &mut ast.tracker, &ast.ident, ast.body)
389    }?;
390
391    let parts = add_top_filter(ast.attrs.filter, parts);
392
393    // We're done!
394    Ok(Impl::new(ast.ident, ast.tracker, parts))
395}
396
397/// Deriving for a enum on which `params` or `no_params` was NOT set directly.
398fn derive_enum_no_params(
399    ctx: Ctx,
400    ut: &mut UseTracker,
401    _self: &Ident,
402    variants: Vec<Variant>,
403) -> DeriveResult<ImplParts> {
404    // Initialize the accumulator:
405    let mut acc = PartsAcc::new(variants.len());
406
407    // Fold into the accumulator the strategies for each variant:
408    for variant in variants {
409        if let Some((weight, ident, fields, attrs)) =
410            keep_inhabited_variant(ctx, _self, variant)?
411        {
412            let path = parse_quote!( #_self::#ident );
413            let (strat, ctor) = if fields.is_empty() {
414                // Unit variant:
415                pair_unit_variant(ctx, &attrs, path)
416            } else {
417                // Not a unit variant:
418                derive_variant_with_fields(
419                    ctx, ut, path, attrs, fields, &mut acc,
420                )?
421            };
422            acc = acc.add_strat((strat, (weight, ctor)));
423        }
424    }
425
426    ensure_union_has_strategies(ctx, &acc.strats);
427
428    // Package the strategies into a union.
429    Ok(acc.finish(ctx))
430}
431
432/// Ensure that there's at least one generatable variant for a union.
433fn ensure_union_has_strategies<C>(ctx: Ctx, strats: &StratAcc<C>) {
434    if strats.is_empty() {
435        // We didn't accumulate any strategies,
436        // so we can't construct any variant.
437        error::uninhabited_enum_because_of_skipped_variants(ctx);
438    }
439}
440
441/// Derive for a variant which has fields and where the
442/// variant or its fields may specify `params` or `no_params`.
443fn derive_variant_with_fields<C>(
444    ctx: Ctx,
445    ut: &mut UseTracker,
446    v_path: Path,
447    attrs: ParsedAttributes,
448    fields: Vec<Field>,
449    acc: &mut PartsAcc<C>,
450) -> DeriveResult<StratPair> {
451    let filter = attrs.filter.clone();
452
453    let pair = match attrs.params {
454        // Parameters were not set on the variant:
455        ParamsMode::Passthrough => match attrs.strategy {
456            // Specific strategy - use the given expr and erase the type:
457            StratMode::Strategy(strat) => {
458                deny_all_attrs_on_fields(ctx, fields)?;
459                pair_existential_self(strat)
460            }
461            // Specific value - use the given expr:
462            StratMode::Value(value) => {
463                deny_all_attrs_on_fields(ctx, fields)?;
464                pair_value_self(value)
465            }
466            StratMode::Regex(regex) => {
467                deny_all_attrs_on_fields(ctx, fields)?;
468                pair_regex_self(regex)
469            }
470            // No explicit strategy, use strategies for variant fields instead:
471            StratMode::Arbitrary => {
472                variant_no_explicit_strategy(ctx, ut, v_path, fields, acc)?
473            }
474        },
475        // no_params set on the variant:
476        ParamsMode::Default => {
477            variant_handle_default_params(ctx, ut, v_path, attrs, fields)?
478        }
479        // params(<type>) set on the variant:
480        ParamsMode::Specified(params_ty) => extract_nparam(
481            acc,
482            params_ty,
483            match attrs.strategy {
484                // Specific strategy - use the given expr and erase the type:
485                StratMode::Strategy(strat) => {
486                    deny_all_attrs_on_fields(ctx, fields)?;
487                    pair_existential_self(strat)
488                }
489                // Specific value - use the given expr in a closure and erase:
490                StratMode::Value(value) => {
491                    deny_all_attrs_on_fields(ctx, fields)?;
492                    pair_value_exist_self(value)
493                }
494                // Logic error by user; Pointless to specify params and regex
495                // because the params can never be used in the regex.
496                StratMode::Regex(regex) => {
497                    error::cant_set_param_and_regex(ctx, error::ENUM_VARIANT);
498                    deny_all_attrs_on_fields(ctx, fields)?;
499                    pair_regex_self(regex)
500                }
501                // Logic error by user. Pointless to specify params and not
502                // the strategy. Bail!
503                StratMode::Arbitrary => {
504                    let ty = self_ty();
505                    error::cant_set_param_but_not_strat(
506                        ctx,
507                        &ty,
508                        error::ENUM_VARIANT,
509                    )?
510                }
511            },
512        ),
513    };
514    let pair = add_filter_self(filter, pair);
515    Ok(pair)
516}
517
518/// Derive for a variant on which params were not set and on which no explicit
519/// strategy was set (or where it doesn't make sense...) and which has fields.
520fn variant_no_explicit_strategy<C>(
521    ctx: Ctx,
522    ut: &mut UseTracker,
523    v_path: Path,
524    fields: Vec<Field>,
525    acc: &mut PartsAcc<C>,
526) -> DeriveResult<StratPair> {
527    // Compute parts for the inner product:
528    let closure = map_closure(v_path, &fields);
529    let fields_acc =
530        derive_product_no_params(ctx, ut, fields, error::ENUM_VARIANT_FIELD)?;
531    let (params, count) = fields_acc.params.consume();
532    let (strat, ctor) = fields_acc.strats.finish(closure);
533
534    // Add params types from inner derive as a single type
535    // in the outer params types.
536    let params_ty = params.into();
537    Ok((
538        strat,
539        if is_unit_type(&params_ty) {
540            ctor
541        } else {
542            let pref = acc.add_param(params_ty);
543            extract_all(ctor, count, FromReg::Num(pref))
544        },
545    ))
546}
547
548/// Determine strategy using "Default" semantics for a variant.
549fn variant_handle_default_params(
550    ctx: Ctx,
551    ut: &mut UseTracker,
552    v_path: Path,
553    attrs: ParsedAttributes,
554    fields: Vec<Field>,
555) -> DeriveResult<StratPair> {
556    let pair = match attrs.strategy {
557        // Specific strategy - use the given expr and erase the type:
558        StratMode::Strategy(strat) => {
559            deny_all_attrs_on_fields(ctx, fields)?;
560            pair_existential_self(strat)
561        }
562        // Specific value - use the given expr:
563        StratMode::Value(value) => {
564            deny_all_attrs_on_fields(ctx, fields)?;
565            pair_value_self(value)
566        }
567        StratMode::Regex(regex) => {
568            deny_all_attrs_on_fields(ctx, fields)?;
569            pair_regex_self(regex)
570        }
571        // Use Arbitrary for the factors (fields) of variant:
572        StratMode::Arbitrary =>
573        // Fields are not allowed to specify params.
574        {
575            derive_product_has_params(
576                ctx,
577                ut,
578                error::ENUM_VARIANT_FIELD,
579                map_closure(v_path, &fields),
580                fields,
581            )?
582        }
583    };
584
585    Ok(pair)
586}
587
588/// Ensures that there are no proptest attributes on any of the fields.
589fn deny_all_attrs_on_fields(ctx: Ctx, fields: Vec<Field>) -> DeriveResult<()> {
590    fields.into_iter().try_for_each(|field| {
591        let f_attr = attr::parse_attributes(ctx, &field.attrs)?;
592        error::if_anything_specified(ctx, &f_attr, error::ENUM_VARIANT_FIELD);
593        Ok(())
594    })
595}
596
597/// Derive for a variant which has fields and where the
598/// variant or its fields may NOT specify `params` or `no_params`.
599fn derive_enum_has_params(
600    ctx: Ctx,
601    ut: &mut UseTracker,
602    _self: &Ident,
603    variants: Vec<Variant>,
604    sty: Option<Type>,
605) -> DeriveResult<ImplParts> {
606    // Initialize the accumulator:
607    let mut acc = StratAcc::new(variants.len());
608
609    // Fold into the accumulator the strategies for each variant:
610    for variant in variants {
611        let parts = keep_inhabited_variant(ctx, _self, variant)?;
612        if let Some((weight, ident, fields, attrs)) = parts {
613            let path = parse_quote!( #_self::#ident );
614            let (strat, ctor) = if fields.is_empty() {
615                // Unit variant:
616                pair_unit_variant(ctx, &attrs, path)
617            } else {
618                // Not a unit variant:
619                let filter = attrs.filter.clone();
620                add_filter_self(
621                    filter,
622                    variant_handle_default_params(
623                        ctx, ut, path, attrs, fields,
624                    )?,
625                )
626            };
627            acc = acc.add((strat, (weight, ctor)));
628        }
629    }
630
631    ensure_union_has_strategies(ctx, &acc);
632
633    Ok(add_top_params(sty, acc.finish(ctx)))
634}
635
636/// Filters out uninhabited and variants that we've been ordered to skip.
637fn keep_inhabited_variant(
638    ctx: Ctx,
639    _self: &Ident,
640    variant: Variant,
641) -> DeriveResult<Option<(u32, Ident, Vec<Field>, ParsedAttributes)>> {
642    let attrs = attr::parse_attributes(ctx, &variant.attrs)?;
643    let fields = fields_to_vec(variant.fields);
644
645    if attrs.skip {
646        // We've been ordered to skip this variant!
647        // Check that all other attributes are not set.
648        ensure_has_only_skip_attr(ctx, &attrs, error::ENUM_VARIANT);
649        fields.into_iter().try_for_each(|field| {
650            let f_attrs = attr::parse_attributes(ctx, &field.attrs)?;
651            error::if_skip_present(ctx, &f_attrs, error::ENUM_VARIANT_FIELD);
652            ensure_has_only_skip_attr(ctx, &f_attrs, error::ENUM_VARIANT_FIELD);
653            Ok(())
654        })?;
655
656        return Ok(None);
657    }
658
659    // If the variant is uninhabited, we can't generate it, so skip it.
660    if (&*fields).is_uninhabited() {
661        return Ok(None);
662    }
663
664    // Compute the weight:
665    let weight = attrs.weight.unwrap_or(1);
666
667    Ok(Some((weight, variant.ident, fields, attrs)))
668}
669
670/// Ensures that no other attributes than skip are present.
671fn ensure_has_only_skip_attr(ctx: Ctx, attrs: &ParsedAttributes, item: &str) {
672    if attrs.params.is_set() {
673        error::skipped_variant_has_param(ctx, item);
674    }
675
676    if attrs.strategy.is_set() {
677        error::skipped_variant_has_strat(ctx, item);
678    }
679
680    if attrs.weight.is_some() {
681        error::skipped_variant_has_weight(ctx, item);
682    }
683
684    if !attrs.filter.is_empty() {
685        error::skipped_variant_has_filter(ctx, item);
686    }
687}
688
689/// Deal with a unit variant.
690fn pair_unit_variant(
691    ctx: Ctx,
692    attrs: &ParsedAttributes,
693    v_path: Path,
694) -> StratPair {
695    error::if_present_on_unit_variant(ctx, attrs);
696    pair_unit_self(&v_path)
697}
698
699//==============================================================================
700// Combined accumulator
701//==============================================================================
702
703/// Combined accumulator for the parameters and strategies.
704struct PartsAcc<C> {
705    /// The accumulator for the parameters.
706    params: ParamAcc,
707    /// The accumulator for the strategies.
708    strats: StratAcc<C>,
709}
710
711impl<C> PartsAcc<C> {
712    /// Constructs a new accumulator with the size
713    /// passed on to the accumulator for the strategies.
714    fn new(size: usize) -> Self {
715        Self {
716            params: ParamAcc::empty(),
717            strats: StratAcc::new(size),
718        }
719    }
720
721    /// Adds a strategy to the accumulator.
722    fn add_strat(self, pair: (Strategy, C)) -> Self {
723        Self {
724            strats: self.strats.add(pair),
725            params: self.params,
726        }
727    }
728
729    /// Adds a parameter type to the accumulator and returns how many types
730    /// there were before adding.
731    fn add_param(&mut self, ty: Type) -> usize {
732        self.params.add(ty)
733    }
734}
735
736impl PartsAcc<Ctor> {
737    /// Finishes off the accumulator by returning the parts needed for
738    /// deriving. The resulting strategy is a mapping of the parts into
739    /// the `Self` type.
740    fn finish(self, closure: MapClosure) -> ImplParts {
741        let (params, count) = self.params.consume();
742        let (strat, ctor) = self.strats.finish(closure);
743        (params, strat, extract_all(ctor, count, FromReg::Top))
744    }
745}
746
747impl PartsAcc<(u32, Ctor)> {
748    /// Finishes off the accumulator by returning the parts needed for
749    /// deriving. The resultant strategy is one that randomly picks
750    /// one of the parts based on the relative weights in the `u32`.
751    fn finish(self, ctx: Ctx) -> ImplParts {
752        let (params, count) = self.params.consume();
753        let (strat, ctor) = self.strats.finish(ctx);
754        (params, strat, extract_all(ctor, count, FromReg::Top))
755    }
756}
757
758//==============================================================================
759// Param accumulator
760//==============================================================================
761
762/// Accumulator of the parameter types.
763struct ParamAcc {
764    /// The accumulated parameters types.
765    types: Params,
766}
767
768impl ParamAcc {
769    /// Returns an empty accumulator.
770    fn empty() -> Self {
771        Self {
772            types: Params::empty(),
773        }
774    }
775
776    /// Adds a type to the accumulator and returns the type count before adding.
777    fn add(&mut self, ty: Type) -> usize {
778        let var = self.types.len();
779        self.types += ty;
780        var
781    }
782
783    /// Consumes the accumulator returning the types and the count.
784    fn consume(self) -> (Params, usize) {
785        let count = self.types.len();
786        (self.types, count)
787    }
788}
789
790//==============================================================================
791// Strategy accumulator
792//==============================================================================
793
794/// Accumulator of a sequence of strategies (both type and constructor).
795struct StratAcc<C> {
796    /// The type half of the accumulator:
797    types: Vec<Strategy>,
798    /// The constructors (Rust expression that makes the strategy) half:
799    ctors: Vec<C>,
800}
801
802impl<C> StratAcc<C> {
803    /// Construct the given accumulator with
804    /// initial capacity according to `size`.
805    fn new(size: usize) -> Self {
806        Self {
807            types: Vec::with_capacity(size),
808            ctors: Vec::with_capacity(size),
809        }
810    }
811
812    /// Add the given type and constructor pair to
813    /// the accumulator which is moved and returned.
814    fn add(mut self, (strat, ctor): (Strategy, C)) -> Self {
815        self.types.push(strat);
816        self.ctors.push(ctor);
817        self
818    }
819
820    /// Consume the accumulator returning the:
821    /// + sequence of strategies
822    /// + sequence of constructors
823    fn consume(self) -> (Vec<Strategy>, Vec<C>) {
824        (self.types, self.ctors)
825    }
826
827    /// Returns `true` iff nothing has been accumulated yet.
828    fn is_empty(&self) -> bool {
829        self.types.is_empty()
830    }
831}
832
833impl StratAcc<Ctor> {
834    /// Finishes off the accumulator by returning
835    /// a `.prop_map(<closure>)` of the strategies.
836    fn finish(self, closure: MapClosure) -> StratPair {
837        pair_map(self.consume(), closure)
838    }
839}
840
841impl StratAcc<(u32, Ctor)> {
842    /// Finishes off the accumulator by returning a union of the
843    /// strategies where the resultant strategy randomly picks
844    /// one of the summands based on the relative weights provided.
845    fn finish(self, ctx: Ctx) -> StratPair {
846        // Check that the weight sum <= u32::MAX
847        if self
848            .ctors
849            .iter()
850            .map(|&(w, _)| w)
851            .try_fold(0u32, |acc, w| acc.checked_add(w))
852            .is_none()
853        {
854            error::weight_overflowing(ctx)
855        }
856
857        pair_oneof(self.consume())
858    }
859}