borrowize 0.2.0

//! Expansion for `#[derive(View)]`.

use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::{
    Data, DataEnum, DataStruct, DeriveInput, Error, Field, Fields, GenericParam, Generics, Ident,
    Lifetime, LifetimeParam, Result, Visibility, parse_quote, spanned::Spanned,
};

use crate::{
    attrs::{self, FieldOptions, StructOptions},
    mapping::{self, FallbackBorrowMode},
};

const GENERATED_BORROW_LIFETIME: &str = "borrowize";

pub(crate) fn expand(input: &DeriveInput) -> Result<TokenStream> {
    let struct_options = attrs::parse_struct_options(&input.attrs)?;
    let borrow_lifetime = generated_borrow_lifetime(&input.generics)?;

    match &input.data {
        Data::Struct(data_struct) => {
            expand_struct(input, data_struct, &struct_options, &borrow_lifetime)
        }
        Data::Enum(data_enum) => expand_enum(input, data_enum, &struct_options, &borrow_lifetime),
        Data::Union(data_union) => Err(Error::new(
            data_union.union_token.span,
            "`View` cannot be derived for unions",
        )),
    }
}

fn expand_struct(
    input: &DeriveInput,
    data_struct: &DataStruct,
    struct_options: &StructOptions,
    borrow_lifetime: &Lifetime,
) -> Result<TokenStream> {
    let named_fields = struct_named_fields(input, data_struct)?;
    let input_ident = &input.ident;
    let view_ident = view_ident(input_ident, struct_options);
    let view_visibility = struct_options
        .view_visibility
        .clone()
        .unwrap_or_else(|| input.vis.clone());
    let view_generics = view_generics(&input.generics, borrow_lifetime.clone());

    let mut view_fields = Vec::new();
    let mut generation_fields = Vec::new();

    for field in &named_fields.named {
        let field_ident = field_ident(field)?;
        let field_parts = struct_field_parts(field, struct_options, borrow_lifetime)?;
        let declaration_prefix = &field_parts.declaration_prefix;
        let borrowed_type = &field_parts.borrowed_type;
        let generation_expression = &field_parts.generation_expression;

        view_fields.push(quote! {
            #declaration_prefix #borrowed_type
        });
        generation_fields.push(quote! {
            #field_ident: #generation_expression
        });
    }

    let method = if struct_options.no_method {
        None
    } else {
        Some(method_tokens(
            input,
            struct_options,
            &view_ident,
            &generation_fields,
        ))
    };

    Ok(quote! {
        #view_visibility struct #view_ident #view_generics {
            #(#view_fields,)*
        }

        #method
    })
}

fn expand_enum(
    input: &DeriveInput,
    data_enum: &DataEnum,
    struct_options: &StructOptions,
    borrow_lifetime: &Lifetime,
) -> Result<TokenStream> {
    reject_enum_struct_options(struct_options)?;

    let input_ident = &input.ident;
    let view_ident = view_ident(input_ident, struct_options);
    let view_visibility = struct_options
        .view_visibility
        .clone()
        .unwrap_or_else(|| input.vis.clone());
    let view_generics = view_generics(&input.generics, borrow_lifetime.clone());

    let mut view_variants = Vec::new();
    let mut match_arms = Vec::new();

    for variant in &data_enum.variants {
        let variant_ident = &variant.ident;

        match &variant.fields {
            Fields::Named(fields) => {
                let mut field_idents = Vec::new();
                let mut view_fields = Vec::new();
                let mut generation_fields = Vec::new();

                for field in &fields.named {
                    let field_ident = field_ident(field)?;
                    let field_parts = enum_field_parts(field, borrow_lifetime)?;
                    let borrowed_type = &field_parts.borrowed_type;
                    let generation_expression = &field_parts.generation_expression;

                    field_idents.push(field_ident.clone());
                    view_fields.push(quote! {
                        #field_ident: #borrowed_type
                    });
                    generation_fields.push(quote! {
                        #field_ident: #generation_expression
                    });
                }

                view_variants.push(quote! {
                    #variant_ident {
                        #(#view_fields,)*
                    }
                });
                match_arms.push(quote! {
                    Self::#variant_ident { #(#field_idents),* } => #view_ident::#variant_ident {
                        #(#generation_fields,)*
                    }
                });
            }
            Fields::Unit => {
                view_variants.push(quote! {
                    #variant_ident
                });
                match_arms.push(quote! {
                    Self::#variant_ident => #view_ident::#variant_ident
                });
            }
            Fields::Unnamed(fields) => {
                return Err(Error::new(
                    fields.span(),
                    "`View` cannot be derived for tuple enum variants yet",
                ));
            }
        }
    }

    let method = if struct_options.no_method {
        None
    } else {
        Some(enum_method_tokens(
            input,
            struct_options,
            &view_ident,
            &match_arms,
        ))
    };

    Ok(quote! {
        #view_visibility enum #view_ident #view_generics {
            #(#view_variants,)*
        }

        #method
    })
}

