path: root/urpc-macro/src/lib.rs

use proc_macro2::{Punct, Spacing, Span, TokenStream};
use quote::{ToTokens, TokenStreamExt};
use syn::parse::Parse;

struct Service {
    vis: syn::Visibility,
    ident: syn::Ident,
    error: syn::Type,
    methods: Vec<Method>,
}

impl Service {
    fn server_trait_ident(&self) -> syn::Ident {
        self.ident.clone()
    }

    fn client_struct_ident(&self) -> syn::Ident {
        syn::Ident::new(&format!("{}Client", self.ident), Span::call_site())
    }
}

impl Parse for Service {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let vis = input.parse()?;
        input.parse::<syn::Token![trait]>()?;
        let ident = input.parse()?;

        let content;
        syn::braced!(content in input);

        content.parse::<syn::Token![type]>()?;
        let error_ident = content.parse::<syn::Ident>()?;
        if error_ident != "Error" {
            return Err(syn::Error::new_spanned(
                error_ident,
                "expected Error associated type",
            ));
        }
        content.parse::<syn::Token![=]>()?;
        let error = content.parse()?;
        content.parse::<syn::Token![;]>()?;

        let mut methods = Vec::default();
        while !content.is_empty() {
            let method = content.parse()?;
            methods.push(method);
        }

        Ok(Self {
            vis,
            ident,
            error,
            methods,
        })
    }
}

struct MethodArg {
    ident: syn::Ident,
    ty: syn::Type,
}

impl ToTokens for MethodArg {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.ident.to_tokens(tokens);
        tokens.append(Punct::new(':', Spacing::Alone));
        self.ty.to_tokens(tokens);
    }
}

impl Parse for MethodArg {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let ident = input.parse()?;
        input.parse::<syn::Token![:]>()?;
        let ty = input.parse()?;
        Ok(Self { ident, ty })
    }
}

struct Method {
    ident: syn::Ident,
    arguments: Vec<MethodArg>,
    return_ty: Option<syn::Type>,
}

impl Parse for Method {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        input.parse::<syn::Token![async]>()?;
        input.parse::<syn::Token![fn]>()?;
        let ident = input.parse::<syn::Ident>()?;
        let arguments_content;
        syn::parenthesized!(arguments_content in input);

        let mut arguments = Vec::default();
        while !arguments_content.is_empty() {
            if !arguments.is_empty() {
                arguments_content.parse::<syn::Token![,]>()?;
            }
            let argument = arguments_content.parse::<MethodArg>()?;
            arguments.push(argument);
        }

        let return_ty = match input.peek(syn::Token![->]) {
            true => {
                input.parse::<syn::Token![->]>()?;
                Some(input.parse()?)
            }
            false => None,
        };

        input.parse::<syn::Token![;]>()?;

        Ok(Self {
            ident,
            arguments,
            return_ty,
        })
    }
}

fn generate_service_into_service_match_arm(service: &Service, method: &Method) -> TokenStream {
    let method_ident = &method.ident;
    let method_name = method.ident.to_string();
    let server_trait = service.server_trait_ident();
    let arg_idents = method
        .arguments
        .iter()
        .map(|arg| &arg.ident)
        .collect::<Vec<_>>();

    quote::quote! {
        #method_name => {
            let ((#(#arg_idents),*), _) =
                urpc::internal::bincode::serde::decode_from_slice(&arguments, urpc::internal::bincode::config::standard()).map_err(std::io::Error::other)?;
            let ctx = Default::default();
            let ret = <T as #server_trait>::#method_ident(&self.0, ctx, #(#arg_idents),*).await;
            let value = urpc::internal::bincode::serde::encode_to_vec(&ret, urpc::internal::bincode::config::standard()).map_err(std::io::Error::other)?;
            Ok(From::from(value))
        }
    }
}

fn generate_service_into_service(service: &Service) -> TokenStream {
    let server_trait = service.server_trait_ident();
    let service_name = service.ident.to_string();
    let match_arms = service
        .methods
        .iter()
        .map(|m| generate_service_into_service_match_arm(service, m))
        .collect::<Vec<_>>();
    quote::quote! {
        fn into_service(self) -> impl urpc::Service {
            struct Adapter<T>(T);

            impl<T> urpc::Service for Adapter<T>
            where
                T: #server_trait,
            {
                fn name() -> &'static str {
                    #service_name
                }

                fn call(
                    &self,
                    method: String,
                    arguments: urpc::internal::bytes::Bytes,
                ) -> std::pin::Pin<Box<dyn std::future::Future<Output = std::io::Result<urpc::internal::bytes::Bytes>> + Send + '_>> {
                    Box::pin(async move {
                        match method.as_str() {
                            #(#match_arms)*
                            _ => panic!("unknown method for service {}", "Router"),
                        }
                    })
                }
            }

            Adapter(self)
        }
    }
}

