mutatis-derive 0.5.2

extern crate proc_macro;

use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::{spanned::Spanned, *};

mod container_attributes;
mod field_attributes;
use container_attributes::ContainerAttributes;
use field_attributes::FieldBehavior;

static MUTATIS_ATTRIBUTE_NAME: &str = "mutatis";

#[proc_macro_derive(Mutate, attributes(mutatis))]
pub fn derive_mutator(tokens: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = syn::parse_macro_input!(tokens as DeriveInput);
    expand_derive_mutator(input)
        .unwrap_or_else(syn::Error::into_compile_error)
        .into()
}

fn expand_derive_mutator(input: DeriveInput) -> Result<TokenStream> {
    let container_attrs = ContainerAttributes::from_derive_input(&input)?;
    let mutator_ty = MutatorType::new(&input, &container_attrs)?;

    let mutator_type_def = gen_mutator_type_def(&input, &mutator_ty, &container_attrs)?;
    let mutator_type_default_impl = gen_mutator_type_default_impl(&mutator_ty)?;
    let mutator_ctor = gen_mutator_ctor(&mutator_ty)?;
    let mutator_impl = gen_mutator_impl(&input, &mutator_ty)?;
    let default_mutator_impl = gen_default_mutator_impl(&mutator_ty, &container_attrs)?;
    let generate_impl = gen_generate_impl(&input, &mutator_ty, &container_attrs)?;

    Ok(quote! {
        #mutator_type_def
        #mutator_type_default_impl
        #mutator_ctor
        #mutator_impl
        #default_mutator_impl
        #generate_impl
    })
}

struct MutatorType {
    ty_name: Ident,

    mutator_name: Ident,

    mutator_fields: Vec<MutatorField>,

    /// A vec of quoted generic parameters, without any bounds but with `const`
    /// defs, e.g. `'a`, `const N: usize`, or `T`.
    ty_impl_generics: Vec<TokenStream>,

    /// A vec of quoted generic parameters, without any bounds and without any
    /// `const` defs, e.t. `'a`, `N`, or `T.
    ty_name_generics: Vec<TokenStream>,

    /// A vec of quoted bounds for the generics above, e.g. `A: Iterator<Item =
    /// B>,`.
    ty_generics_bounds: Vec<TokenStream>,
}