fn struct_named_fields<'a>(
    input: &DeriveInput,
    data_struct: &'a DataStruct,
) -> Result<&'a syn::FieldsNamed> {
    match &data_struct.fields {
        Fields::Named(named_fields) => Ok(named_fields),
        Fields::Unnamed(fields) => Err(Error::new(
            fields.span(),
            "`View` can only be derived for structs with named fields",
        )),
        Fields::Unit => Err(Error::new(
            input.ident.span(),
            "`View` cannot be derived for unit structs",
        )),
    }
}

fn reject_enum_struct_options(options: &StructOptions) -> Result<()> {
    if let Some(field_visibility) = &options.field_visibility {
        return Err(Error::new(
            field_visibility.span(),
            "`field_visibility` cannot be used when deriving `View` for enums",
        ));
    }

    Ok(())
}

fn field_ident(field: &Field) -> Result<&Ident> {
    field
        .ident
        .as_ref()
        .ok_or_else(|| Error::new(field.span(), "expected a named field"))
}

fn struct_field_parts(
    field: &Field,
    struct_options: &StructOptions,
    borrow_lifetime: &Lifetime,
) -> Result<FieldParts> {
    let field_ident = field_ident(field)?;
    let field_options = attrs::parse_field_options(&field.attrs)?;
    let field_visibility = field_visibility(&field.vis, struct_options, &field_options);
    let default_mapping = mapping::field_mapping(
        &field.ty,
        borrow_lifetime,
        parse_quote!(self.#field_ident),
        FallbackBorrowMode::NeedsBorrow,
    );

    Ok(FieldParts {
        declaration_prefix: quote!(#field_visibility #field_ident:),
        borrowed_type: field_options
            .borrowed_type
            .clone()
            .unwrap_or(default_mapping.borrowed_type),
        generation_expression: field_options
            .generation_expression
            .clone()
            .unwrap_or(default_mapping.generation_expression),
    })
}

fn enum_field_parts(field: &Field, borrow_lifetime: &Lifetime) -> Result<FieldParts> {
    let field_ident = field_ident(field)?;
    let field_options = attrs::parse_field_options(&field.attrs)?;

    if let Some(visibility) = &field_options.visibility {
        return Err(Error::new(
            visibility.span(),
            "`visibility` cannot be used on enum variant fields",
        ));
    }

    let default_mapping = mapping::field_mapping(
        &field.ty,
        borrow_lifetime,
        parse_quote!(#field_ident),
        FallbackBorrowMode::MatchedBorrow,
    );

    Ok(FieldParts {
        declaration_prefix: quote!(#field_ident:),
        borrowed_type: field_options
            .borrowed_type
            .clone()
            .unwrap_or(default_mapping.borrowed_type),
        generation_expression: field_options
            .generation_expression
            .clone()
            .unwrap_or(default_mapping.generation_expression),
    })
}

struct FieldParts {
    declaration_prefix: TokenStream,
    borrowed_type: syn::Type,
    generation_expression: syn::Expr,
}

fn view_ident(input_ident: &Ident, options: &StructOptions) -> Ident {
    options
        .view_name
        .clone()
        .unwrap_or_else(|| format_ident!("{input_ident}View"))
}

fn field_visibility(
    source_visibility: &Visibility,
    struct_options: &StructOptions,
    field_options: &FieldOptions,
) -> Visibility {
    field_options
        .visibility
        .clone()
        .or_else(|| struct_options.field_visibility.clone())
        .unwrap_or_else(|| source_visibility.clone())
}

fn method_tokens(
    input: &DeriveInput,
    options: &StructOptions,
    view_ident: &Ident,
    generation_fields: &[TokenStream],
) -> TokenStream {
    let input_ident = &input.ident;
    let method_visibility = options
        .method_visibility
        .clone()
        .unwrap_or_else(|| input.vis.clone());
    let method_ident = options
        .method_name
        .clone()
        .unwrap_or_else(|| Ident::new("view", input_ident.span()));
    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    let view_type_arguments = view_type_arguments(&input.generics);

    quote! {
        impl #impl_generics #input_ident #type_generics #where_clause {
            #method_visibility fn #method_ident(&self) -> #view_ident<#(#view_type_arguments),*> {
                #view_ident {
                    #(#generation_fields,)*
                }
            }
        }
    }
}

fn enum_method_tokens(
    input: &DeriveInput,
    options: &StructOptions,
    view_ident: &Ident,
    match_arms: &[TokenStream],
) -> TokenStream {
    let input_ident = &input.ident;
    let method_visibility = options
        .method_visibility
        .clone()
        .unwrap_or_else(|| input.vis.clone());
    let method_ident = options
        .method_name
        .clone()
        .unwrap_or_else(|| Ident::new("view", input_ident.span()));
    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    let view_type_arguments = view_type_arguments(&input.generics);

    quote! {
        impl #impl_generics #input_ident #type_generics #where_clause {
            #method_visibility fn #method_ident(&self) -> #view_ident<#(#view_type_arguments),*> {
                match self {
                    #(#match_arms,)*
                }
            }
        }
    }
}

