rust-nex/macros/src/protos.rs

use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, ToTokens};
use syn::{
    parse::{Parse, ParseStream},
    punctuated::Punctuated,
    Attribute, FnArg, ItemTrait, LitInt, LitStr, Meta, Pat, ReturnType, Token, TraitItem, Type,
};

use crate::util::fold_tokenable;

pub struct ProtoInputParams {
    proto_num: LitInt,
    properties: Option<(Token![,], Punctuated<Ident, Token![,]>)>,
}

impl Parse for ProtoInputParams {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let proto_num = input.parse()?;

        if let Some(seperator) = input.parse()? {
            let mut punctuated = Punctuated::new();
            loop {
                punctuated.push_value(input.parse()?);
                if let Some(punct) = input.parse()? {
                    punctuated.push_punct(punct);
                } else {
                    return Ok(Self {
                        proto_num,
                        properties: Some((seperator, punctuated)),
                    });
                }
            }
        } else {
            Ok(Self {
                proto_num,
                properties: None,
            })
        }
    }
}
pub struct ProtoMethodData {
    pub id: LitInt,
    pub name: Ident,
    pub attributes: Vec<Attribute>,
    pub parameters: Vec<(Ident, Type, Vec<Attribute>)>,
    pub ret_val: ReturnType,
}

/// This is a representation of the code generated by `rmc_proto` it serves to split the logic of
/// acquiring data from the actual generation to tidy up the process into first getting then
/// generating.
///
/// Use the [`ToTokens`] trait to generate the actual code.
pub struct RmcProtocolData {
    pub has_returns: bool,
    pub id: LitInt,
    pub name: Ident,
    pub methods: Vec<ProtoMethodData>,
}

impl RmcProtocolData {
    pub fn new(params: ProtoInputParams, input: &ItemTrait) -> Self {
        let ProtoInputParams {
            proto_num,
            properties,
        } = params;

        let no_return_data =
            properties.is_some_and(|p| p.1.iter().any(|i| i.to_string() == "NoReturn"));

        // gigantic ass struct initializer (to summarize this gets all of the data)
        RmcProtocolData {
            has_returns: !no_return_data,
            name: input.ident.clone(),
            id: proto_num,
            methods: input
                .items
                .iter()
                .filter_map(|v| match v {
                    TraitItem::Fn(v) => Some(v),
                    _ => None,
                })
                .map(|func| {
                    let Some(attr) = func.attrs.iter().find(|a| {
                        a.path()
                            .segments
                            .last()
                            .is_some_and(|s| s.ident.to_string() == "method_id")
                    }) else {
                        panic!("every function inside of an rmc protocol must have a method id");
                    };

                    let Ok(id): Result<LitInt, _> = attr.parse_args() else {
                        panic!("todo: put a propper error message here");
                    };

                    let funcs = func
                        .sig
                        .inputs
                        .iter()
                        .skip(1)
                        .map(|f| {
                            let FnArg::Typed(t) = f else {
                                panic!("what");
                            };
                            let Pat::Ident(i) = &*t.pat else {
                                panic!(
                                    "unable to handle non identifier patterns as parameter bindings"
                                );
                            };

                            (i.ident.clone(), t.ty.as_ref().clone(), t.attrs.clone())
                        })
                        .collect();

                    ProtoMethodData {
                        id,
                        name: func.sig.ident.clone(),
                        parameters: funcs,
                        ret_val: func.sig.output.clone(),
                        attributes: func
                            .attrs
                            .iter()
                            .filter(|a| match &a.meta {
                                Meta::NameValue(v) => {
                                    if let Some(i) = v.path.get_ident() {
                                        i.to_string() != "doc"
                                    } else {
                                        true
                                    }
                                }
                                Meta::List(l) => {
                                    if let Some(seg) = l.path.segments.last() {
                                        seg.ident.to_string() != "method_id"
                                    } else {
                                        true
                                    }
                                }
                                _ => true,
                            })
                            .cloned()
                            .collect(),
                    }
                })
                .collect(),
        }
    }

