kstool-helper-generator 0.7.1

use proc_macro2::TokenStream;
use syn::{DataEnum, Variant, Visibility, spanned::Spanned};

use proc_macro2::Ident;
use quote::{ToTokens, quote};

use crate::basic::{EnumDefinition, do_expand, extract_enum};

/* fn check_is_enum(input: &TokenStream) -> bool {
    for token in input.clone().into_iter() {
        if let proc_macro2::TokenTree::Ident(ident) = token {
            if ident.eq("enum") {
                return true;
            }
        }
    }
    return false;
}
 */
/*/// Extract visibility and name
 fn header(input: TokenStream) -> syn::Result<(Vec<TokenTree>, TokenTree, TokenTree)> {
    let mut v = Vec::new();
    let span = input.span();
    let mut iter = input.into_iter();
    while let Some(token) = iter.next() {
        if let proc_macro2::TokenTree::Ident(ident) = &token {
            if ident.eq("pub") {
                v.push(token.clone());
                break;
            }
        }
    }

    while let Some(token) = iter.next() {
        if let proc_macro2::TokenTree::Ident(ident) = &token {
           //eprintln!("ident: {}", token);
            if ident.eq("enum") {
                break;
            }
        }
        v.push(token.clone());
    }
    let name = iter
        .next()
        .ok_or_else(|| syn::Error::new(span, "Should have a name"))?;

    Ok((
        v,
        name,
        iter.next()
            .ok_or_else(|| syn::Error::new(span, "Should have body"))?,
    ))
}

fn parse_group(input: Group) -> syn::Result<TokenStream> {
    for input in input.stream() {
       //eprintln!("{:?}", input);
    }
    Ok(Default::default())
}

fn build_tree(input: TokenStream) -> syn::Result<TokenStream> {
    for token in input.clone().into_iter() {
       //eprintln!("{:?}", token);
    }
    if !check_is_enum(&input) {
        return Err(syn::Error::new_spanned(input, "Not enum"));
    }
    let (visibility, name, group) = header(input.clone())?;
   //eprintln!("{:?} {}", visibility, name);

    if let proc_macro2::TokenTree::Group(group) = group {
       //eprintln!("--------------------------------");
        parse_group(group);
    } else {
        return Err(syn::Error::new_spanned(
            input,
            "Missing group after enum name",
        ));
    }

    Ok(Default::default())
}
 */

/// Generates async and/or blocking send methods for a variant without a return type.
///
/// Identical to the basic helper's function generation: sends the variant through the
/// MPSC channel and returns `Option<()>`.
fn generate_normal_function(
    basic: &Ident,
    variant: &Variant,
    block: bool,
    no_async: bool,
    vis: &Visibility,
) -> syn::Result<TokenStream> {
    let mut ret = TokenStream::new();
    let definition = EnumDefinition::try_from(variant)?;
    let arg_def = definition.fields().to_arg_def(variant.span());
    let arg = definition.fields().to_arg(variant.span());
    let function_name = definition.get_name(variant.span());
    let member = &variant.ident;
    if !no_async {
        let result = quote! {
            #vis async fn #function_name (&self, #arg_def) -> std::option::Option<()> {
                self.sender
                    .send(#basic::#member #arg)
                    .await
                    .ok()
            }
        };
        //eprintln!("{:#?}", result.to_string());
        ret.extend(result);
    }
    if block {
        let function_name = if no_async {
            function_name
        } else {
            definition.get_name_block(variant.span())
        };
        let result = quote! {
            #vis fn #function_name (&self, #arg_def) -> std::option::Option<()> {
                self.sender
                    .blocking_send(#basic::#member #arg)
                    .ok()
            }
        };
        ret.extend(result);
    }

    Ok(ret)
}

/// Generates async and/or blocking methods for a variant with a `#[ret(Type)]` return type.
///
/// Creates a oneshot channel internally, sends the variant (with the oneshot sender
/// attached) through the MPSC channel, and awaits (or blocking-receives) the response.
/// Returns `Option<ReturnType>`.
fn generate_waitable_function(
    basic: &Ident,
    variant: &Variant,
    return_type: TokenStream,
    block: bool,
    no_async: bool,
    vis: &Visibility,
) -> syn::Result<TokenStream> {
    let mut ret = TokenStream::new();
    let definition = EnumDefinition::try_from(variant)?;
    let arg_def = definition.fields().to_arg_def(variant.span());
    let arg = definition.fields().enchant_arg(variant.span());
    let function_name = definition.get_name(variant.span());
    let member = &variant.ident;
    if !no_async {
        let q = quote! {
            #vis async fn #function_name (&self, #arg_def) -> std::option::Option< #return_type > {
                let (__private_sender, __private_r) = tokio::sync::oneshot::channel();
                self.sender
                    .send(#basic::#member #arg)
                    .await
                    .ok();
                __private_r.await.ok()
            }
        };
        ret.extend(q);
    };

    if block {
        let function_name = if no_async {
            function_name
        } else {
            definition.get_name_block(variant.span())
        };
        let q = quote! {
            #vis fn #function_name (&self, #arg_def) -> std::option::Option< #return_type > {
                let (__private_sender, __private_r) = tokio::sync::oneshot::channel();
                self.sender
                    .blocking_send(#basic::#member #arg)
                    .ok();
                __private_r.blocking_recv().ok()
            }
        };
        ret.extend(q);
    }
    Ok(ret)
}