fn generate_service_server_method(service: &Service, method: &Method) -> TokenStream {
    let ident = &method.ident;
    let service_error = &service.error;
    let args = &method.arguments;
    let return_ty = match &method.return_ty {
        Some(t) => quote::quote! { std::result::Result<#t, #service_error> },
        None => quote::quote! { std::result::Result<(), #service_error> },
    };
    quote::quote! {
        fn #ident(&self, ctx: urpc::Context, #(#args),*) -> impl std::future::Future<Output = #return_ty> + Send;
    }
}

fn generate_service_server(service: &Service) -> TokenStream {
    let vis = &service.vis;
    let ident = service.server_trait_ident();
    let methods = service
        .methods
        .iter()
        .map(|m| generate_service_server_method(service, m));
    let into_service = generate_service_into_service(service);

    quote::quote! {
        #vis trait #ident : Sized + Send + Sync + 'static {
            #(#methods)*
            #into_service
        }
    }
}

fn generate_service_client_method(service: &Service, method: &Method) -> TokenStream {
    let service_name = service.ident.to_string();
    let service_error = &service.error;
    let method_ident = &method.ident;
    let method_string = method.ident.to_string();
    let method_arg_idents = method
        .arguments
        .iter()
        .map(|arg| &arg.ident)
        .collect::<Vec<_>>();
    let method_arg_types = method.arguments.iter().map(|arg| &arg.ty);
    let method_return_type = match &method.return_ty {
        Some(ty) => quote::quote! { #ty },
        None => quote::quote! { () },
    };
    quote::quote! {
        pub async fn #method_ident(&self, #(#method_arg_idents: #method_arg_types),*) -> std::result::Result<#method_return_type, urpc::protocol::RpcError<#service_error>> {
            let service = String::from(#service_name);
            let method = String::from(#method_string);
            let arguments = (#(#method_arg_idents),*);
            let arguments = match urpc::internal::bincode::serde::encode_to_vec(&arguments, urpc::internal::bincode::config::standard()) {
                Ok(arguments) => From::from(arguments),
                Err(err) => return Err(urpc::protocol::RpcError::Transport(std::io::Error::other(err))),
            };
            let response = match self.channel.call(service, method, arguments).await {
                Ok(response) => response,
                Err(err) => return Err(urpc::protocol::RpcError::Transport(err)),
            };
            match urpc::internal::bincode::serde::decode_from_slice::<std::result::Result<#method_return_type, #service_error>, _>(&response, urpc::internal::bincode::config::standard()) {
                Ok((result, _)) => result.map_err(urpc::protocol::RpcError::Remote),
                Err(err) => Err(urpc::protocol::RpcError::Transport(std::io::Error::other(err))),
            }
        }
    }
}

fn generate_service_client(service: &Service) -> TokenStream {
    let vis = &service.vis;
    let client_ident = service.client_struct_ident();
    let methods = service
        .methods
        .iter()
        .map(|m| generate_service_client_method(service, m));
    quote::quote! {
        #vis struct #client_ident {
            channel: urpc::ClientChannel
        }

        impl #client_ident {
            pub fn new(channel: urpc::ClientChannel) -> Self {
                Self { channel }
            }

            #(#methods)*
        }
    }
}

#[proc_macro_attribute]
pub fn service(
    _attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let service = syn::parse_macro_input!(item as Service);
    let service_server = generate_service_server(&service);
    let service_client = generate_service_client(&service);
    From::from(quote::quote! {
        #service_server
        #service_client
    })
}