use std::ops::{Add, AddAssign};
use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, TokenStreamExt};
use syn;
use syn::spanned::Spanned;
use crate::error::{Ctx, DeriveResult};
use crate::use_tracking::UseTracker;
use crate::util::self_ty;
const UNION_CHUNK_SIZE: usize = 9;
const NESTED_TUPLE_CHUNK_SIZE: usize = 9;
const TOP_PARAM_NAME: &str = "_top";
const API_PARAM_NAME: &str = "params";
pub struct Impl {
typ: syn::Ident,
tracker: UseTracker,
parts: ImplParts,
}
pub type ImplParts = (Params, Strategy, Ctor);
impl Impl {
pub fn new(typ: syn::Ident, tracker: UseTracker, parts: ImplParts) -> Self {
Self {
typ,
tracker,
parts,
}
}
pub fn into_tokens(self, ctx: Ctx) -> DeriveResult<TokenStream> {
let Impl {
typ,
mut tracker,
parts: (params, strategy, ctor),
} = self;
fn debug_bound() -> syn::TypeParamBound {
parse_quote!(::std::fmt::Debug)
}
fn arbitrary_bound() -> syn::TypeParamBound {
parse_quote!(_proptest::arbitrary::Arbitrary)
}
tracker.add_bounds(ctx, &arbitrary_bound(), Some(debug_bound()))?;
let generics = tracker.consume();
let (impl_generics, ty_generics, where_clause) =
generics.split_for_impl();
let _top = call_site_ident(TOP_PARAM_NAME);
let _const = call_site_ident(&format!("_IMPL_ARBITRARY_FOR_{}", typ));
let q = quote! {
#[allow(non_upper_case_globals)]
#[allow(clippy::arc_with_non_send_sync)]
const #_const: () = {
extern crate proptest as _proptest;
impl #impl_generics _proptest::arbitrary::Arbitrary
for #typ #ty_generics #where_clause {
type Parameters = #params;
type Strategy = #strategy;
fn arbitrary_with(#_top: Self::Parameters) -> Self::Strategy {
#ctor
}
}
};
};
Ok(q)
}
}
pub type StratPair = (Strategy, Ctor);
pub fn pair_any(ty: syn::Type, span: Span) -> StratPair {
let q = Ctor::Arbitrary(ty.clone(), None, span);
(Strategy::Arbitrary(ty, span), q)
}
pub fn pair_any_with(ty: syn::Type, var: usize, span: Span) -> StratPair {
let q = Ctor::Arbitrary(ty.clone(), Some(var), span);
(Strategy::Arbitrary(ty, span), q)
}
pub fn pair_existential(ty: syn::Type, strat: syn::Expr) -> StratPair {
(Strategy::Existential(ty), Ctor::Existential(strat))
}
pub fn pair_value(ty: syn::Type, val: syn::Expr) -> StratPair {
(Strategy::Value(ty), Ctor::Value(val))
}
pub fn pair_existential_self(strat: syn::Expr) -> StratPair {
pair_existential(self_ty(), strat)
}
pub fn pair_value_self(val: syn::Expr) -> StratPair {
pair_value(self_ty(), val)
}
pub fn pair_value_exist(ty: syn::Type, strat: syn::Expr) -> StratPair {
(Strategy::Existential(ty), Ctor::ValueExistential(strat))
}
pub fn pair_value_exist_self(strat: syn::Expr) -> StratPair {
pair_value_exist(self_ty(), strat)
}
pub fn pair_unit_self(path: &syn::Path) -> StratPair {
pair_value_self(parse_quote!( #path {} ))
}
pub fn pair_regex(ty: syn::Type, regex: syn::Expr) -> StratPair {
(Strategy::Regex(ty.clone()), Ctor::Regex(ty, regex))
}
pub fn pair_regex_self(regex: syn::Expr) -> StratPair {
pair_regex(self_ty(), regex)
}
pub fn pair_map(
(strats, ctors): (Vec<Strategy>, Vec<Ctor>),
closure: MapClosure,
) -> StratPair {
(
Strategy::Map(strats.into()),
Ctor::Map(ctors.into(), closure),
)
}
pub fn pair_oneof(
(strats, ctors): (Vec<Strategy>, Vec<(u32, Ctor)>),
) -> StratPair {
(Strategy::Union(strats.into()), Ctor::Union(ctors.into()))
}
pub fn pair_filter(
filter: Vec<syn::Expr>,
ty: syn::Type,
pair: StratPair,
) -> StratPair {
filter.into_iter().fold(pair, |(strat, ctor), filter| {
(
Strategy::Filter(Box::new(strat), ty.clone()),
Ctor::Filter(Box::new(ctor), filter),
)
})
}
pub struct Params(Vec<syn::Type>);
impl Params {
pub fn empty() -> Self {
Params(Vec::new())
}
pub fn len(&self) -> usize {
self.0.len()
}
}
impl From<Params> for syn::Type {
fn from(x: Params) -> Self {
let tys = x.0;
parse_quote!( (#(#tys),*) )
}
}
impl Add<syn::Type> for Params {
type Output = Params;
fn add(mut self, rhs: syn::Type) -> Self::Output {
self.0.push(rhs);
self
}
}
impl AddAssign<syn::Type> for Params {
fn add_assign(&mut self, rhs: syn::Type) {
self.0.push(rhs);
}
}
impl ToTokens for Params {
fn to_tokens(&self, tokens: &mut TokenStream) {
NestedTuple(self.0.as_slice()).to_tokens(tokens)
}
}
pub fn arbitrary_param(ty: &syn::Type) -> syn::Type {
parse_quote!(<#ty as _proptest::arbitrary::Arbitrary>::Parameters)
}
pub enum Strategy {
Arbitrary(syn::Type, Span),
Regex(syn::Type),
Existential(syn::Type),
Value(syn::Type),
Map(Box<[Strategy]>),
Union(Box<[Strategy]>),
Filter(Box<Strategy>, syn::Type),
}
macro_rules! quote_append {
($tokens: expr, $($quasi: tt)*) => {
$tokens.append_all(quote!($($quasi)*))
};
}
impl Strategy {
fn types(&self) -> Vec<syn::Type> {
use self::Strategy::*;
match self {
Arbitrary(ty, _) => vec![ty.clone()],
Regex(ty) => vec![ty.clone()],
Existential(ty) => vec![ty.clone()],
Value(ty) => vec![ty.clone()],
Map(strats) => strats.iter().flat_map(|s| s.types()).collect(),
Union(strats) => strats.iter().flat_map(|s| s.types()).collect(),
Filter(_, ty) => vec![ty.clone()],
}
}
}
impl ToTokens for Strategy {
fn to_tokens(&self, tokens: &mut TokenStream) {
use self::Strategy::*;
match self {
Arbitrary(ty, span) => tokens.append_all(quote_spanned!(*span=>
<#ty as _proptest::arbitrary::Arbitrary>::Strategy
)),
Regex(ty) => quote_append!(tokens,
<#ty as _proptest::string::StrategyFromRegex>::Strategy
),
Existential(ty) => quote_append!(tokens,
_proptest::strategy::BoxedStrategy<#ty>
),
Value(ty) => quote_append!(tokens, fn() -> #ty ),
Map(strats) => {
let types = self.types();
let field_tys = NestedTuple(&types);
let strats = NestedTuple(&strats);
quote_append!(tokens,
_proptest::strategy::Map< ( #strats ),
fn( #field_tys ) -> Self
>
)
}
Union(strats) => union_strat_to_tokens(tokens, strats),
Filter(strat, ty) => quote_append!(tokens,
_proptest::strategy::Filter<#strat, fn(&#ty) -> bool>
),
}
}
}
pub enum FromReg {
Top,
Num(usize),
}
pub enum ToReg {
Range(usize),
API,
}
pub enum Ctor {
Arbitrary(syn::Type, Option<usize>, Span),
Regex(syn::Type, syn::Expr),
Existential(syn::Expr),
Value(syn::Expr),
ValueExistential(syn::Expr),
Map(Box<[Ctor]>, MapClosure),
Union(Box<[(u32, Ctor)]>),
Extract(Box<Ctor>, ToReg, FromReg),
Filter(Box<Ctor>, syn::Expr),
}
pub fn extract_all(c: Ctor, to: usize, from: FromReg) -> Ctor {
extract(c, ToReg::Range(to), from)
}
pub fn extract_api(c: Ctor, from: FromReg) -> Ctor {
extract(c, ToReg::API, from)
}
impl ToTokens for FromReg {
fn to_tokens(&self, tokens: &mut TokenStream) {
match self {
FromReg::Top => call_site_ident(TOP_PARAM_NAME).to_tokens(tokens),
FromReg::Num(reg) => param(*reg).to_tokens(tokens),
}
}
}
impl ToTokens for ToReg {
fn to_tokens(&self, tokens: &mut TokenStream) {
match *self {
ToReg::Range(to) if to == 1 => param(0).to_tokens(tokens),
ToReg::Range(to) => {
let params: Vec<_> = (0..to).map(param).collect();
NestedTuple(¶ms).to_tokens(tokens)
}
ToReg::API => call_site_ident(API_PARAM_NAME).to_tokens(tokens),
}
}
}
impl ToTokens for Ctor {
fn to_tokens(&self, tokens: &mut TokenStream) {
use self::Ctor::*;
match self {
Filter(ctor, filter) => quote_append!(tokens,
_proptest::strategy::Strategy::prop_filter(
#ctor, stringify!(#filter), #filter)
),
Extract(ctor, to, from) => quote_append!(tokens, {
let #to = #from; #ctor
}),
Arbitrary(ty, fv, span) => {
tokens.append_all(if let Some(fv) = fv {
let args = param(*fv);
quote_spanned!(*span=>
_proptest::arbitrary::any_with::<#ty>(#args)
)
} else {
quote_spanned!(*span=>
_proptest::arbitrary::any::<#ty>()
)
})
}
Regex(ty, regex) => quote_append!(tokens,
<#ty as _proptest::string::StrategyFromRegex>::from_regex(#regex)
),
Existential(expr) => quote_append!(tokens,
_proptest::strategy::Strategy::boxed( #expr ) ),
Value(expr) => quote_append!(tokens, || #expr ),
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),
Union(ctors) => union_ctor_to_tokens(tokens, ctors),
}
}
}
struct NestedTuple<'a, T>(&'a [T]);
impl<'a, T: ToTokens> ToTokens for NestedTuple<'a, T> {
fn to_tokens(&self, tokens: &mut TokenStream) {
let NestedTuple(elems) = self;
if elems.is_empty() {
quote_append!(tokens, ());
} else if let [x] = elems {
x.to_tokens(tokens);
} else {
let chunks = elems.chunks(NESTED_TUPLE_CHUNK_SIZE);
Recurse(&chunks).to_tokens(tokens);
}
struct Recurse<'a, T: ToTokens>(&'a ::std::slice::Chunks<'a, T>);
impl<'a, T: ToTokens> ToTokens for Recurse<'a, T> {
fn to_tokens(&self, tokens: &mut TokenStream) {
let mut chunks = self.0.clone();
if let Some(head) = chunks.next() {
if let [c] = head {
quote_append!(tokens, #c);
} else {
let tail = Recurse(&chunks);
quote_append!(tokens, (#(#head,)* #tail));
}
}
}
}
}
}
fn map_ctor_to_tokens(
tokens: &mut TokenStream,
ctors: &[Ctor],
closure: &MapClosure,
) {
let ctors = NestedTuple(ctors);
quote_append!(tokens,
_proptest::strategy::Strategy::prop_map(
#ctors,
#closure
)
);
}
fn union_ctor_to_tokens(tokens: &mut TokenStream, ctors: &[(u32, Ctor)]) {
if ctors.is_empty() {
return;
}
if let [(_, ctor)] = ctors {
ctor.to_tokens(tokens);
return;
}
let mut chunks = ctors.chunks(UNION_CHUNK_SIZE);
let chunk = chunks.next().unwrap();
let head = chunk.iter().map(wrap_arc);
let tail = Recurse(weight_sum(ctors) - weight_sum(chunk), chunks);
quote_append!(tokens,
_proptest::strategy::TupleUnion::new(( #(#head,)* #tail ))
);
struct Recurse<'a>(u32, ::std::slice::Chunks<'a, (u32, Ctor)>);
impl<'a> ToTokens for Recurse<'a> {
fn to_tokens(&self, tokens: &mut TokenStream) {
let (tweight, mut chunks) = (self.0, self.1.clone());
if let Some(chunk) = chunks.next() {
if let [(w, c)] = chunk {
quote_append!(tokens, (#w, ::std::sync::Arc::new(#c)) );
} else {
let head = chunk.iter().map(wrap_arc);
let tail = Recurse(tweight - weight_sum(chunk), chunks);
quote_append!(tokens,
(#tweight, ::std::sync::Arc::new(
_proptest::strategy::TupleUnion::new((
#(#head,)* #tail
))))
);
}
}
}
}
fn weight_sum(ctors: &[(u32, Ctor)]) -> u32 {
use std::num::Wrapping;
let Wrapping(x) = ctors.iter().map(|&(w, _)| Wrapping(w)).sum();
x
}
fn wrap_arc(arg: &(u32, Ctor)) -> TokenStream {
let (w, c) = arg;
quote!( (#w, ::std::sync::Arc::new(#c)) )
}
}
fn union_strat_to_tokens(tokens: &mut TokenStream, strats: &[Strategy]) {
if strats.is_empty() {
return;
}
if let [strat] = strats {
strat.to_tokens(tokens);
return;
}
let mut chunks = strats.chunks(UNION_CHUNK_SIZE);
let chunk = chunks.next().unwrap();
let head = chunk.iter().map(wrap_arc);
let tail = Recurse(chunks);
quote_append!(tokens,
_proptest::strategy::TupleUnion<( #(#head,)* #tail )>
);
struct Recurse<'a>(::std::slice::Chunks<'a, Strategy>);
impl<'a> ToTokens for Recurse<'a> {
fn to_tokens(&self, tokens: &mut TokenStream) {
let mut chunks = self.0.clone();
if let Some(chunk) = chunks.next() {
if let [s] = chunk {
quote_append!(tokens, (u32, ::std::sync::Arc<#s>) );
} else {
let head = chunk.iter().map(wrap_arc);
let tail = Recurse(chunks);
quote_append!(tokens,
(u32,
::std::sync::Arc<_proptest::strategy::TupleUnion<(
#(#head,)* #tail
)>>)
);
}
}
}
}
fn wrap_arc(s: &Strategy) -> TokenStream {
quote!( (u32, ::std::sync::Arc<#s>) )
}
}
fn extract(c: Ctor, to: ToReg, from: FromReg) -> Ctor {
Ctor::Extract(Box::new(c), to, from)
}
fn param<'a>(fv: usize) -> FreshVar<'a> {
fresh_var("param", fv)
}
pub fn map_closure(path: syn::Path, fs: &[syn::Field]) -> MapClosure {
MapClosure(path, fs.to_owned())
}
#[derive(Debug)]
pub struct MapClosure(syn::Path, Vec<syn::Field>);
impl ToTokens for MapClosure {
fn to_tokens(&self, tokens: &mut TokenStream) {
fn tmp_var<'a>(idx: usize) -> FreshVar<'a> {
fresh_var("tmp", idx)
}
let MapClosure(path, fields) = self;
let count = fields.len();
let tmps: Vec<_> = (0..count).map(tmp_var).collect();
let inits = fields.iter().enumerate().map(|(idx, field)| {
let tv = tmp_var(idx);
if let Some(name) = &field.ident {
quote_spanned!(field.span()=> #name: #tv )
} else {
let name = syn::Member::Unnamed(syn::Index::from(idx));
quote_spanned!(field.span()=> #name: #tv )
}
});
let tmps = NestedTuple(&tmps);
quote_append!(tokens, | #tmps | #path { #(#inits),* } );
}
}
fn fresh_var(prefix: &str, count: usize) -> FreshVar {
FreshVar { prefix, count }
}
struct FreshVar<'a> {
prefix: &'a str,
count: usize,
}
impl<'a> ToTokens for FreshVar<'a> {
fn to_tokens(&self, tokens: &mut TokenStream) {
let ident = format!("{}_{}", self.prefix, self.count);
call_site_ident(&ident).to_tokens(tokens)
}
}
fn call_site_ident(ident: &str) -> syn::Ident {
syn::Ident::new(ident, Span::call_site())
}