fn view_generics(source_generics: &Generics, borrow_lifetime: Lifetime) -> Generics {
    let mut view_generics = source_generics.clone();
    view_generics.params.insert(
        0,
        GenericParam::Lifetime(LifetimeParam::new(borrow_lifetime)),
    );
    view_generics
}

fn view_type_arguments(source_generics: &Generics) -> Vec<TokenStream> {
    let mut arguments = vec![quote!('_)];

    arguments.extend(
        source_generics
            .params
            .iter()
            .map(|parameter| match parameter {
                GenericParam::Lifetime(lifetime) => {
                    let lifetime = &lifetime.lifetime;
                    quote!(#lifetime)
                }
                GenericParam::Type(type_parameter) => {
                    let ident = &type_parameter.ident;
                    quote!(#ident)
                }
                GenericParam::Const(const_parameter) => {
                    let ident = &const_parameter.ident;
                    quote!(#ident)
                }
            }),
    );

    arguments
}

/// Return the stable generated view lifetime, rejecting user lifetime collisions.
///
/// The generated view type always uses `'borrowize`; `borrowed_type`
/// overrides can refer to that same lifetime directly. A source struct that
/// already defines `'borrowize` would make that contract ambiguous, so the
/// derive fails instead of choosing a fresh name.
fn generated_borrow_lifetime(generics: &Generics) -> Result<Lifetime> {
    if let Some(existing_lifetime) = generics
        .lifetimes()
        .find(|lifetime| lifetime.lifetime.ident == GENERATED_BORROW_LIFETIME)
    {
        return Err(Error::new(
            existing_lifetime.lifetime.span(),
            "`View` reserves the lifetime name `'borrowize` for the generated view borrow",
        ));
    }

    Ok(Lifetime::new("'borrowize", proc_macro2::Span::call_site()))
}