impl MutatorType {
    fn new(input: &DeriveInput, container_attrs: &ContainerAttributes) -> Result<Self> {
        let ty_name = input.ident.clone();

        let mutator_name = container_attrs
            .mutator_name
            .clone()
            .unwrap_or_else(|| Ident::new(&format!("{}Mutator", input.ident), input.ident.span()));

        let mutator_fields = get_mutator_fields(&input)?;

        let mut ty_impl_generics = vec![];
        let mut ty_name_generics = vec![];
        let mut ty_generics_bounds = vec![];

        for gen in &input.generics.params {
            match gen {
                GenericParam::Lifetime(l) => {
                    if !l.bounds.is_empty() {
                        ty_generics_bounds.push(quote! { #l });
                    }

                    let l = &l.lifetime;
                    ty_impl_generics.push(quote! { #l });
                    ty_name_generics.push(quote! { #l });
                }
                GenericParam::Const(c) => {
                    ty_impl_generics.push(quote! { #c });
                    let c = &c.ident;
                    ty_name_generics.push(quote! { #c });
                }
                GenericParam::Type(t) => {
                    if !t.bounds.is_empty() {
                        ty_generics_bounds.push(quote! { #t });
                    }
                    let t = &t.ident;
                    ty_impl_generics.push(quote! { #t });
                    ty_name_generics.push(quote! { #t });
                }
            }
        }

        if let Some(wc) = &input.generics.where_clause {
            for bound in wc.predicates.iter() {
                ty_generics_bounds.push(quote! { #bound });
            }
        }

        Ok(Self {
            ty_name,
            mutator_name,
            mutator_fields,
            ty_impl_generics,
            ty_name_generics,
            ty_generics_bounds,
        })
    }

    fn mutator_impl_generics_iter(&self) -> impl Iterator<Item = TokenStream> + '_ {
        self.ty_impl_generics.iter().cloned().chain(
            self.mutator_fields
                .iter()
                .filter_map(|f| f.generic.as_ref().map(|g| quote! { #g })),
        )
    }

    /// All the `impl` generic parameters for this mutator, including those
    /// inherited from the type that it is a mutator for.
    fn mutator_impl_generics(&self) -> TokenStream {
        let impl_generics = self
            .ty_impl_generics
            .iter()
            .cloned()
            .chain(
                self.mutator_fields
                    .iter()
                    .filter_map(|f| f.generic.as_ref().map(|g| quote! { #g })),
            )
            .collect::<Vec<_>>();
        if impl_generics.is_empty() {
            quote! {}
        } else {
            quote! { < #( #impl_generics ),* > }
        }
    }

    /// All the named (i.e. just the "N" and excluding "const", ":", and "usize"
    /// in `const N: usize` generics) generic parameters for this mutator,
    /// including those inherited from the type that it is a mutator for.
    fn mutator_name_generics_iter(&self) -> impl Iterator<Item = TokenStream> + '_ {
        self.ty_name_generics.iter().cloned().chain(
            self.mutator_fields
                .iter()
                .filter_map(|f| f.generic.as_ref().map(|g| quote! { #g })),
        )
    }

    fn mutator_impl_generics_with_defaults_iter(&self) -> impl Iterator<Item = TokenStream> + '_ {
        self.ty_impl_generics
            .iter()
            .cloned()
            .chain(self.mutator_fields.iter().filter_map(move |f| {
                f.generic.as_ref().map(|g| {
                    let for_ty = &f.for_ty;
                    quote! { #g = <#for_ty as mutatis::DefaultMutate>::DefaultMutate }
                })
            }))
    }

    fn ty_name_with_generics(&self) -> TokenStream {
        let ty_name = &self.ty_name;
        if self.ty_name_generics.is_empty() {
            quote! { #ty_name }
        } else {
            let ty_generics = self.ty_name_generics.iter();
            quote! { #ty_name < #( #ty_generics ),* > }
        }
    }

    fn mutator_name_with_generics(&self, kind: MutatorNameGenericsKind) -> TokenStream {
        let mutator_name = &self.mutator_name;

        let generics = match kind {
            MutatorNameGenericsKind::Generics => {
                self.mutator_name_generics_iter().collect::<Vec<_>>()
            }
            MutatorNameGenericsKind::Impl {
                impl_default: false,
            } => self.mutator_impl_generics_iter().collect::<Vec<_>>(),
            MutatorNameGenericsKind::Impl { impl_default: true } => self
                .mutator_impl_generics_with_defaults_iter()
                .collect::<Vec<_>>(),
            MutatorNameGenericsKind::JustTyGenerics => self.ty_name_generics.clone(),
        };

        if generics.is_empty() {
            quote! { #mutator_name }
        } else {
            quote! { #mutator_name < #( #generics ),* > }
        }
    }

    fn where_clause(&self, kind: WhereClauseKind) -> TokenStream {
        let mut bounds = self.ty_generics_bounds.clone();

        match kind {
            WhereClauseKind::NoMutateBounds => {}
            WhereClauseKind::MutateBounds => {
                for f in &self.mutator_fields {
                    let for_ty = &f.for_ty;
                    if let Some(g) = f.generic.as_ref() {
                        bounds.push(quote! { #g: mutatis::Mutate<#for_ty> });
                    } else {
                        debug_assert_eq!(f.behavior, FieldBehavior::DefaultMutate);
                        bounds.push(quote! { #for_ty: mutatis::DefaultMutate });
                    }
                }
            }
            WhereClauseKind::MutateAndGenerateBounds => {
                for f in &self.mutator_fields {
                    let for_ty = &f.for_ty;
                    if let Some(g) = f.generic.as_ref() {
                        bounds.push(
                            quote! { #g: mutatis::Mutate<#for_ty> + mutatis::Generate<#for_ty> },
                        );
                    } else {
                        debug_assert_eq!(f.behavior, FieldBehavior::DefaultMutate);
                        bounds.push(quote! { #for_ty: mutatis::DefaultMutate });
                        bounds.push(quote! { <#for_ty as mutatis::DefaultMutate>::DefaultMutate: mutatis::Generate<#for_ty> });
                    }
                }
            }
            WhereClauseKind::DefaultBounds => {
                for f in &self.mutator_fields {
                    if let Some(g) = f.generic.as_ref() {
                        bounds.push(quote! { #g: Default });
                    } else {
                        let for_ty = &f.for_ty;
                        debug_assert_eq!(f.behavior, FieldBehavior::DefaultMutate);
                        bounds.push(quote! { #for_ty: mutatis::DefaultMutate });
                    }
                }
            }
            WhereClauseKind::DefaultMutateBounds => {
                for f in &self.mutator_fields {
                    let for_ty = &f.for_ty;
                    bounds.push(quote! { #for_ty: mutatis::DefaultMutate });
                }
            }
        }

        if bounds.is_empty() {
            quote! {}
        } else {
            quote! { where #( #bounds ),* }
        }
    }

    fn phantom_fields_defs<'a>(
        &self,
        input: &'a DeriveInput,
    ) -> impl Iterator<Item = TokenStream> + 'a {
        let make_phantom_field = |i, ty| {
            let ident = Ident::new(&format!("_phantom{i}"), Span::call_site());
            quote! { #ident : core::marker::PhantomData<#ty> , }
        };

        input
            .generics
            .params
            .iter()
            .enumerate()
            .map(move |(i, g)| match g {
                GenericParam::Lifetime(l) => {
                    let l = &l.lifetime;
                    make_phantom_field(i, quote! { & #l () })
                }
                GenericParam::Const(c) => {
                    let c = &c.ident;
                    make_phantom_field(i, quote! { [(); #c] })
                }
                GenericParam::Type(t) => {
                    let t = &t.ident;
                    make_phantom_field(i, quote! { #t })
                }
            })
    }

    fn phantom_fields_literals(&self) -> impl Iterator<Item = TokenStream> + '_ {
        (0..self.ty_name_generics.len()).map(|i| {
            let ident = Ident::new(&format!("_phantom{i}"), Span::call_site());
            quote! { #ident : core::marker::PhantomData, }
        })
    }
}

#[derive(Clone, Copy)]
enum WhereClauseKind {
    NoMutateBounds,
    MutateBounds,
    MutateAndGenerateBounds,
    DefaultBounds,
    DefaultMutateBounds,
}

#[derive(Clone, Copy)]
enum MutatorNameGenericsKind {
    Generics,
    Impl { impl_default: bool },
    JustTyGenerics,
}

struct MutatorField {
    /// The identifier for this field inside the mutator struct.
    ident: Ident,
    /// The generic type parameter for this field, if any.
    generic: Option<Ident>,
    /// The behavior for this field.
    behavior: FieldBehavior,
    /// The type that this field is a mutator for.
    for_ty: Type,
}

fn get_mutator_fields(input: &DeriveInput) -> Result<Vec<MutatorField>> {
    let mut i = 0;
    let mut generic = |b: &FieldBehavior| -> Option<Ident> {
        if b.needs_generic() {
            let g = Ident::new(&format!("MutatorT{}", i), Span::call_site());
            i += 1;
            Some(g)
        } else {
            None
        }
    };

    match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(fields) => fields
                .named
                .iter()
                .filter_map(|f| {
                    FieldBehavior::for_field(f)
                        .map(|b| {
                            b.map(|b| MutatorField {
                                ident: f.ident.clone().unwrap(),
                                generic: generic(&b),
                                behavior: b,
                                for_ty: f.ty.clone(),
                            })
                        })
                        .transpose()
                })
                .collect(),
            Fields::Unnamed(fields) => fields
                .unnamed
                .iter()
                .enumerate()
                .filter_map(|(i, f)| {
                    FieldBehavior::for_field(f)
                        .map(|b| {
                            b.map(|b| MutatorField {
                                ident: Ident::new(&format!("field{}", i), f.span()),
                                generic: generic(&b),
                                behavior: b,
                                for_ty: f.ty.clone(),
                            })
                        })
                        .transpose()
                })
                .collect(),
            Fields::Unit => Ok(vec![]),
        },
        Data::Enum(data) => Ok(data
            .variants
            .iter()
            .map(|v| {
                let prefix = v.ident.to_string().to_lowercase();
                match v.fields {
                    Fields::Named(ref fields) => fields
                        .named
                        .iter()
                        .filter_map(|f| {
                            FieldBehavior::for_field(f)
                                .map(|b| {
                                    b.map(|b| MutatorField {
                                        ident: Ident::new(
                                            &format!("{prefix}_{}", f.ident.clone().unwrap()),
                                            f.span(),
                                        ),
                                        generic: generic(&b),
                                        behavior: b,
                                        for_ty: f.ty.clone(),
                                    })
                                })
                                .transpose()
                        })
                        .collect::<Result<Vec<_>>>(),
                    Fields::Unnamed(ref fields) => fields
                        .unnamed
                        .iter()
                        .enumerate()
                        .filter_map(|(i, f)| {
                            FieldBehavior::for_field(f)
                                .map(|b| {
                                    b.map(|b| MutatorField {
                                        ident: Ident::new(&format!("{prefix}{i}"), f.span()),
                                        generic: generic(&b),
                                        behavior: b,
                                        for_ty: f.ty.clone(),
                                    })
                                })
                                .transpose()
                        })
                        .collect::<Result<Vec<_>>>(),
                    Fields::Unit => Ok(vec![]),
                }
            })
            .collect::<Result<Vec<_>>>()?
            .into_iter()
            .flat_map(|fs| fs)
            .collect()),
        Data::Union(_) => Err(Error::new_spanned(
            input,
            "cannot `derive(Mutate)` on a union",
        )),
    }
}

fn gen_mutator_type_def(
    input: &DeriveInput,
    mutator_ty: &MutatorType,
    container_attrs: &ContainerAttributes,
) -> Result<TokenStream> {
    let vis = &input.vis;
    let name = &input.ident;

    let impl_default = container_attrs.default_mutate.unwrap_or(true);
    let mutator_name =
        mutator_ty.mutator_name_with_generics(MutatorNameGenericsKind::Impl { impl_default });

    let mut temp: Option<LitStr> = None;
    let doc = container_attrs.mutator_doc.as_deref().unwrap_or_else(|| {
        temp = Some(LitStr::new(
            &format!(" A mutator for the `{name}` type."),
            input.ident.span(),
        ));
        std::slice::from_ref(temp.as_ref().unwrap())
    });

    let where_clause = mutator_ty.where_clause(WhereClauseKind::NoMutateBounds);

    let fields = mutator_ty
        .mutator_fields
        .iter()
        .map(|f| {
            let ident = &f.ident;
            if let Some(g) = f.generic.as_ref() {
                quote! { #ident: #g , }
            } else {
                let for_ty = &f.for_ty;
                debug_assert_eq!(f.behavior, FieldBehavior::DefaultMutate);
                quote! { #ident: <#for_ty as mutatis::DefaultMutate>::DefaultMutate, }
            }
        })
        .collect::<Vec<_>>();

    let phantoms = mutator_ty.phantom_fields_defs(input);

    Ok(quote! {
        #( #[doc = #doc] )*
        // #[derive(Clone, Debug)]
        #vis struct #mutator_name #where_clause {
            #( #fields )*
            #( #phantoms )*
            _private: (),
        }
    })
}

fn gen_mutator_type_default_impl(mutator_ty: &MutatorType) -> Result<TokenStream> {
    let impl_generics = mutator_ty.mutator_impl_generics();
    let mutator_name = mutator_ty.mutator_name_with_generics(MutatorNameGenericsKind::Generics);
    let where_clause = mutator_ty.where_clause(WhereClauseKind::DefaultBounds);

    let fields = mutator_ty
        .mutator_fields
        .iter()
        .map(|f| {
            let ident = &f.ident;
            quote! { #ident: Default::default(), }
        })
        .collect::<Vec<_>>();

    let phantoms = mutator_ty.phantom_fields_literals();

    Ok(quote! {
        #[automatically_derived]
        impl #impl_generics Default for #mutator_name #where_clause {
            fn default() -> Self {
                Self {
                    #( #fields )*
                    #( #phantoms )*
                    _private: (),
                }
            }
        }
    })
}

fn gen_mutator_ctor(mutator_ty: &MutatorType) -> Result<TokenStream> {
    let impl_generics = mutator_ty.mutator_impl_generics();

    let params = mutator_ty
        .mutator_fields
        .iter()
        .filter_map(|f| {
            f.generic.as_ref().map(|g| {
                let ident = &f.ident;
                quote! { #ident: #g , }
            })
        })
        .collect::<Vec<_>>();

    let fields = mutator_ty
        .mutator_fields
        .iter()
        .map(|f| {
            let ident = &f.ident;
            if f.generic.is_some() {
                quote! { #ident , }
            } else {
                let for_ty = &f.for_ty;
                debug_assert_eq!(f.behavior, FieldBehavior::DefaultMutate);
                quote! { #ident: mutatis::mutators::default::<#for_ty>() , }
            }
        })
        .collect::<Vec<_>>();

    let name = &mutator_ty.mutator_name_with_generics(MutatorNameGenericsKind::Generics);
    let doc = format!("Construct a new `{name}` instance.");
    let where_clause = mutator_ty.where_clause(WhereClauseKind::NoMutateBounds);
    let phantoms = mutator_ty.phantom_fields_literals();

    Ok(quote! {
        impl #impl_generics #name #where_clause {
            #[doc = #doc]
            #[inline]
            pub fn new( #( #params )* ) -> Self {
                Self {
                    #( #fields )*
                    #( #phantoms )*
                    _private: (),
                }
            }
        }
    })
}

fn gen_mutator_impl(input: &DeriveInput, mutator_ty: &MutatorType) -> Result<TokenStream> {
    let impl_generics = mutator_ty.mutator_impl_generics();

    let ty_name = mutator_ty.ty_name_with_generics();

    let is_multi_variant_enum = matches!(&input.data, Data::Enum(data) if data.variants.len() > 1);
    let where_clause = if is_multi_variant_enum {
        mutator_ty.where_clause(WhereClauseKind::MutateAndGenerateBounds)
    } else {
        mutator_ty.where_clause(WhereClauseKind::MutateBounds)
    };

    let mut fields_iter = mutator_ty.mutator_fields.iter();
    let mut make_mutation = |value| {
        let ident = &fields_iter.next().unwrap().ident;
        quote! { self.#ident.mutate(mutations, #value)?; }
    };

    let mutation_body = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(fields) => {
                let mutations = fields
                    .named
                    .iter()
                    .filter(|f| FieldBehavior::for_field(f).unwrap().is_some())
                    .map(|f| {
                        let ident = &f.ident;
                        make_mutation(quote! { &mut value.#ident })
                    });
                quote! {
                    #( #mutations )*
                }
            }
            Fields::Unnamed(fields) => {
                let mutations = fields
                    .unnamed
                    .iter()
                    .enumerate()
                    .filter(|(_i, f)| FieldBehavior::for_field(f).unwrap().is_some())
                    .map(|(i, f)| {
                        let index = Index {
                            index: u32::try_from(i).unwrap(),
                            span: f.span(),
                        };
                        make_mutation(quote! { &mut value.#index })
                    });
                quote! {
                    #( #mutations )*
                }
            }
            Fields::Unit => quote! {},
        },

        Data::Enum(data) => {
            // Build the existing field-mutation match arms.
            let mut field_mutation_arms = vec![];
            for v in data.variants.iter() {
                let variant_ident = &v.ident;
                match &v.fields {
                    Fields::Named(fields) => {
                        let mut patterns = vec![];
                        let mutates = fields
                            .named
                            .iter()
                            .filter_map(|f| {
                                let ident = &f.ident;
                                if FieldBehavior::for_field(f).unwrap().is_some() {
                                    patterns.push(quote! { ref mut #ident , });
                                    Some(make_mutation(quote! { #ident }))
                                } else {
                                    patterns.push(quote! { #ident: _ , });
                                    None
                                }
                            })
                            .collect::<Vec<_>>();
                        field_mutation_arms.push(quote! {
                            #ty_name::#variant_ident { #( #patterns )* } => {
                                #( #mutates )*
                            }
                        });
                    }

                    Fields::Unnamed(fields) => {
                        let mut patterns = vec![];
                        let mutates = fields
                            .unnamed
                            .iter()
                            .enumerate()
                            .filter_map(|(i, f)| {
                                if FieldBehavior::for_field(f).unwrap().is_some() {
                                    let binding = Ident::new(&format!("field{}", i), f.span());
                                    patterns.push(quote! { ref mut #binding , });
                                    Some(make_mutation(quote! { #binding }))
                                } else {
                                    patterns.push(quote! { _ , });
                                    None
                                }
                            })
                            .collect::<Vec<_>>();
                        field_mutation_arms.push(quote! {
                            #ty_name::#variant_ident( #( #patterns )* ) => {
                                #( #mutates )*
                            }
                        });
                    }

                    Fields::Unit => {
                        field_mutation_arms.push(quote! {
                            #ty_name::#variant_ident => {}
                        });
                    }
                }
            }

            // Build variant-switching mutations for enums with multiple
            // variants.
            let variant_switching = if data.variants.len() > 1 {
                // Build a match to determine the current variant index.
                let index_arms: Vec<_> = data
                    .variants
                    .iter()
                    .enumerate()
                    .map(|(v_idx, v)| {
                        let variant_ident = &v.ident;
                        match &v.fields {
                            Fields::Named(_) => {
                                quote! { #ty_name::#variant_ident { .. } => #v_idx, }
                            }
                            Fields::Unnamed(_) => {
                                quote! { #ty_name::#variant_ident(..) => #v_idx, }
                            }
                            Fields::Unit => quote! { #ty_name::#variant_ident => #v_idx, },
                        }
                    })
                    .collect();

                // For each variant, build an expression that constructs a new
                // instance of that variant. For non-ignored fields, use the
                // mutator's `generate` method. For ignored fields, use
                // `Default::default()`.
                let mut variant_mutations = vec![];
                let mut mutator_field_offset = 0usize;
                for (v_idx, v) in data.variants.iter().enumerate() {
                    let variant_ident = &v.ident;

                    let construction = match &v.fields {
                        Fields::Named(fields) => {
                            let field_exprs: Vec<_> = fields
                                .named
                                .iter()
                                .map(|f| {
                                    let ident = &f.ident;
                                    if let Some(_behavior) = FieldBehavior::for_field(f).unwrap() {
                                        let mutator_ident =
                                            &mutator_ty.mutator_fields[mutator_field_offset].ident;
                                        mutator_field_offset += 1;
                                        quote! { #ident: self.#mutator_ident.generate(ctx)? }
                                    } else {
                                        quote! { #ident: Default::default() }
                                    }
                                })
                                .collect();
                            quote! {
                                *value = #ty_name::#variant_ident { #( #field_exprs ),* };
                            }
                        }
                        Fields::Unnamed(fields) => {
                            let field_exprs: Vec<_> = fields
                                .unnamed
                                .iter()
                                .map(|f| {
                                    if let Some(_behavior) = FieldBehavior::for_field(f).unwrap() {
                                        let mutator_ident =
                                            &mutator_ty.mutator_fields[mutator_field_offset].ident;
                                        mutator_field_offset += 1;
                                        quote! { self.#mutator_ident.generate(ctx)? }
                                    } else {
                                        quote! { Default::default() }
                                    }
                                })
                                .collect();
                            quote! {
                                *value = #ty_name::#variant_ident( #( #field_exprs ),* );
                            }
                        }
                        Fields::Unit => {
                            quote! {
                                *value = #ty_name::#variant_ident;
                            }
                        }
                    };

                    variant_mutations.push((v_idx, construction));
                }

                let num_variants = data.variants.len();
                let group_count = num_variants - 1;
                let group_arms: Vec<_> = variant_mutations
                    .iter()
                    .map(|(v_idx, construction)| {
                        quote! {
                            #v_idx => {
                                #construction
                                Ok(())
                            }
                        }
                    })
                    .collect();

                quote! {
                    let _variant_index: usize = match value {
                        #( #index_arms )*
                    };
                    mutations.mutation_group(#group_count as u32, |ctx, _which| {
                        let _target = if (_which as usize) >= _variant_index {
                            _which as usize + 1
                        } else {
                            _which as usize
                        };
                        match _target {
                            #( #group_arms )*
                            _ => unreachable!(),
                        }
                    })?;
                }
            } else {
                quote! {}
            };

            quote! {
                #variant_switching
                match *value {
                    #( #field_mutation_arms )*
                }
            }
        }

        Data::Union(_) => {
            return Err(Error::new_spanned(
                input,
                "cannot `derive(Mutate)` on a union",
            ))
        }
    };

    let mutate_method = quote! {
        fn mutate(
            &mut self,
            mutations: &mut mutatis::Candidates,
            value: &mut #ty_name,
        ) -> mutatis::Result<()> {
            #mutation_body

            // Silence unused-variable warnings if every field was marked
            // `ignore`. Allow unreachable for empty enums.
            #[allow(unreachable_code)]
            let _ = (mutations, value);

            Ok(())
        }
    };

    let mutation_count_body = match &input.data {
        Data::Struct(data) => {
            let mut field_counts = vec![];
            let mut mutator_field_idx = 0usize;
            match &data.fields {
                Fields::Named(fields) => {
                    for f in fields.named.iter() {
                        if FieldBehavior::for_field(f).unwrap().is_some() {
                            let field_ident = &f.ident;
                            let mutator_ident = &mutator_ty.mutator_fields[mutator_field_idx].ident;
                            mutator_field_idx += 1;
                            field_counts.push(quote! {
                                _count += self.#mutator_ident.mutation_count(
                                    &value.#field_ident,
                                    shrink,
                                )?;
                            });
                        }
                    }
                }
                Fields::Unnamed(fields) => {
                    for (i, f) in fields.unnamed.iter().enumerate() {
                        if FieldBehavior::for_field(f).unwrap().is_some() {
                            let index = Index {
                                index: u32::try_from(i).unwrap(),
                                span: f.span(),
                            };
                            let mutator_ident = &mutator_ty.mutator_fields[mutator_field_idx].ident;
                            mutator_field_idx += 1;
                            field_counts.push(quote! {
                                _count += self.#mutator_ident.mutation_count(
                                    &value.#index,
                                    shrink,
                                )?;
                            });
                        }
                    }
                }
                Fields::Unit => {}
            }
            quote! {
                let mut _count = 0u32;
                #(#field_counts)*
                Some(_count)
            }
        }

        Data::Enum(data) => {
            let variant_switch_count = if data.variants.len() > 1 {
                data.variants.len() - 1
            } else {
                0
            };

            let mut count_arms = vec![];
            let mut mutator_field_idx = 0usize;
            for v in data.variants.iter() {
                let variant_ident = &v.ident;
                match &v.fields {
                    Fields::Named(fields) => {
                        let mut patterns = vec![];
                        let mut fld_counts = vec![];
                        for f in fields.named.iter() {
                            let ident = &f.ident;
                            if FieldBehavior::for_field(f).unwrap().is_some() {
                                patterns.push(quote! { ref #ident, });
                                let mutator_ident =
                                    &mutator_ty.mutator_fields[mutator_field_idx].ident;
                                mutator_field_idx += 1;
                                fld_counts.push(quote! {
                                    _count += self.#mutator_ident.mutation_count(
                                        #ident,
                                        shrink,
                                    )?;
                                });
                            } else {
                                patterns.push(quote! { #ident: _, });
                            }
                        }
                        count_arms.push(quote! {
                            #ty_name::#variant_ident { #(#patterns)* } => {
                                #(#fld_counts)*
                            }
                        });
                    }
                    Fields::Unnamed(fields) => {
                        let mut patterns = vec![];
                        let mut fld_counts = vec![];
                        for (i, f) in fields.unnamed.iter().enumerate() {
                            if FieldBehavior::for_field(f).unwrap().is_some() {
                                let binding = Ident::new(&format!("field{}", i), f.span());
                                patterns.push(quote! { ref #binding, });
                                let mutator_ident =
                                    &mutator_ty.mutator_fields[mutator_field_idx].ident;
                                mutator_field_idx += 1;
                                fld_counts.push(quote! {
                                    _count += self.#mutator_ident.mutation_count(
                                        #binding,
                                        shrink,
                                    )?;
                                });
                            } else {
                                patterns.push(quote! { _, });
                            }
                        }
                        count_arms.push(quote! {
                            #ty_name::#variant_ident(#(#patterns)*) => {
                                #(#fld_counts)*
                            }
                        });
                    }
                    Fields::Unit => {
                        count_arms.push(quote! {
                            #ty_name::#variant_ident => {}
                        });
                    }
                }
            }

            if count_arms.is_empty() {
                quote! { Some(0u32) }
            } else {
                quote! {
                    let mut _count = #variant_switch_count as u32;
                    match *value {
                        #(#count_arms)*
                    }
                    Some(_count)
                }
            }
        }

        Data::Union(_) => quote! { None },
    };

    let mutation_count_method = quote! {
        #[inline]
        fn mutation_count(&self, value: &#ty_name, shrink: bool) -> core::option::Option<u32> {
            #mutation_count_body
        }
    };

    let mutator_name = &mutator_ty.mutator_name_with_generics(MutatorNameGenericsKind::Generics);

    Ok(quote! {
        #[automatically_derived]
        impl #impl_generics mutatis::Mutate<#ty_name> for #mutator_name
            #where_clause
        {
            #mutate_method
            #mutation_count_method
        }
    })
}

fn gen_default_mutator_impl(
    mutator_ty: &MutatorType,
    container_attrs: &ContainerAttributes,
) -> Result<TokenStream> {
    let impl_default = container_attrs.default_mutate.unwrap_or(true);
    if !impl_default {
        return Ok(quote! {});
    }

    let ty_generics = if mutator_ty.ty_impl_generics.is_empty() {
        quote! {}
    } else {
        let gens = &mutator_ty.ty_impl_generics;
        quote! { < #( #gens ),* > }
    };

    let ty_name = mutator_ty.ty_name_with_generics();
    let where_clause = mutator_ty.where_clause(WhereClauseKind::DefaultMutateBounds);
    let mutator_name =
        &mutator_ty.mutator_name_with_generics(MutatorNameGenericsKind::JustTyGenerics);

    Ok(quote! {
        #[automatically_derived]
        impl #ty_generics mutatis::DefaultMutate for #ty_name
            #where_clause
        {
            type DefaultMutate = #mutator_name;
        }
    })
}

fn gen_generate_impl(
    input: &DeriveInput,
    mutator_ty: &MutatorType,
    container_attrs: &ContainerAttributes,
) -> Result<TokenStream> {
    let impl_generate = container_attrs.generate.unwrap_or(true);
    if !impl_generate {
        return Ok(quote! {});
    }

    let impl_generics = mutator_ty.mutator_impl_generics();
    let ty_name = mutator_ty.ty_name_with_generics();
    let mutator_name = &mutator_ty.mutator_name_with_generics(MutatorNameGenericsKind::Generics);
    let where_clause = mutator_ty.where_clause(WhereClauseKind::MutateAndGenerateBounds);

    let mut fields_iter = mutator_ty.mutator_fields.iter();
    let mut next_field_generate = || -> TokenStream {
        let mf = fields_iter.next().unwrap();
        let ident = &mf.ident;
        quote! { self.#ident.generate(cx)? }
    };

    // For struct literals, we can't include generic args (e.g. `Foo<T> { .. }`
    // is invalid). Use the bare name and let Rust infer the type params.
    let bare_ty_name = &mutator_ty.ty_name;

    let generate_body = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(fields) => {
                let field_exprs: Vec<_> = fields
                    .named
                    .iter()
                    .map(|f| {
                        let ident = &f.ident;
                        if FieldBehavior::for_field(f).unwrap().is_some() {
                            let expr = next_field_generate();
                            quote! { #ident: #expr }
                        } else {
                            quote! { #ident: Default::default() }
                        }
                    })
                    .collect();
                quote! { Ok(#bare_ty_name { #( #field_exprs ),* }) }
            }
            Fields::Unnamed(fields) => {
                let field_exprs: Vec<_> = fields
                    .unnamed
                    .iter()
                    .map(|f| {
                        if FieldBehavior::for_field(f).unwrap().is_some() {
                            next_field_generate()
                        } else {
                            quote! { Default::default() }
                        }
                    })
                    .collect();
                quote! { Ok(#bare_ty_name( #( #field_exprs ),* )) }
            }
            Fields::Unit => {
                quote! { Ok(#bare_ty_name) }
            }
        },

        Data::Enum(data) => {
            if data.variants.is_empty() {
                quote! { unreachable!() }
            } else {
                let num_variants = data.variants.len();
                let mut mutator_field_offset = 0usize;
                let match_arms: Vec<_> = data
                    .variants
                    .iter()
                    .enumerate()
                    .map(|(v_idx, v)| {
                        let variant_ident = &v.ident;
                        let construction = match &v.fields {
                            Fields::Named(fields) => {
                                let field_exprs: Vec<_> = fields
                                    .named
                                    .iter()
                                    .map(|f| {
                                        let ident = &f.ident;
                                        if FieldBehavior::for_field(f).unwrap().is_some() {
                                            let mf =
                                                &mutator_ty.mutator_fields[mutator_field_offset];
                                            let mutator_ident = &mf.ident;
                                            mutator_field_offset += 1;
                                            quote! { #ident: self.#mutator_ident.generate(cx)? }
                                        } else {
                                            quote! { #ident: Default::default() }
                                        }
                                    })
                                    .collect();
                                quote! { #ty_name::#variant_ident { #( #field_exprs ),* } }
                            }
                            Fields::Unnamed(fields) => {
                                let field_exprs: Vec<_> = fields
                                    .unnamed
                                    .iter()
                                    .map(|f| {
                                        if FieldBehavior::for_field(f).unwrap().is_some() {
                                            let mf =
                                                &mutator_ty.mutator_fields[mutator_field_offset];
                                            let mutator_ident = &mf.ident;
                                            mutator_field_offset += 1;
                                            quote! { self.#mutator_ident.generate(cx)? }
                                        } else {
                                            quote! { Default::default() }
                                        }
                                    })
                                    .collect();
                                quote! { #ty_name::#variant_ident( #( #field_exprs ),* ) }
                            }
                            Fields::Unit => {
                                quote! { #ty_name::#variant_ident }
                            }
                        };
                        quote! { Some(#v_idx) => Ok(#construction), }
                    })
                    .collect();

                quote! {
                    match cx.rng().gen_index(#num_variants) {
                        #( #match_arms )*
                        _ => unreachable!(),
                    }
                }
            }
        }

        Data::Union(_) => {
            return Err(Error::new_spanned(
                input,
                "cannot `derive(Mutate)` on a union",
            ))
        }
    };

    Ok(quote! {
        #[automatically_derived]
        impl #impl_generics mutatis::Generate<#ty_name> for #mutator_name
            #where_clause
        {
            fn generate(&mut self, cx: &mut mutatis::Context) -> mutatis::Result<#ty_name> {
                #generate_body
            }
        }
    })
}