path: root/mingling_macros/src/attr/completion.rs

// Doc Not Optimize
use crate::res_injection::{ResourceInjection, generate_immut_resource_bindings};
use proc_macro::TokenStream;
use quote::quote;
use syn::spanned::Spanned;
use syn::{FnArg, Ident, ItemFn, Pat, PatType, Type, TypePath, parse_macro_input};

#[cfg(feature = "comp")]
#[allow(clippy::too_many_lines)]
pub(crate) fn completion_attr(attr: TokenStream, item: TokenStream) -> TokenStream {
    use crate::get_global_set;

    let previous_type_path: TypePath = if attr.is_empty() {
        return syn::Error::new(
            proc_macro2::Span::call_site(),
            "completion attribute requires a previous type argument, e.g. #[completion(HelloEntry)]",
        )
        .to_compile_error()
        .into();
    } else {
        parse_macro_input!(attr as TypePath)
    };
    let previous_type_ident = &previous_type_path.path.segments.last().unwrap().ident;

    let input_fn = parse_macro_input!(item as ItemFn);

    if input_fn.sig.asyncness.is_some() {
        return syn::Error::new(input_fn.sig.span(), "Completion function cannot be async")
            .to_compile_error()
            .into();
    }

    let sig = &input_fn.sig;
    let inputs = &sig.inputs;
    let output = &sig.output;

    // Parameter classification:
    // - Owned (non-reference) parameters are **Shell sources**: each is derived
    //   from `&ShellContext` via `From<&ShellContext>` (identity `From` covers
    //   `ShellContext` itself).
    // - `&T` / `&mut T` parameters are resource injections (unchanged).
    // - `&ShellContext` is rejected with a helpful message: use the owned
    //   `ShellContext` (or any other `From<&ShellContext>` type) instead.
    let mut derived_stmts: Vec<proc_macro2::TokenStream> = Vec::new();
    let mut call_args: Vec<proc_macro2::TokenStream> = Vec::new();
    let mut resources = Vec::new();

    for (idx, arg) in inputs.iter().enumerate() {
        match arg {
            FnArg::Typed(PatType { pat, ty, .. }) => {
                if let Type::Reference(ref_type) = &**ty {
                    // `&ShellContext` is no longer allowed: it clashes with the
                    // resource-injection semantics of references.
                    if is_shell_context_path(&ref_type.elem) {
                        return syn::Error::new(
                            ty.span(),
                            "`&ShellContext` is not supported; use the owned `ShellContext` \
                             (or any other type implementing `From<&ShellContext>`) as a value \
                             parameter",
                        )
                        .to_compile_error()
                        .into();
                    }

                    // Reference: resource injection (requires a named binding).
                    let var_name = match &**pat {
                        Pat::Ident(pat_ident) => pat_ident.ident.clone(),
                        _ => {
                            return syn::Error::new(
                                pat.span(),
                                "Resource injection parameter must be a simple identifier",
                            )
                            .to_compile_error()
                            .into();
                        }
                    };

                    // Reference: resource injection.
                    let (inner_type, is_mut) = match &*ref_type.elem {
                        Type::Path(type_path) => {
                            let is_mut = ref_type.mutability.is_some();
                            (type_path.clone(), is_mut)
                        }
                        _ => {
                            return syn::Error::new(
                                ty.span(),
                                "Reference resource type must be a type path",
                            )
                            .to_compile_error()
                            .into();
                        }
                    };
                    resources.push(ResourceInjection {
                        var_name: var_name.clone(),
                        full_type: (**ty).clone(),
                        inner_type,
                        is_ref: true,
                        is_mut,
                    });
                    call_args.push(quote! { #var_name });
                } else {
                    // Owned value: derive from `&ShellContext`. The parameter
                    // name is irrelevant (anonymous `_` is fine) since the
                    // derived binding is generated by this macro.
                    let derived_ident = Ident::new(&format!("__ctx_derived_{idx}"), pat.span());
                    derived_stmts.push(quote! {
                        let #derived_ident: #ty =
                            <#ty as ::std::convert::From<&::mingling::ShellContext>>::from(ctx);
                    });
                    call_args.push(quote! { #derived_ident });
                }
            }
            FnArg::Receiver(_) => {
                return syn::Error::new(
                    arg.span(),
                    "Completion function cannot have self parameter",
                )
                .to_compile_error()
                .into();
            }
        }
    }

    let fn_body = &input_fn.block;

    let mut fn_attrs = input_fn.attrs.clone();
    fn_attrs.retain(|attr| !attr.path().is_ident("completion"));

    let vis = &input_fn.vis;
    let fn_name = &sig.ident;

    let internal_name = format!(
        "__internal_completion_{}",
        just_fmt::snake_case!(fn_name.to_string())
    );
    let struct_name = Ident::new(&internal_name, fn_name.span());

    let program_type = crate::default_program_path();
    let mut_resources: Vec<_> = resources.iter().filter(|r| r.is_mut).collect();

    let immut_resource_stmts = generate_immut_resource_bindings(resources.iter(), &program_type);

    let fn_call = quote! { #fn_name(#(#call_args),*) };

    let inner_call = if mut_resources.is_empty() {
        fn_call
    } else {
        let mut wrapped = fn_call;
        for res in mut_resources.iter().rev() {
            let var_name = &res.var_name;
            let inner_type = &res.inner_type;
            wrapped = quote! {
                ::mingling::this::<#program_type>().modify_res(|#var_name: &mut #inner_type| {
                    #wrapped
                })
            };
        }
        wrapped
    };

    let comp_body = quote! {
        #(#derived_stmts)*
        #(#immut_resource_stmts)*
        #inner_call
    };

    // A `()` return (or no return type) means "no suggestions": map it to an
    // empty `Suggest` instead of requiring `Into<Suggest>`.
    let returns_unit = match &sig.output {
        syn::ReturnType::Default => true,
        syn::ReturnType::Type(_, ty) => {
            matches!(ty.as_ref(), syn::Type::Tuple(t) if t.elems.is_empty())
        }
    };
    let return_stmt = if returns_unit {
        quote! {
            { #comp_body };
            ::mingling::Suggest::new()
        }
    } else {
        quote! {
            let __completion_result = { #comp_body };
            ::std::convert::Into::into(__completion_result)
        }
    };

    let expanded: proc_macro2::TokenStream = quote! {
        #(#fn_attrs)*
        #[doc(hidden)]
        #[allow(non_camel_case_types)]
        #vis struct #struct_name;

        impl ::mingling::Completion for #struct_name {
            type Previous = #previous_type_path;

            fn comp(ctx: &::mingling::ShellContext) -> ::mingling::Suggest {
                #return_stmt
            }
        }

        // Keep the original function for internal use
        #(#fn_attrs)*
        #vis fn #fn_name(#inputs) #output {
            #fn_body
        }
    };

    let completion_entry = quote! {
        Self::#previous_type_ident => <#struct_name as ::mingling::Completion>::comp(ctx),
    };

    let completion_str = completion_entry.to_string();
    let variant_name = previous_type_ident.to_string();
    let span = previous_type_path.span();

    let mut completions = get_global_set(&crate::COMPLETIONS).lock().unwrap();
    if let Err(err) = crate::check_duplicate_variant(
        &completions,
        &completion_str,
        &variant_name,
        "completion",
        span,
    ) {
        return err.into();
    }
    completions.insert(completion_str);
    drop(completions);

    expanded.into()
}

/// Returns `true` when the type is a path whose last segment is `ShellContext`
/// (e.g. `ShellContext` or `mingling::ShellContext`).
fn is_shell_context_path(ty: &Type) -> bool {
    if let Type::Path(type_path) = ty {
        type_path
            .path
            .segments
            .last()
            .is_some_and(|seg| seg.ident == "ShellContext")
    } else {
        false
    }
}