bauer-macros 0.6.0

use std::collections::HashMap;

use proc_macro2::TokenStream;
use quote::{ToTokens, format_ident, quote, quote_spanned};
use syn::{
    DeriveInput, Ident, Token, TraitBound, Type, TypePath, parse_quote, parse_quote_spanned,
    spanned::Spanned,
};

use crate::{
    BuilderAttr, BuilderField, Len, Repeat,
    attr::field::{FieldIdents, WrappedType},
    builder_args,
    type_state::generics::{CustomImplGenerics, CustomTypeGenerics},
    util::{ReplaceTrait, ensure_no_conflict, known_idents, parallel_assign},
};

mod generics;

fn build_fn(
    builder: &Ident,
    builder_attr: &BuilderAttr,
    fields: &[BuilderField],
    generic_fields: &[&BuilderField],
    len_structs: &HashMap<usize, TraitBound>,
    input: &DeriveInput,
    inner: &Ident,
) -> TokenStream {
    let ident = &input.ident;
    let private_module = builder_attr.private_module();

    let names = fields.iter().map(|f| &f.ident);

    let not_skipped_field_values = fields.iter().filter(|f| !f.should_skip()).map(|field| {
        let name = &field.ident;
        let wrapped_ty = &field.ty;
        let pascal = &field.idents.pascal;
        let field_i = field.tuple_index();

        if field.is_associated() {
            return quote! {{
                let #name: #wrapped_ty = inner.#field_i;
                #name
            }};
        };

        let value = match field.wrapped_ty {
            WrappedType::None => {
                if let Some(default) = &field.attr.default {
                    let default = default.to_value(field.attr.into);

                    quote_spanned! {field.ty.span()=>
                        // TODO: make this a function once const traits are stable
                        if <#pascal as #private_module::state::BuilderState>::SET {
                            // SAFETY: If #pascal::SET is true, then we have already set #field_i
                            unsafe { inner.#field_i.assume_init() }
                        } else {
                            #default
                        }
                    }
                } else {
                    quote! {
                        // SAFETY: This function is only accessible if all required fields are set.  This
                        // is enusred by the type bounds.
                        unsafe { inner.#field_i.assume_init() }
                    }
                }
            }
            WrappedType::Flag => quote! { inner.#field_i },
            WrappedType::Option(_) => quote! { inner.#field_i },
            WrappedType::Repeat(
                ref inner_ty,
                Repeat {
                    len: Len::Int { .. },
                    array,
                    ref collector,
                },
            ) => {
                let value = if array {
                    quote! { array }
                } else {
                    assert!(!builder_attr.konst);
                    assert!(!array);
                    collector.collect(parse_quote_spanned! { inner_ty.span()=>
                        array.into_iter()
                    })
                };

                quote_spanned! {
                    inner_ty.span() =>
                    {
                        // SAFETY: The build function can only be called once the array has been
                        // filled and is fully initialised.
                        let array = unsafe { inner.#field_i.assume_init() };
                        #value
                    }
                }
            }
            WrappedType::Repeat(
                ref inner_ty,
                Repeat {
                    array,
                    ref collector,
                    ..
                },
            ) => {
                assert!(!builder_attr.konst);
                assert!(!array);
                collector.collect(parse_quote_spanned! { inner_ty.span()=> {
                    let _: &::std::vec::Vec<_> = &inner.#field_i; // assert that the types are correct
                    inner.#field_i.into_iter()
                }})
            }
        };

        quote! {{
            let #name: #wrapped_ty = #value; // assert the type
            #name
        }}
    });

    let not_skipped_fields: Vec<_> = fields
        .iter()
        .filter(|f| !f.should_skip())
        .map(|f| &f.ident)
        .collect();

    let set_not_skipped_fields = parallel_assign(
        not_skipped_fields.iter().copied(),
        not_skipped_field_values,
        quote! {
            let inner = self.#inner;
        },
    );

    let set_skipped_fields = parallel_assign(
        fields.iter().filter(|f| f.should_skip()).map(|f| &f.ident),
        fields.iter().filter_map(BuilderField::skipped_field_value),
        quote! {
            #[allow(unused)]
            let (#(#not_skipped_fields),*) = (#(&#not_skipped_fields),*);
        },
    );

    let build_impl_generics = generic_fields.iter().enumerate().filter_map(|(i, f)| {
        let pascal = &f.idents.pascal;
        if f.optional() {
            Some(quote! {
                #pascal: #private_module::state::BuilderState
            })
        } else if f.optional() || len_structs.contains_key(&i) {
            Some(quote! { #pascal })
        } else {
            None
        }
    });

    let build_generics = generic_fields.iter().enumerate().map(|(i, f)| {
        let FieldIdents {
            count, pascal, set, ..
        } = &f.idents;
        let ty: Type = if len_structs.contains_key(&i) {
            parse_quote! { #count<#pascal> }
        } else if f.optional() {
            parse_quote! { #pascal }
        } else {
            parse_quote! { #set<true> }
        };
        ty
    });

    let impl_generics = CustomImplGenerics::new(&input.generics, build_impl_generics);
    let ty_generics = CustomTypeGenerics::new(&input.generics, build_generics);

    let (_, _, where_clause) = input.generics.split_for_impl();

    let mut builder_where = where_clause.to_token_stream();
    if let Some(where_clause) = where_clause {
        if !where_clause.predicates.trailing_punct() {
            <Token![,]>::default().to_tokens(&mut builder_where);
        }
    } else {
        <Token![where]>::default().to_tokens(&mut builder_where);
    }

    for (&i, len) in len_structs {
        let FieldIdents { pascal, count, .. } = &generic_fields[i].idents;
        builder_where.extend(quote! {
            #count<#pascal>: #len,
        });
    }

    let (_, default_ty_generics, _) = input.generics.split_for_impl();

    let konst = builder_attr.konst_kw();

    let build_fn_attributes = &builder_attr.build_fn.attributes;
    let build_fn_name = &builder_attr.build_fn.name;

    let mut ret_ty = quote! { #ident #default_ty_generics };
    let mut ret_val = quote! { val };

    if let Some((param, ty, body)) = &builder_attr.build_fn.mapper {
        ret_val = quote! {
            #[allow(unused_braces)]
            #konst fn __private_mapper(#param: #ret_ty) -> #ty {
                #body
            }
            __private_mapper(val)
        };
        ret_ty = ty.to_token_stream();
    };

    let mut from = if builder_attr.error.force {
        ret_ty = quote! { ::core::result::Result<#ret_ty, ::core::convert::Infallible> };
        ret_val = quote! { Ok(#ret_val) };
        Some(quote! {
            let Ok(built) = builder.#build_fn_name();
            built
        })
    } else {
        Some(quote! { builder.#build_fn_name() })
    };

    if builder_attr.build_fn.mapper.is_some() {
        from = None;
    }

    let from = from.map(|from| quote! {
        impl #impl_generics ::core::convert::From<#builder #ty_generics> for #ident #default_ty_generics #builder_where {
            fn from(builder: #builder #ty_generics) -> Self  {
                #from
            }
        }
    });

    let vis = builder_attr.build_fn.vis(builder_attr);

    quote! {
        impl #impl_generics #builder #ty_generics #builder_where {
            #(#build_fn_attributes)*
            #vis #konst fn #build_fn_name(self) -> #ret_ty {
                #[allow(deprecated)] // #inner is set to deprecated
                let val = {
                    #set_not_skipped_fields
                    #set_skipped_fields
                    #ident {
                        #(#names),*
                    }
                };
                #ret_val
            }
        }

        #from
    }
}

pub fn type_state_builder(
    builder_attr: &BuilderAttr,
    input: &DeriveInput,
    mut fields: Vec<BuilderField>,
) -> TokenStream {
    let ident = &input.ident;
    let builder = format_ident!("{}Builder", ident);
    let builder_vis = &builder_attr.vis;

    let known_idents = known_idents(input, &fields);
    for f in &mut fields {
        ensure_no_conflict(&mut f.idents.pascal, &known_idents);
        ensure_no_conflict(&mut f.idents.set, &known_idents);
        ensure_no_conflict(&mut f.idents.count, &known_idents);
    }

    let generic_fields: Vec<_> = fields
        .iter()
        .filter(|f| {
            !f.should_skip()
                && !f.is_associated()
                && f.wrapped_ty.repeat().is_none_or(|(_, r)| r.len.is_some())
        })
        .collect();

    let mut out = TokenStream::new();

    builder_attr.assert_crate().to_tokens(&mut out);

    let private_module = builder_attr.private_module();
    out.extend(generic_fields.iter().map(|&f| {
        let FieldIdents { count, set, .. } = &f.idents;
        if f.wrapped_ty.repeat().is_some_and(|(_, r)| r.len.is_some()) {
            quote! {
                #[doc(hidden)]
                #[allow(non_camel_case_types)]
                #[non_exhaustive]
                #builder_vis struct #count<T>(T); // never constructed, so doesn't really need to be PhantomData
            }
        } else {
            quote! {
                #[doc(hidden)]
                #[allow(non_camel_case_types)]
                #[non_exhaustive]
                #builder_vis struct #set<const SET: bool>;
                impl<const SET: bool> #private_module::sealed::Sealed for #set<SET> {}
                impl<const SET: bool> #private_module::state::BuilderState for #set<SET> {
                    const SET: bool = SET;
                }
            }
        }
    }));

    let mut len_structs = HashMap::new();

    for (i, &f) in generic_fields.iter().enumerate() {
        let WrappedType::Repeat(_, repeat) = &f.wrapped_ty else {
            continue;
        };

        if repeat.len.is_none() {
            continue;
        }

        let ident = match repeat
            .len
            .to_trait(&builder_attr.krate, &f.idents.count, &mut out)
        {
            Ok(i) => i,
            Err(e) => return e.to_compile_error(),
        };
        len_structs.insert(i, ident);
    }

    let (default_impl_generics, default_ty_generics, where_clause) =
        input.generics.split_for_impl();

    let field_decls = fields.iter().filter(|f| !f.should_skip()).map(|f| {
        if f.is_associated() {
            return f.ty.to_token_stream();
        }

        match &f.wrapped_ty {
            WrappedType::None => {
                let ty = &f.ty;
                quote! { ::core::mem::MaybeUninit<#ty> }
            }
            WrappedType::Flag => {
                quote! { bool }
            }
            WrappedType::Option(ty) => quote! { ::core::option::Option<#ty> },
            WrappedType::Repeat(
                ty,
                Repeat {
                    len: Len::Int { len },
                    ..
                },
            ) => quote! { ::core::mem::MaybeUninit<[#ty; #len]> },
            WrappedType::Repeat(ty, _) => quote! { ::std::vec::Vec<#ty> },
        }
    });

    let init = fields.iter().filter(|f| !f.should_skip()).map(|f| {
        if f.is_associated() {
            let (_, value) = f.attr.to_args_and_value(&f.ty, f.arg_name());
            return value;
        }

        match f.wrapped_ty {
            WrappedType::None => quote! { ::core::mem::MaybeUninit::uninit() },
            WrappedType::Flag => quote! { false },
            WrappedType::Option(_) => quote! { ::core::option::Option::None },
            WrappedType::Repeat(
                _,
                Repeat {
                    len: Len::Int { .. },
                    ..
                },
            ) => quote! { ::core::mem::MaybeUninit::uninit() },
            WrappedType::Repeat(_, _) => quote! { ::std::vec::Vec::new() },
        }
    });

    let inner = format_ident!("__unsafe_builder_content");
    let state = format_ident!("__unsafe_builder_state");

    let phantom_generics = generic_fields.iter().map(|f| &f.idents.pascal);
    let phantom = quote! {
        #state: ::core::marker::PhantomData<(#(#phantom_generics,)*)>
    };

    let new_generics = generic_fields.iter().map(|f| {
        let FieldIdents { count, set, .. } = &f.idents;
        if f.wrapped_ty.repeat().is_some_and(|(_, r)| r.len.is_some()) {
            quote! { #count<()> }
        } else {
            quote! { #set<false> }
        }
    });

    let struct_generics = CustomImplGenerics::new(
        &input.generics,
        generic_fields.iter().map(|f| &f.idents.pascal),
    );
    let new_generics = CustomTypeGenerics::new(&input.generics, new_generics);

    let konst = builder_attr.konst_kw();
    let attributes = &builder_attr.attributes;

    let builder_fn = {
        let vis = builder_attr.builder_fn.vis(builder_attr);
        let name = &builder_attr.builder_fn.name;
        let attributes = &builder_attr.builder_fn.attributes;
        let (names, args, _) = builder_args(&fields);
        quote! {
            impl #default_impl_generics #ident #default_ty_generics #where_clause {
                #(#attributes)*
                #vis #konst fn #name(#(#args),*) -> #builder #new_generics {
                    #builder::new(#(#names),*)
                }
            }
        }
    };

    let new_fn = {
        let (_, args, _) = builder_args(&fields);
        quote! {
            impl #default_impl_generics #builder #new_generics #where_clause {
                #konst fn new(#(#args),*) -> Self {
                    Self {
                        #inner: (#(#init,)*),
                        #state: ::core::marker::PhantomData,
                    }
                }
            }
        }
    };

    out.extend(quote! {
        #(#attributes)*
        #[allow(clippy::type_complexity)]
        #[must_use = "The builder doesn't construct its type until `.build()` is called"]
        #builder_vis struct #builder #struct_generics #where_clause {
            #[deprecated = "This field is for internal use only; you almost certainly don't need to touch this. If you encounter a bug or missing feature, file an issue on the repo."]
            #[doc(hidden)]
            #inner: (#(#field_decls,)*),
            #[deprecated = "This field is for internal use only; you almost certainly don't need to touch this. If you encounter a bug or missing feature, file an issue on the repo."]
            #[doc(hidden)]
            #phantom
        }

        #builder_fn
        #new_fn
    });

    // add `fn build()`
    out.extend(build_fn(
        &builder,
        builder_attr,
        &fields,
        &generic_fields,
        &len_structs,
        input,
        &inner,
    ));

    let mut i = 0;
    for f in &fields {
        if f.should_skip() || f.is_associated() {
            continue; // skip
        }

        let (args, value, value_ty) = if let WrappedType::Flag = f.wrapped_ty {
            (TokenStream::new(), quote! { true }, &parse_quote! { bool })
        } else {
            let ty = f.arg_ty().expect("None iff wrapped_ty != Flag");
            let (args, val) = f.attr.to_args_and_value(ty, &f.ident);
            (args, val, ty)
        };
        let fn_ident = f.function_ident(builder_attr);

        fn ident_to_type(ident: &Ident) -> Type {
            Type::Path(TypePath {
                attrs: Vec::new(),
                qself: None,
                path: ident.clone().into(),
            })
        }

        let fn_attributes = &f.attr.attributes;

        let field_i = f.tuple_index();
        let vis = &f.attr.vis;

        let fun = match f.wrapped_ty {
            WrappedType::None | WrappedType::Option(_) | WrappedType::Flag => {
                // NOTE: be sure to update match below
                let impl_generics_fields = CustomImplGenerics::new(
                    &input.generics,
                    generic_fields[..i]
                        .iter()
                        .chain(generic_fields[i + 1..].iter())
                        .map(|f| &f.idents.pascal),
                );

                let FieldIdents { set, .. } = &generic_fields[i].idents;
                let struct_generics_fields = CustomTypeGenerics::new(
                    &input.generics,
                    generic_fields
                        .iter()
                        .map(|f| ident_to_type(&f.idents.pascal))
                        .replace(i, parse_quote! { #set<false> }),
                );

                let return_struct_generics_fields = CustomTypeGenerics::new(
                    &input.generics,
                    generic_fields
                        .iter()
                        .map(|f| ident_to_type(&f.idents.pascal))
                        .replace(i, parse_quote! { #set<true> }),
                );

                let setter = match f.wrapped_ty {
                    WrappedType::None => quote! {
                        this.#inner.#field_i.write(value);
                    },
                    WrappedType::Flag => quote! {
                        this.#inner.#field_i = value;
                    },
                    WrappedType::Option(_) => quote! {
                        this.#inner.#field_i = Some(value);
                    },
                    _ => unreachable!("Checked above"),
                };

                quote_spanned! {
                    fn_ident.span() =>
                    impl #impl_generics_fields #builder #struct_generics_fields #where_clause {
                        #(#fn_attributes)*
                        #[allow(clippy::type_complexity)]
                        #vis #konst fn #fn_ident(self, #args) -> #builder #return_struct_generics_fields {
                            let value: #value_ty = #value;
                            let mut this = self; // rather than have `mut self` in the signature
                            #[allow(deprecated)] // #inner is set to deprecated
                            {
                                #setter
                                #builder {
                                    #inner: this.#inner,
                                    #state: ::core::marker::PhantomData,
                                }
                            }
                        }
                    }
                }
            }
            WrappedType::Repeat(_, Repeat { len: Len::None, .. }) => {
                let impl_generics = CustomImplGenerics::new(
                    &input.generics,
                    generic_fields.iter().map(|f| &f.idents.pascal),
                );
                let ty_generics = CustomTypeGenerics::new(
                    &input.generics,
                    generic_fields.iter().map(|f| &f.idents.pascal),
                );
                quote_spanned! {
                    fn_ident.span() =>
                    impl #impl_generics #builder #ty_generics #where_clause {
                        #(#fn_attributes)*
                        #[allow(clippy::type_complexity)]
                        #vis #konst fn #fn_ident(self, #args) -> #builder #ty_generics {
                            let value: #value_ty = #value;
                            let mut this = self; // rather than have `mut self` in the signature
                            #[allow(deprecated)] // #inner is set to deprecated
                            {
                                this.#inner.#field_i.push(value);
                                #builder {
                                    #inner: this.#inner,
                                    #state: ::core::marker::PhantomData,
                                }
                            }
                        }
                    }
                }
            }
            WrappedType::Repeat(_, Repeat { ref len, .. }) => {
                let FieldIdents { count, pascal, .. } = &generic_fields[i].idents;

                let impl_generics = CustomImplGenerics::new(
                    &input.generics,
                    generic_fields.iter().map(|f| &f.idents.pascal),
                );
                let ty_generics = CustomTypeGenerics::new(
                    &input.generics,
                    generic_fields
                        .iter()
                        .map(|f| ident_to_type(&f.idents.pascal))
                        .replace(i, parse_quote! { #count<#pascal> }),
                );

                let ret_ty_generics = CustomTypeGenerics::new(
                    &input.generics,
                    generic_fields
                        .iter()
                        .map(|f| ident_to_type(&f.idents.pascal))
                        .replace(i, parse_quote! { #count<(#pascal, ())> }),
                );

                let mut field_where = where_clause.to_token_stream();
                if let Some(where_clause) = where_clause {
                    if !where_clause.predicates.trailing_punct() {
                        <Token![,]>::default().to_tokens(&mut field_where);
                    }
                } else {
                    <Token![where]>::default().to_tokens(&mut field_where);
                }

                let add = if let Len::Int { .. } = len {
                    field_where.extend(quote! {
                        #private_module::state::Count::<#pascal>: #private_module::state::Countable,
                    });
                    quote! {
                        let ptr = this.#inner.#field_i.as_mut_ptr();
                        // SAFETY: ptr points to a valid location created by the MaybeUninit
                        unsafe {
                            let ptr: *mut _ = &raw mut (*ptr)[<#private_module::state::Count::<#pascal> as #private_module::state::Countable>::COUNT];
                            ptr.write(value);
                        }
                    }
                } else {
                    quote! { this.#inner.#field_i.push(value); }
                };

                quote_spanned! {
                    fn_ident.span() =>
                    impl #impl_generics #builder #ty_generics #field_where {
                        #(#fn_attributes)*
                        #[allow(clippy::type_complexity)]
                        #vis #konst fn #fn_ident(self, #args) -> #builder #ret_ty_generics {
                            let value: #value_ty = #value;
                            let mut this = self; // rather than have `mut self` in the signature
                            #[allow(deprecated)] // #inner is set to deprecated
                            {
                                #add
                                #builder {
                                    #inner: this.#inner,
                                    #state: ::core::marker::PhantomData,
                                }
                            }
                        }
                    }
                }
            }
        };

        out.extend(fun);
        if f.wrapped_ty.repeat().is_none_or(|(_, r)| r.len.is_some()) {
            i += 1;
        }
    }

    out
}