/// Generates helper methods for all variants in a oneshot-enhanced enum.
///
/// Dispatches each variant to either [`generate_waitable_function`] (if it has a
/// `#[ret(...)]` attribute) or [`generate_normal_function`] (otherwise).
fn generate_function(
    st: &syn::DeriveInput,
    de: &DataEnum,
    block: bool,
    no_async: bool,
    vis: &Visibility,
) -> syn::Result<TokenStream> {
    let mut ret = TokenStream::new();
    let basic = &st.ident;
    for variant in &de.variants {
        let return_type = parse_return_type(variant);

        let token = if let Some(ret) = return_type {
            generate_waitable_function(basic, variant, ret, block, no_async, vis)
        } else {
            generate_normal_function(basic, variant, block, no_async, vis)
        }?;

        //eprintln!("{:?}", token);

        ret.extend(token);
    }
    Ok(ret)
}

/// Extracts the return type from a variant's `#[ret(Type)]` attribute, if present.
///
/// Returns `None` if the variant has no `#[ret(...)]` attribute.
fn parse_return_type(variant: &Variant) -> Option<TokenStream> {
    for attr in &variant.attrs {
        match &attr.meta {
            syn::Meta::List(list) => {
                if !list.path.segments.first()?.ident.eq("ret") {
                    continue;
                }
                return Some(list.tokens.clone());
            }
            _ => continue,
        }
    }
    None
}

/// Re-generates enum variant definitions, injecting oneshot sender fields where needed.
///
/// For variants with `#[ret(Type)]`, appends a `tokio::sync::oneshot::Sender<Type>` field.
/// Variants without `#[ret(...)]` are emitted unchanged.
fn generate_member(_st: &syn::DeriveInput, de: &DataEnum) -> syn::Result<TokenStream> {
    let mut ret = TokenStream::new();
    //eprintln!("{:?}", st);
    for variant in &de.variants {
        //eprintln!("{:?}", parse_return_type(variant));
        let return_type = parse_return_type(variant);
        let definition = EnumDefinition::try_from(variant)?;
        let arg_def = definition.fields().enum_arg_def(return_type);
        //eprintln!("{:?}", arg_def);
        let enum_name = definition.get_normal_name(variant.span());
        let result = quote! {
            #enum_name #arg_def ,
        };
        //eprintln!("{:#?}", result.to_string());
        ret.extend(result);
    }
    Ok(ret)
}

/// First expansion pass: re-emits the enum definition with injected oneshot sender fields.
///
/// Strips `#[helper(...)]` attributes from the enum and delegates to [`generate_member`]
/// to rebuild the variant list with additional sender fields for `#[ret(...)]` variants.
fn do_first_expand(st: &syn::DeriveInput) -> syn::Result<TokenStream> {
    //eprintln!("{} {}", block, no_async);
    let data_enum = extract_enum(st)?;
    //print_fields();

    let enum_ident = &st.ident;

    let mut attrs = TokenStream::new();
    st.attrs.iter().for_each(|x| {
        //eprintln!("{x:?}");
        if let syn::Attribute {
            meta:
                syn::Meta::List(syn::MetaList {
                    path: syn::Path { segments, .. },
                    ..
                }),
            ..
        } = x
            && segments.first().is_some_and(|x| x.ident.eq("helper"))
        {
            return;
        }
        x.to_tokens(&mut attrs)
    });

    let member_function = generate_member(st, data_enum)?;

    let ret = quote! {

        #attrs
        pub enum #enum_ident {
            #member_function
        }

    };

    Ok(ret)
}

/// Entry point for the `oneshot_helper!` macro.
///
/// Performs two expansion passes:
/// 1. Re-emits the enum with oneshot sender fields injected ([`do_first_expand`])
/// 2. Generates the helper struct and methods via [`do_expand`] with the oneshot-aware
///    [`generate_function`]
pub(crate) fn handle_new(input: syn::DeriveInput) -> syn::Result<TokenStream> {
    //eprintln!("{:?}", early_st);

    crate::basic::early_check(&input)?;

    do_first_expand(&input).and_then(|mut stream| {
        stream.extend(do_expand(&input, Some(generate_function))?);
        Ok(stream)
    })
}