    fn generate_raw_trait(&self) -> TokenStream {
        let Self {
            has_returns,
            name,
            id,
            methods,
        } = self;
        let generate_raw_method = |method: &ProtoMethodData| -> TokenStream {
            let ProtoMethodData {
                name,
                parameters,
                attributes,
                ..
            } = method;

            let attribs = fold_tokenable(attributes.iter());

            let raw_name = Ident::new(&format!("raw_{}", name), name.span());

            let optional_return = if self.has_returns {
                quote! {
                    -> ::core::result::Result<Vec<u8>, ::rnex_core::rmc::response::ErrorCode>
                }
            } else {
                quote! {}
            }
            .into_token_stream();

            let deser_params =
                fold_tokenable(parameters.iter().map(|(param_name, param_type, attribs)| {
                    let error_msg = LitStr::new(
                        &format!("an error occurred whilest deserializing {}", param_name),
                        Span::call_site(),
                    );
                    let return_from_deser_error = if self.has_returns {
                        quote! {
                            return Err(::rnex_core::rmc::response::ErrorCode::Core_InvalidArgument);
                        }
                    } else {
                        quote! {
                            return;
                        }
                    };
                    let attribs = fold_tokenable(attribs.iter());
                    quote! {
                        #attribs
                        let Ok(#param_name) =
                            <#param_type as rnex_core::rmc::structures::RmcSerialize>::deserialize(
                            &mut cursor
                        ) else{
                            log::error!(#error_msg);
                            #return_from_deser_error
                        };
                    }
                }));

            let call_params = fold_tokenable(parameters.iter().map(|(param_name, _, attribs)| {
                let attribs = fold_tokenable(attribs.iter());
                quote! {
                    #attribs
                    #param_name,
                }
            }));

            let optional_method_return = if *has_returns {
                quote! {
                    let retval = retval?;
                    let mut vec = Vec::new();
                    rnex_core::rmc::structures::RmcSerialize::serialize(&retval, &mut vec).ok();
                    Ok(vec)
                }
            } else {
                quote! {}
            };

            quote! {
                #[inline(always)]
                #attribs
                async fn #raw_name (&self, data: ::std::vec::Vec<u8>) #optional_return{
                    let mut cursor = ::std::io::Cursor::new(data);
                    #deser_params
                    let retval = self.#name(#call_params).await;
                    #optional_method_return
                }

            }
        };

        let generate_rmc_call_proto = || {
            let method_entries = fold_tokenable(methods.iter().map(|m| {
                let ProtoMethodData {
                    id,
                    name,
                    attributes,
                    ..
                } = m;

                let attribs = fold_tokenable(attributes.iter());
                let raw_name = Ident::new(&format!("raw_{}", name), name.span());

                quote! {
                    #attribs
                    #id => self.#raw_name(data).await,
                }
            }));

            let optional_notimpl_return = if self.has_returns {
                quote! {
                    Err(rnex_core::rmc::response::ErrorCode::Core_NotImplemented)
                }
            } else {
                quote! {}
            };

            let optional_result_sendback = if *has_returns {
                quote! {
                    rnex_core::rmc::response::send_result(
                        remote_response_connection,
                        ret,
                        #id,
                        method_id,
                        call_id,
                    ).await
                }
            } else {
                quote! {}
            };

            quote! {
                #[inline(always)]
                async fn rmc_call_proto(
                    &self,
                    remote_response_connection: &rnex_core::util::SendingBufferConnection,
                    method_id: u32,
                    call_id: u32,
                    data: Vec<u8>,
                ){
                    let ret = match method_id{
                        #method_entries
                        v => {
                            log::error!("(protocol {})unimplemented method id called on protocol: {}", #id, v);
                            #optional_notimpl_return
                        }
                    };
                    #optional_result_sendback
                }
            }
        };

        // this gives us the name which the identifier of the corresponding Raw trait
        let raw_name = Ident::new(&format!("Raw{}", name), name.span());
        let proto_raw_methods = fold_tokenable(self.methods.iter().map(|m| generate_raw_method(m)));
        let rmc_call_proto = generate_rmc_call_proto();

        // boilerplate tokens which all raw traits need
        quote! {
            #[doc(hidden)]
            #[allow(unused_must_use)]
            pub trait #raw_name: #name{
                #proto_raw_methods
                #rmc_call_proto
            }
            impl<T: #name> #raw_name for T{}
        }
        .to_token_stream()
    }

