proptest_derive/
ast.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//! High level IR and abstract syntax tree (AST) of impls.
10//!
11//! We compile to this AST and then linearise that to Rust code.
12
13use std::ops::{Add, AddAssign};
14
15use proc_macro2::{Span, TokenStream};
16use quote::{ToTokens, TokenStreamExt};
17use syn::spanned::Spanned;
18
19use crate::error::{Ctx, DeriveResult};
20use crate::use_tracking::UseTracker;
21use crate::util::self_ty;
22
23//==============================================================================
24// Config
25//==============================================================================
26
27/// The `MAX - 1` number of strategies that `TupleUnion` supports.
28/// Increase this if the behaviour is changed in `proptest`.
29/// Keeping this lower than what `proptest` supports will also work
30/// but for optimality this should follow what `proptest` supports.
31#[cfg(not(feature = "boxed_union"))]
32const UNION_CHUNK_SIZE: usize = 9;
33
34/// The `MAX - 1` tuple length `Arbitrary` is implemented for. After this number,
35/// tuples are expanded as nested tuples of up to `MAX` elements. The value should
36/// be kept in sync with the largest impl in `proptest/src/arbitrary/tuples.rs`.
37const NESTED_TUPLE_CHUNK_SIZE: usize = 9;
38
39/// The name of the top parameter variable name given in `arbitrary_with`.
40/// Changing this is not a breaking change because a user is expected not
41/// to rely on this (and the user shouldn't be able to..).
42const TOP_PARAM_NAME: &str = "_top";
43
44/// The name of the variable name used for user facing parameter types
45/// specified in a `#[proptest(params = "<type>")]` attribute.
46///
47/// Changing the value of this constant constitutes a breaking change!
48const API_PARAM_NAME: &str = "params";
49
50//==============================================================================
51// AST Root
52//==============================================================================
53
54/// Top level AST and everything required to implement `Arbitrary` for any
55/// given type. Linearizing this AST gives you the impl wrt. Rust code.
56pub struct Impl {
57    /// Name of the type.
58    typ: syn::Ident,
59    /// Tracker for uses of Arbitrary trait for a generic type.
60    tracker: UseTracker,
61    /// The three main parts, see description of `ImplParts` for details.
62    parts: ImplParts,
63}
64
65/// The three main parts to deriving `Arbitrary` for a type.
66/// That is: the associated items `Parameters` (`Params`),
67/// `Strategy` (`Strategy`) as well as the construction of the
68/// strategy itself (`Ctor`).
69pub type ImplParts = (Params, Strategy, Ctor);
70
71impl Impl {
72    /// Constructs a new `Impl` from the parts as described on the type.
73    pub fn new(typ: syn::Ident, tracker: UseTracker, parts: ImplParts) -> Self {
74        Self {
75            typ,
76            tracker,
77            parts,
78        }
79    }
80
81    /// Linearises the impl into a sequence of tokens.
82    /// This produces the actual Rust code for the impl.
83    pub fn into_tokens(self, ctx: Ctx) -> DeriveResult<TokenStream> {
84        let Impl {
85            typ,
86            mut tracker,
87            parts: (params, strategy, ctor),
88        } = self;
89
90        /// A `Debug` bound on a type variable.
91        fn debug_bound() -> syn::TypeParamBound {
92            parse_quote!(::std::fmt::Debug)
93        }
94
95        /// An `Arbitrary` bound on a type variable.
96        fn arbitrary_bound() -> syn::TypeParamBound {
97            parse_quote!(_proptest::arbitrary::Arbitrary)
98        }
99
100        // Add bounds and get generics for the impl.
101        tracker.add_bounds(ctx, &arbitrary_bound(), Some(debug_bound()))?;
102        let generics = tracker.consume();
103        let (impl_generics, ty_generics, where_clause) =
104            generics.split_for_impl();
105
106        let _top = call_site_ident(TOP_PARAM_NAME);
107
108        // Linearise everything. We're done after this.
109        //
110        // NOTE: The clippy::arc_with_non_send_sync lint is disabled here because the strategies
111        // generated are often not Send or Sync, such as BoxedStrategy.
112        //
113        // The double-curly-braces are not strictly required, but allow the expression to be
114        // annotated with an attribute.
115        let q = quote! {
116            #[allow(non_local_definitions)]
117            #[allow(non_upper_case_globals)]
118            #[allow(clippy::arc_with_non_send_sync)]
119            const _: () = {
120            use proptest as _proptest;
121
122            impl #impl_generics _proptest::arbitrary::Arbitrary
123            for #typ #ty_generics #where_clause {
124                type Parameters = #params;
125
126                type Strategy = #strategy;
127
128                fn arbitrary_with(#_top: Self::Parameters) -> Self::Strategy {
129                    #ctor
130                }
131            }
132
133            };
134        };
135
136        Ok(q)
137    }
138}
139
140//==============================================================================
141// Smart construcors, StratPair
142//==============================================================================
143
144/// A pair of `Strategy` and `Ctor`. These always come in pairs.
145pub type StratPair = (Strategy, Ctor);
146
147/// The type and constructor for `any::<Type>()`.
148pub fn pair_any(ty: syn::Type, span: Span) -> StratPair {
149    let q = Ctor::Arbitrary(ty.clone(), None, span);
150    (Strategy::Arbitrary(ty, span), q)
151}
152
153/// The type and constructor for `any_with::<Type>(parameters)`.
154pub fn pair_any_with(ty: syn::Type, var: usize, span: Span) -> StratPair {
155    let q = Ctor::Arbitrary(ty.clone(), Some(var), span);
156    (Strategy::Arbitrary(ty, span), q)
157}
158
159/// The type and constructor for a specific strategy value constructed by the
160/// given expression. Currently, the type is erased and a `BoxedStrategy<Type>`
161/// is given back instead.
162///
163/// This is a temporary restriction. Once `impl Trait` is stabilized,
164/// the boxing and dynamic dispatch can be replaced with a statically
165/// dispatched anonymous type instead.
166pub fn pair_existential(ty: syn::Type, strat: syn::Expr) -> StratPair {
167    (Strategy::Existential(ty), Ctor::Existential(strat))
168}
169
170/// The type and constructor for a strategy that always returns the value
171/// provided in the expression `val`.
172/// This is statically dispatched since no erasure is needed or used.
173pub fn pair_value(ty: syn::Type, val: syn::Expr) -> StratPair {
174    (Strategy::Value(ty), Ctor::Value(val))
175}
176
177/// Same as `pair_existential` for the `Self` type.
178pub fn pair_existential_self(strat: syn::Expr) -> StratPair {
179    pair_existential(self_ty(), strat)
180}
181
182/// Same as `pair_value` for the `Self` type.
183pub fn pair_value_self(val: syn::Expr) -> StratPair {
184    pair_value(self_ty(), val)
185}
186
187/// Erased strategy for a fixed value.
188pub fn pair_value_exist(ty: syn::Type, strat: syn::Expr) -> StratPair {
189    (Strategy::Existential(ty), Ctor::ValueExistential(strat))
190}
191
192/// Erased strategy for a fixed value.
193pub fn pair_value_exist_self(strat: syn::Expr) -> StratPair {
194    pair_value_exist(self_ty(), strat)
195}
196
197/// Same as `pair_value` but for a unit variant or unit struct.
198pub fn pair_unit_self(path: &syn::Path) -> StratPair {
199    pair_value_self(parse_quote!( #path {} ))
200}
201
202/// The type and constructor for `#[proptest(regex(..))]`.
203pub fn pair_regex(ty: syn::Type, regex: syn::Expr) -> StratPair {
204    (Strategy::Regex(ty.clone()), Ctor::Regex(ty, regex))
205}
206
207/// Same as `pair_regex` for the `Self` type.
208pub fn pair_regex_self(regex: syn::Expr) -> StratPair {
209    pair_regex(self_ty(), regex)
210}
211
212/// The type and constructor for .prop_map:ing a set of strategies
213/// into the type we are implementing for. The closure for the
214/// `.prop_map(<closure>)` must also be given.
215pub fn pair_map(
216    (strats, ctors): (Vec<Strategy>, Vec<Ctor>),
217    closure: MapClosure,
218) -> StratPair {
219    (
220        Strategy::Map(strats.into()),
221        Ctor::Map(ctors.into(), closure),
222    )
223}
224
225/// The type and constructor for a union of strategies which produces a new
226/// strategy that used the given strategies with probabilities based on the
227/// assigned relative weights for each strategy.
228pub fn pair_oneof(
229    (strats, ctors): (Vec<Strategy>, Vec<(u32, Ctor)>),
230) -> StratPair {
231    (Strategy::Union(strats.into()), Ctor::Union(ctors.into()))
232}
233
234/// Potentially apply a filter to a strategy type and its constructor.
235pub fn pair_filter(
236    filter: Vec<syn::Expr>,
237    ty: syn::Type,
238    pair: StratPair,
239) -> StratPair {
240    filter.into_iter().fold(pair, |(strat, ctor), filter| {
241        (
242            Strategy::Filter(Box::new(strat), ty.clone()),
243            Ctor::Filter(Box::new(ctor), filter),
244        )
245    })
246}
247
248//==============================================================================
249// Parameters
250//==============================================================================
251
252/// Represents the associated item of `Parameters` of an `Arbitrary` impl.
253pub struct Params(Vec<syn::Type>);
254
255impl Params {
256    /// Construct an `empty` list of parameters.
257    /// This is equivalent to the unit type `()`.
258    pub fn empty() -> Self {
259        Params(Vec::new())
260    }
261
262    /// Computes and returns the number of parameter types.
263    pub fn len(&self) -> usize {
264        self.0.len()
265    }
266}
267
268impl From<Params> for syn::Type {
269    fn from(x: Params) -> Self {
270        let tys = x.0;
271        parse_quote!( (#(#tys),*) )
272    }
273}
274
275impl Add<syn::Type> for Params {
276    type Output = Params;
277
278    fn add(mut self, rhs: syn::Type) -> Self::Output {
279        self.0.push(rhs);
280        self
281    }
282}
283
284impl AddAssign<syn::Type> for Params {
285    fn add_assign(&mut self, rhs: syn::Type) {
286        self.0.push(rhs);
287    }
288}
289
290impl ToTokens for Params {
291    fn to_tokens(&self, tokens: &mut TokenStream) {
292        NestedTuple(self.0.as_slice()).to_tokens(tokens)
293    }
294}
295
296/// Returns for a given type `ty` the associated item `Parameters` of the
297/// type's `Arbitrary` implementation.
298pub fn arbitrary_param(ty: &syn::Type) -> syn::Type {
299    parse_quote!(<#ty as _proptest::arbitrary::Arbitrary>::Parameters)
300}
301
302//==============================================================================
303// Strategy
304//==============================================================================
305
306/// The type of a given `Strategy`.
307pub enum Strategy {
308    /// Assuming the metavariable `$ty` for a given type, this models the
309    /// strategy type `<$ty as Arbitrary>::Strategy`.
310    Arbitrary(syn::Type, Span),
311    /// This models <$ty as StrategyFromRegex>::Strategy.
312    Regex(syn::Type),
313    /// Assuming the metavariable `$ty` for a given type, this models the
314    /// strategy type `BoxedStrategy<$ty>`, i.e: an existentially typed strategy.
315    ///
316    /// The dynamic dispatch used here is an implementation detail that may be
317    /// changed. Such a change does not count as a breakage semver wise.
318    Existential(syn::Type),
319    /// Assuming the metavariable `$ty` for a given type, this models a
320    /// non-shrinking strategy that simply always returns a value of the
321    /// given type.
322    Value(syn::Type),
323    /// Assuming a sequence of strategies, this models a mapping from that
324    /// sequence to `Self`.
325    Map(Box<[Strategy]>),
326    /// Assuming a sequence of relative-weighted strategies, this models a
327    /// weighted choice of those strategies. The resultant strategy will in
328    /// other words randomly pick one strategy with probabilities based on the
329    /// specified weights.
330    Union(Box<[Strategy]>),
331    /// A filtered strategy with `.prop_filter`.
332    Filter(Box<Strategy>, syn::Type),
333}
334
335macro_rules! quote_append {
336    ($tokens: expr, $($quasi: tt)*) => {
337        $tokens.append_all(quote!($($quasi)*))
338    };
339}
340
341impl Strategy {
342    fn types(&self) -> Vec<syn::Type> {
343        use self::Strategy::*;
344        match self {
345            Arbitrary(ty, _) => vec![ty.clone()],
346            Regex(ty) => vec![ty.clone()],
347            Existential(ty) => vec![ty.clone()],
348            Value(ty) => vec![ty.clone()],
349            Map(strats) => strats.iter().flat_map(|s| s.types()).collect(),
350            Union(strats) => strats.iter().flat_map(|s| s.types()).collect(),
351            Filter(_, ty) => vec![ty.clone()],
352        }
353    }
354}
355
356impl ToTokens for Strategy {
357    fn to_tokens(&self, tokens: &mut TokenStream) {
358        // The logic of each of these are pretty straight forward save for
359        // union which is described separately.
360        use self::Strategy::*;
361        match self {
362            Arbitrary(ty, span) => tokens.append_all(quote_spanned!(*span=>
363                <#ty as _proptest::arbitrary::Arbitrary>::Strategy
364            )),
365            Regex(ty) => quote_append!(tokens,
366                <#ty as _proptest::string::StrategyFromRegex>::Strategy
367            ),
368            Existential(ty) => quote_append!(tokens,
369                _proptest::strategy::BoxedStrategy<#ty>
370            ),
371            Value(ty) => quote_append!(tokens, fn() -> #ty ),
372            Map(strats) => {
373                let types = self.types();
374                let field_tys = NestedTuple(&types);
375                let strats = NestedTuple(&strats);
376                quote_append!(tokens,
377                    _proptest::strategy::Map< ( #strats ),
378                        fn( #field_tys ) -> Self
379                    >
380                )
381            }
382            #[cfg(not(feature = "boxed_union"))]
383            Union(strats) => union_strat_to_tokens(tokens, strats),
384            #[cfg(feature = "boxed_union")]
385            Union(strats) => union_strat_to_tokens_boxed(tokens, strats),
386            Filter(strat, ty) => quote_append!(tokens,
387                _proptest::strategy::Filter<#strat, fn(&#ty) -> bool>
388            ),
389        }
390    }
391}
392
393//==============================================================================
394// Constructor
395//==============================================================================
396
397/// The right hand side (RHS) of a let binding of parameters.
398pub enum FromReg {
399    /// Denotes a move from the top parameter given in the arguments of
400    /// `arbitrary_with`.
401    Top,
402    /// Denotes a move from a variable `params_<x>` where `<x>` is the given
403    /// number.
404    Num(usize),
405}
406
407/// The left hand side (LHS) of a let binding of parameters.
408pub enum ToReg {
409    /// Denotes a move and declaration to a sequence of variables from
410    /// `params_0` to `params_x`.
411    Range(usize),
412    /// Denotes a move and declaration of a special variable `params` that is
413    /// user facing and is ALWAYS named `params`.
414    ///
415    /// To change the name this linearises to is considered a breaking change
416    /// wrt. semver.
417    API,
418}
419
420/// Models an expression that generates a proptest `Strategy`.
421pub enum Ctor {
422    /// A strategy generated by using the `Arbitrary` impl for the given `Ty´.
423    /// If `Some(idx)` is specified, then a parameter at `params_<idx>` is used
424    /// and provided to `any_with::<Ty>(params_<idx>)`.
425    Arbitrary(syn::Type, Option<usize>, Span),
426    /// A strategy that is generated by a mapping a regex in the form of a
427    /// string slice to the actual regex.
428    Regex(syn::Type, syn::Expr),
429    /// An exact strategy value given by the expression.
430    Existential(syn::Expr),
431    /// A strategy that always produces the given expression.
432    Value(syn::Expr),
433    /// A strategy that always produces the given expression but which is erased.
434    ValueExistential(syn::Expr),
435    /// A strategy that maps from a sequence of strategies into `Self`.
436    Map(Box<[Ctor]>, MapClosure),
437    /// A strategy that randomly selects one of the given relative-weighted
438    /// strategies.
439    Union(Box<[(u32, Ctor)]>),
440    /// A let binding that moves to and declares the `ToReg` from the `FromReg`
441    /// as well as the strategy that uses the `ToReg`.
442    Extract(Box<Ctor>, ToReg, FromReg),
443    /// A filtered strategy with `.prop_filter`.
444    Filter(Box<Ctor>, syn::Expr),
445}
446
447/// Wraps the given strategy producing expression with a move into
448/// `params_<to>` from `FromReg`. This is used when the given `c` expects
449/// `params_<to>` to be there.
450pub fn extract_all(c: Ctor, to: usize, from: FromReg) -> Ctor {
451    extract(c, ToReg::Range(to), from)
452}
453
454/// Wraps the given strategy producing expression with a move into `params`
455/// (literally named like that) from `FromReg`. This is used when the given
456/// `c` expects `params` to be there.
457pub fn extract_api(c: Ctor, from: FromReg) -> Ctor {
458    extract(c, ToReg::API, from)
459}
460
461impl ToTokens for FromReg {
462    fn to_tokens(&self, tokens: &mut TokenStream) {
463        match self {
464            FromReg::Top => call_site_ident(TOP_PARAM_NAME).to_tokens(tokens),
465            FromReg::Num(reg) => param(*reg).to_tokens(tokens),
466        }
467    }
468}
469
470impl ToTokens for ToReg {
471    fn to_tokens(&self, tokens: &mut TokenStream) {
472        match *self {
473            ToReg::Range(to) if to == 1 => param(0).to_tokens(tokens),
474            ToReg::Range(to) => {
475                let params: Vec<_> = (0..to).map(param).collect();
476                NestedTuple(&params).to_tokens(tokens)
477            }
478            ToReg::API => call_site_ident(API_PARAM_NAME).to_tokens(tokens),
479        }
480    }
481}
482
483impl ToTokens for Ctor {
484    fn to_tokens(&self, tokens: &mut TokenStream) {
485        // The logic of each of these are pretty straight forward save for
486        // union which is described separately.
487        use self::Ctor::*;
488        match self {
489            Filter(ctor, filter) => quote_append!(tokens,
490                _proptest::strategy::Strategy::prop_filter(
491                    #ctor, stringify!(#filter), #filter)
492            ),
493            Extract(ctor, to, from) => quote_append!(tokens, {
494                let #to = #from; #ctor
495            }),
496            Arbitrary(ty, fv, span) => {
497                tokens.append_all(if let Some(fv) = fv {
498                    let args = param(*fv);
499                    quote_spanned!(*span=>
500                        _proptest::arbitrary::any_with::<#ty>(#args)
501                    )
502                } else {
503                    quote_spanned!(*span=>
504                        _proptest::arbitrary::any::<#ty>()
505                    )
506                })
507            }
508            Regex(ty, regex) => quote_append!(tokens,
509                <#ty as _proptest::string::StrategyFromRegex>::from_regex(#regex)
510            ),
511            Existential(expr) => quote_append!(tokens,
512                _proptest::strategy::Strategy::boxed( #expr ) ),
513            Value(expr) => quote_append!(tokens, (|| #expr) as fn() -> _),
514            ValueExistential(expr) => quote_append!(tokens,
515                _proptest::strategy::Strategy::boxed(
516                    _proptest::strategy::LazyJust::new(move || #expr)
517                )
518            ),
519            Map(ctors, closure) => map_ctor_to_tokens(tokens, &ctors, closure),
520            #[cfg(not(feature = "boxed_union"))]
521            Union(ctors) => union_ctor_to_tokens(tokens, ctors),
522            #[cfg(feature = "boxed_union")]
523            Union(ctors) => union_ctor_to_tokens_boxed(tokens, ctors),
524        }
525    }
526}
527
528struct NestedTuple<'a, T>(&'a [T]);
529
530impl<'a, T: ToTokens> ToTokens for NestedTuple<'a, T> {
531    fn to_tokens(&self, tokens: &mut TokenStream) {
532        let NestedTuple(elems) = self;
533        if elems.is_empty() {
534            quote_append!(tokens, ());
535        } else if let [x] = elems {
536            x.to_tokens(tokens);
537        } else {
538            let chunks = elems.chunks(NESTED_TUPLE_CHUNK_SIZE);
539            Recurse(&chunks).to_tokens(tokens);
540        }
541
542        struct Recurse<'a, T: ToTokens>(&'a ::std::slice::Chunks<'a, T>);
543
544        impl<'a, T: ToTokens> ToTokens for Recurse<'a, T> {
545            fn to_tokens(&self, tokens: &mut TokenStream) {
546                let mut chunks = self.0.clone();
547                if let Some(head) = chunks.next() {
548                    if let [c] = head {
549                        // Only one element left - no need to nest.
550                        quote_append!(tokens, #c);
551                    } else {
552                        let tail = Recurse(&chunks);
553                        quote_append!(tokens, (#(#head,)* #tail));
554                    }
555                }
556            }
557        }
558    }
559}
560
561fn map_ctor_to_tokens(
562    tokens: &mut TokenStream,
563    ctors: &[Ctor],
564    closure: &MapClosure,
565) {
566    let ctors = NestedTuple(ctors);
567
568    quote_append!(tokens,
569        _proptest::strategy::Strategy::prop_map(
570            #ctors,
571            #closure
572        )
573    );
574}
575
576/// Tokenizes a weighted list of `Ctor`.
577///
578/// The logic is that the output should be as linear as possible while still
579/// supporting enums with an unbounded number of variants without any boxing
580/// (erasure) or dynamic dispatch.
581///
582/// As `TupleUnion` is (currently) limited to 10 summands in the coproduct
583/// we can't just emit the entire thing linearly as this will fail on the 11:th
584/// variant.
585///
586/// A naive approach to solve might be to simply use a cons-list like so:
587///
588/// ```ignore
589/// TupleUnion::new(
590///     (w_1, s_1),
591///     (w_2 + w_3 + w_4 + w_5,
592///      TupleUnion::new(
593///         (w_2, s_2),
594///         (w_3 + w_4 + w_5,
595///          TupleUnion::new(
596///             (w_3, s_3),
597///             (w_4 + w_5,
598///              TupleUnion::new(
599///                 (w_4, s_4),
600///                 (w_5, s_5),
601///             ))
602///         ))
603///     ))
604/// )
605/// ```
606///
607/// However, we can do better by being linear for the `10 - 1` first
608/// strategies and then switch to nesting like so:
609///
610/// ```ignore
611/// (1, 2, 3, 4, 5, 6, 7, 8, 9,
612///     (10, 11, 12, 13, 14, 15, 16, 17, 18,
613///         (19, ..)))
614/// ```
615#[cfg(not(feature = "boxed_union"))]
616fn union_ctor_to_tokens(tokens: &mut TokenStream, ctors: &[(u32, Ctor)]) {
617    if ctors.is_empty() {
618        return;
619    }
620
621    if let [(_, ctor)] = ctors {
622        // This is not a union at all - user provided an enum with one variant.
623        ctor.to_tokens(tokens);
624        return;
625    }
626
627    let mut chunks = ctors.chunks(UNION_CHUNK_SIZE);
628    let chunk = chunks.next().unwrap();
629    let head = chunk.iter().map(wrap_arc);
630    let tail = Recurse(weight_sum(ctors) - weight_sum(chunk), chunks);
631
632    quote_append!(tokens,
633        _proptest::strategy::TupleUnion::new(( #(#head,)* #tail ))
634    );
635
636    struct Recurse<'a>(u32, ::std::slice::Chunks<'a, (u32, Ctor)>);
637
638    impl<'a> ToTokens for Recurse<'a> {
639        fn to_tokens(&self, tokens: &mut TokenStream) {
640            let (tweight, mut chunks) = (self.0, self.1.clone());
641
642            if let Some(chunk) = chunks.next() {
643                if let [(w, c)] = chunk {
644                    // Only one element left - no need to nest.
645                    quote_append!(tokens, (#w, ::std::sync::Arc::new(#c)) );
646                } else {
647                    let head = chunk.iter().map(wrap_arc);
648                    let tail = Recurse(tweight - weight_sum(chunk), chunks);
649                    quote_append!(tokens,
650                        (#tweight, ::std::sync::Arc::new(
651                            _proptest::strategy::TupleUnion::new((
652                                #(#head,)* #tail
653                            ))))
654                    );
655                }
656            }
657        }
658    }
659
660    fn weight_sum(ctors: &[(u32, Ctor)]) -> u32 {
661        use std::num::Wrapping;
662        let Wrapping(x) = ctors.iter().map(|&(w, _)| Wrapping(w)).sum();
663        x
664    }
665
666    fn wrap_arc(arg: &(u32, Ctor)) -> TokenStream {
667        let (w, c) = arg;
668        quote!( (#w, ::std::sync::Arc::new(#c)) )
669    }
670}
671
672/// Tokenizes a weighted list of `Strategy`.
673/// For details, see `union_ctor_to_tokens`.
674#[cfg(not(feature = "boxed_union"))]
675fn union_strat_to_tokens(tokens: &mut TokenStream, strats: &[Strategy]) {
676    if strats.is_empty() {
677        return;
678    }
679
680    if let [strat] = strats {
681        // This is not a union at all - user provided an enum with one variant.
682        strat.to_tokens(tokens);
683        return;
684    }
685
686    let mut chunks = strats.chunks(UNION_CHUNK_SIZE);
687    let chunk = chunks.next().unwrap();
688    let head = chunk.iter().map(wrap_arc);
689    let tail = Recurse(chunks);
690
691    quote_append!(tokens,
692        _proptest::strategy::TupleUnion<( #(#head,)* #tail )>
693    );
694
695    struct Recurse<'a>(::std::slice::Chunks<'a, Strategy>);
696
697    impl<'a> ToTokens for Recurse<'a> {
698        fn to_tokens(&self, tokens: &mut TokenStream) {
699            let mut chunks = self.0.clone();
700
701            if let Some(chunk) = chunks.next() {
702                if let [s] = chunk {
703                    // Only one element left - no need to nest.
704                    quote_append!(tokens, (u32, ::std::sync::Arc<#s>) );
705                } else {
706                    let head = chunk.iter().map(wrap_arc);
707                    let tail = Recurse(chunks);
708                    quote_append!(tokens,
709                        (u32,
710                         ::std::sync::Arc<_proptest::strategy::TupleUnion<(
711                             #(#head,)* #tail
712                         )>>)
713                    );
714                }
715            }
716        }
717    }
718
719    fn wrap_arc(s: &Strategy) -> TokenStream {
720        quote!( (u32, ::std::sync::Arc<#s>) )
721    }
722}
723
724/// Tokenizes a weighted list of `Ctor`.
725///
726/// This can be used instead of `union_ctor_to_tokens` to generate a boxing
727/// macro.
728#[cfg(feature = "boxed_union")]
729fn union_ctor_to_tokens_boxed(tokens: &mut TokenStream, ctors: &[(u32, Ctor)]) {
730    if ctors.is_empty() {
731        return;
732    }
733
734    if let [(_, ctor)] = ctors {
735        // This is not a union at all - user provided an enum with one variant.
736        ctor.to_tokens(tokens);
737        return;
738    }
739
740    let ctors_boxed = ctors.iter().map(wrap_boxed);
741
742    quote_append!(
743        tokens,
744        _proptest::strategy::Union::new_weighted(vec![ #(#ctors_boxed,)* ])
745    );
746
747    fn wrap_boxed(arg: &(u32, Ctor)) -> TokenStream {
748        let (w, c) = arg;
749        quote!( (#w, _proptest::strategy::Strategy::boxed(#c)) )
750    }
751}
752
753/// Tokenizes a weighted list of `Strategy`.
754/// For details, see `union_ctor_to_tokens_boxed`.
755#[cfg(feature = "boxed_union")]
756fn union_strat_to_tokens_boxed(tokens: &mut TokenStream, strats: &[Strategy]) {
757    if strats.is_empty() {
758        return;
759    }
760
761    if let [strat] = strats {
762        // This is not a union at all - user provided an enum with one variant.
763        strat.to_tokens(tokens);
764        return;
765    }
766
767    quote_append!(
768        tokens,
769        _proptest::strategy::Union<_proptest::strategy::BoxedStrategy<Self>>
770    );
771}
772
773/// Wraps a `Ctor` that expects the `to` "register" to be filled with
774/// contents of the `from` register. The correctness of this wrt. the
775/// generated Rust code has to be verified externally by checking the
776/// construction of the particular `Ctor`.
777fn extract(c: Ctor, to: ToReg, from: FromReg) -> Ctor {
778    Ctor::Extract(Box::new(c), to, from)
779}
780
781/// Construct a `FreshVar` prefixed by `param_`.
782fn param<'a>(fv: usize) -> FreshVar<'a> {
783    fresh_var("param", fv)
784}
785
786//==============================================================================
787// MapClosure
788//==============================================================================
789
790/// Constructs a `MapClosure` for the given `path` and a list of fields.
791pub fn map_closure(path: syn::Path, fs: &[syn::Field]) -> MapClosure {
792    MapClosure(path, fs.to_owned())
793}
794
795/// A `MapClosure` models the closure part inside a `.prop_map(..)` call.
796#[derive(Debug)]
797pub struct MapClosure(syn::Path, Vec<syn::Field>);
798
799impl ToTokens for MapClosure {
800    fn to_tokens(&self, tokens: &mut TokenStream) {
801        fn tmp_var<'a>(idx: usize) -> FreshVar<'a> {
802            fresh_var("tmp", idx)
803        }
804
805        let MapClosure(path, fields) = self;
806        let count = fields.len();
807        let tmps: Vec<_> = (0..count).map(tmp_var).collect();
808        let inits = fields.iter().enumerate().map(|(idx, field)| {
809            let tv = tmp_var(idx);
810            if let Some(name) = &field.ident {
811                quote_spanned!(field.span()=> #name: #tv )
812            } else {
813                let name = syn::Member::Unnamed(syn::Index::from(idx));
814                quote_spanned!(field.span()=> #name: #tv )
815            }
816        });
817        let tmps = NestedTuple(&tmps);
818        quote_append!(tokens, | #tmps | #path { #(#inits),* } );
819    }
820}
821
822//==============================================================================
823// FreshVar
824//==============================================================================
825
826/// Construct a `FreshVar` with the given `prefix` and the number it has in the
827/// count of temporaries for that prefix.
828fn fresh_var(prefix: &str, count: usize) -> FreshVar {
829    FreshVar { prefix, count }
830}
831
832/// A `FreshVar` is an internal implementation detail and models a temporary
833/// variable on the stack.
834struct FreshVar<'a> {
835    prefix: &'a str,
836    count: usize,
837}
838
839impl<'a> ToTokens for FreshVar<'a> {
840    fn to_tokens(&self, tokens: &mut TokenStream) {
841        let ident = format!("{}_{}", self.prefix, self.count);
842        call_site_ident(&ident).to_tokens(tokens)
843    }
844}
845
846fn call_site_ident(ident: &str) -> syn::Ident {
847    syn::Ident::new(ident, Span::call_site())
848}