diff options
Diffstat (limited to 'mingling_macros/src/attr/completion.rs')
| -rw-r--r-- | mingling_macros/src/attr/completion.rs | 250 |
1 files changed, 250 insertions, 0 deletions
diff --git a/mingling_macros/src/attr/completion.rs b/mingling_macros/src/attr/completion.rs new file mode 100644 index 0000000..3ced091 --- /dev/null +++ b/mingling_macros/src/attr/completion.rs @@ -0,0 +1,250 @@ +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")] +pub(crate) fn completion_attr(attr: TokenStream, item: TokenStream) -> TokenStream { + // Parse the attribute arguments such as HelloEntry or crate::EntryFine from #[completion(crate::EntryFine)] + 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; + + // Parse the function item + let input_fn = parse_macro_input!(item as ItemFn); + + // Validate the function is not async + if input_fn.sig.asyncness.is_some() { + return syn::Error::new(input_fn.sig.span(), "Completion function cannot be async") + .to_compile_error() + .into(); + } + + // Get the function signature parts + let sig = &input_fn.sig; + let inputs = &sig.inputs; + let output = &sig.output; + + // Must have at least one parameter ctx + if inputs.is_empty() { + return syn::Error::new( + inputs.span(), + "Completion function must have at least one parameter: `ctx: &ShellContext`", + ) + .to_compile_error() + .into(); + } + + // Extract the first param pattern and type for the ctx parameter + let first_arg = &inputs[0]; + let _ctx_type = match first_arg { + FnArg::Typed(PatType { ty, .. }) => (**ty).clone(), + FnArg::Receiver(_) => { + return syn::Error::new( + first_arg.span(), + "Completion function cannot have self parameter", + ) + .to_compile_error() + .into(); + } + }; + let fixed_ctx: Pat = syn::parse_quote!(ctx); + + // Extract resources from params 2 through N, skipping ctx + let resources = match extract_resources_from_args(sig, 1) { + Ok(r) => r, + Err(e) => return e.to_compile_error().into(), + }; + + // Get the function body + let fn_body = &input_fn.block; + + // Get function attributes excluding the completion attribute + let mut fn_attrs = input_fn.attrs.clone(); + fn_attrs.retain(|attr| !attr.path().is_ident("completion")); + + // Get function visibility + let vis = &input_fn.vis; + + // Get function name + let fn_name = &sig.ident; + + // Generate internal name from function name using snake_case + 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 has_resources = !resources.is_empty(); + let mut_resources: Vec<_> = resources.iter().filter(|r| r.is_mut).collect(); + + // Generate immutable resource bindings + let immut_resource_stmts = generate_immut_resource_bindings(resources.iter(), &program_type); + + // Build the call to the original function with resource arguments injected + let resource_args: Vec<_> = resources + .iter() + .map(|res| { + let var_name = &res.var_name; + quote! { #var_name } + }) + .collect(); + + let fn_call = if has_resources { + quote! { #fn_name(#fixed_ctx, #(#resource_args),*) } + } else { + quote! { #fn_name(#fixed_ctx) } + }; + + // Wrap the function call with modify_res for mutable resources + 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 = if has_resources { + quote! { + #(#immut_resource_stmts)* + #inner_call + } + } else { + quote! { #inner_call } + }; + + // Generate the struct and implementation + // The `comp` trait method only takes `ctx` as the first parameter; resources are injected internally + + let expanded = 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(#fixed_ctx: &::mingling::ShellContext) #output { + #comp_body + } + } + + // 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 mut completions = get_global_set(&crate::COMPLETIONS).lock().unwrap(); + let completion_str = completion_entry.to_string(); + + // Check for duplicate variant before inserting + let variant_name = previous_type_ident.to_string(); + if let Err(err) = crate::check_duplicate_variant( + &completions, + &completion_str, + &variant_name, + "completion", + previous_type_path.span(), + ) { + return err.into(); + } + + completions.insert(completion_str); + + expanded.into() +} + +/// Extract resource injection parameters from function arguments (skipping the first N params). +fn extract_resources_from_args( + sig: &syn::Signature, + skip: usize, +) -> syn::Result<Vec<ResourceInjection>> { + let mut resources = Vec::new(); + for arg in sig.inputs.iter().skip(skip) { + match arg { + FnArg::Typed(PatType { pat, ty, .. }) => { + let var_name = match &**pat { + Pat::Ident(pat_ident) => pat_ident.ident.clone(), + _ => { + return Err(syn::Error::new( + pat.span(), + "Resource injection parameter must be a simple identifier", + )); + } + }; + + let full_type = *(*ty).clone(); + + let (inner_type, is_ref, is_mut) = match &full_type { + Type::Reference(ref_type) => match &*ref_type.elem { + Type::Path(type_path) => { + let is_mut = ref_type.mutability.is_some(); + (type_path.clone(), true, is_mut) + } + _ => { + return Err(syn::Error::new( + ty.span(), + "Reference resource type must be a type path", + )); + } + }, + Type::Path(_) => { + return Err(syn::Error::new( + ty.span(), + "Resource injection parameter must be a reference (`&T` or `&mut T`)", + )); + } + _ => { + return Err(syn::Error::new( + ty.span(), + "Resource injection type must be a type path or reference", + )); + } + }; + + resources.push(ResourceInjection { + var_name, + full_type, + inner_type, + is_ref, + is_mut, + }); + } + FnArg::Receiver(_) => { + return Err(syn::Error::new( + arg.span(), + "Resource injection parameter cannot be self", + )); + } + } + } + Ok(resources) +} |