    fn generate_raw_remote_trait(&self) -> TokenStream {
        let Self {
            has_returns,
            name,
            id: proto_id,
            methods,
            ..
        } = self;

        // this gives us the name which the identifier of the corresponding Raw trait
        let remote_name = Ident::new(&format!("Remote{}", name), name.span());
        let generate_remote_method = |m: &ProtoMethodData| -> TokenStream {
            let ProtoMethodData {
                name,
                parameters,
                ret_val,
                attributes,
                id: method_id,
            } = m;

            let params = fold_tokenable(parameters.iter().map(|(ident, ty, attr)| {
                let attrs = fold_tokenable(attr.iter());
                quote! { #attrs #ident: #ty, }
            }));

            let optional_questionmark_operator = if self.has_returns {
                quote! {
                    ?
                }
            } else {
                quote! {}
            };

            let param_serialize = fold_tokenable(parameters.iter().map(|(name, ty, attrs)|{
                let attrs = fold_tokenable(attrs.iter());
                quote!{
                    #attrs
                    rnex_core::result::ResultExtension::display_err_or_some(
                        <#ty as rnex_core::rmc::structures::RmcSerialize>::serialize(
                            &#name,
                            &mut cursor
                        )
                    ).ok_or(rnex_core::rmc::response::ErrorCode::Core_InvalidArgument)#optional_questionmark_operator ;
                }
            }));

            let make_call = if *has_returns {
                quote! {
                    rnex_core::result::ResultExtension::display_err_or_some(
                        rmc_conn.make_raw_call(&message).await
                    ).ok_or(rnex_core::rmc::response::ErrorCode::Core_Exception)
                }
            } else {
                quote! {
                    rnex_core::result::ResultExtension::display_err_or_some(
                        rmc_conn.make_raw_call_no_response(&message).await
                    );
                }
            };

            let attribs = fold_tokenable(attributes.iter());

            quote! {
                #attribs
                async fn #name(&self, #params) #ret_val{
                    let mut send_data = ::std::vec::Vec::new();
                    let mut cursor = ::std::io::Cursor::new(&mut send_data);
                    #param_serialize

                    let call_id = rand::random();

                    let message = rnex_core::rmc::message::RMCMessage{
                        call_id,
                        method_id: #method_id,
                        protocol_id: #proto_id,
                        rest_of_data: send_data
                    };

                    let rmc_conn = <Self as rnex_core::rmc::protocols::HasRmcConnection>::get_connection(self);

                    #make_call
                }
            }
        };

        let remote_methods = fold_tokenable(methods.iter().map(|m| generate_remote_method(m)));

        quote! {
            #[doc(hidden)]
            #[allow(unused_must_use)]
            pub trait #remote_name: rnex_core::rmc::protocols::HasRmcConnection{
                #remote_methods
            }
        }
    }

    fn generate_raw_info(&self) -> TokenStream {
        let Self { name, id, .. } = self;

        let raw_info_name = Ident::new(&format!("Raw{}Info", name), Span::call_site());

        quote! {
            #[doc(hidden)]
            pub struct #raw_info_name;

            impl #raw_info_name {
                pub const PROTOCOL_ID: u16 = #id;
            }
        }
    }
}

impl ToTokens for RmcProtocolData {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.generate_raw_trait().to_tokens(tokens);
        self.generate_raw_info().to_tokens(tokens);
        self.generate_raw_remote_trait().to_tokens(tokens);
    }
}