Initial commit - extracted urpc from monorepo

1 files changed, 282 insertions, 0 deletions

diff --git a/urpc-macro/src/lib.rs b/urpc-macro/src/lib.rs
new file mode 100644
index 0000000..4facc60
--- /dev/null
+++ b/urpc-macro/src/lib.rs
@@ -0,0 +1,282 @@
+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
+    })
+}