dremoc-macro 0.1.6

Procedural macros for dremoc
Documentation
use crate::utils::{attribute_tokens, to_pascal_case};
use proc_macro2::TokenStream as TokenStream2;
use quote::{quote, ToTokens, TokenStreamExt};
use syn::{
    braced, parenthesized, parse::Parse, punctuated::Punctuated, token, Attribute, Block, Ident,
    ReturnType, Token, Type, TypePath,
};

#[derive(Debug)]
pub struct ServiceMethod {
    pub attrs: Vec<Attribute>,
    pub sig: ServiceMethodSignature,
    pub default: Option<Block>,
    pub semi_token: Option<Token![;]>,
}

impl ServiceMethod {
    pub fn request_enum_variant(&self) -> TokenStream2 {
        let ident = to_pascal_case(&self.sig.ident);
        let entries = self.entries(true);

        quote! { #ident { #entries }, }
    }

    pub fn dispatch_variant(&self) -> TokenStream2 {
        let ident = &self.sig.ident;
        let enum_ident = to_pascal_case(ident);
        let args = self.args(false);

        let call = quote! {
            ::dremoc::remoc::rtc::select! {
                biased;
                () = __result_tx.closed() => (),
                result = target.#ident(#args) => {
                    let _ = __result_tx.send(result);
                }
            }
        };

        quote! {
            Self::#enum_ident { #args __result_tx } => {
                #call
            },
        }
    }

    pub fn variant_to_req(
        &self,
        request_type: &Ident,
        ref_request_type: &Ident,
        ref_mut_request_type: &Ident,
    ) -> TokenStream2 {
        let ident = to_pascal_case(&self.sig.ident);
        let entries = self.entries(false);

        let req = match self.sig.self_kind {
            SelfKind::Ref { .. } => {
                quote! { ::dremoc::remoc::rtc::Req::Ref(#ref_request_type::#ident { #entries }) }
            }
            SelfKind::RefMut { .. } => {
                quote! { ::dremoc::remoc::rtc::Req::RefMut(#ref_mut_request_type::#ident { #entries }) }
            }
        };

        quote! { #request_type::#ident { #entries } => #req, }
    }

    pub fn handle_method(&self, host_feature: &str) -> TokenStream2 {
        let ident = &self.sig.ident;
        let args = self.args(true);
        let call_args = self.args(false);
        let return_type = self.sig.return_type();

        let self_kind = match self.sig.self_kind {
            SelfKind::Ref { .. } => quote! { &self },
            SelfKind::RefMut { .. } => quote! { &mut self },
        };
        let as_ref = match self.sig.self_kind {
            SelfKind::Ref { .. } => quote! { as_ref() },
            SelfKind::RefMut { .. } => quote! { as_mut() },
        };

        quote! {
            #[inline]
            async fn #ident(#self_kind, #args) -> #return_type {
                #[cfg(feature = #host_feature)]
                if let Some(__target) = self.target.#as_ref {
                    return __target.#ident(#call_args).await;
                }

                self.client.#ident(#call_args).await
            }
        }
    }

    pub fn target_method(
        &self,
        trait_path: &TypePath,
        variants: impl Iterator<Item = Ident>,
    ) -> TokenStream2 {
        let ident = &self.sig.ident;
        let args = self.args(true);
        let call_args = self.args(false);
        let return_type = self.sig.return_type();

        let self_kind = match self.sig.self_kind {
            SelfKind::Ref { .. } => quote! { &self },
            SelfKind::RefMut { .. } => quote! { &mut self },
        };

        quote! {
            #[inline(always)]
            async fn #ident(#self_kind, #args) -> #return_type {
                match self {
                    #(Self::#variants(__target) => #trait_path::#ident(__target, #call_args).await),*,
                    Self::__Phantom(_) => unreachable!()
                }
            }
        }
    }

    pub fn client_method(&self, req_ty: &Ident) -> TokenStream2 {
        let ident = &self.sig.ident;
        let enum_ident = to_pascal_case(ident);
        let entries = self.entries(false);
        let args = self.args(true);
        let return_type = self.sig.return_type();

        let self_kind = match self.sig.self_kind {
            SelfKind::Ref { .. } => quote! { &self },
            SelfKind::RefMut { .. } => quote! { &mut self },
        };

        quote! {
            async fn #ident(#self_kind, #args) -> #return_type {
                let (mut __result_tx, __result_rx) = ::dremoc::sync::oneshot::channel();
                __result_tx.set_max_item_size(self.max_reply_size);

                let __req = #req_ty::#enum_ident { #entries };

                self.req_tx
                    .send(__req.into())
                    .await
                    .map_err(::dremoc::remoc::rtc::CallError::from)?;

                __result_rx
                    .await
                    .map_err(::dremoc::remoc::rtc::CallError::from)?
            }
        }
    }

    fn entries(&self, with_type: bool) -> TokenStream2 {
        let mut entries = quote! {};
        for ServiceMethodArgument {
            attrs,
            ident,
            colon_token,
            ty,
        } in &self.sig.args
        {
            if with_type {
                let attrs = attribute_tokens(attrs);
                entries.append_all(quote! {
                    #attrs
                    #ident #colon_token #ty,
                });
            } else {
                entries.append_all(quote! { #ident, });
            }
        }

        if with_type {
            let return_type = self.sig.return_type();
            entries
                .append_all(quote! { __result_tx: ::dremoc::sync::oneshot::Sender<#return_type> });
        } else {
            entries.append_all(quote! { __result_tx });
        }

        entries
    }

    fn args(&self, with_type: bool) -> TokenStream2 {
        let mut args = quote! {};
        for ServiceMethodArgument { ident, ty, .. } in &self.sig.args {
            if with_type {
                args.append_all(quote! { #ident: #ty, });
            } else {
                args.append_all(quote! { #ident, });
            }
        }

        args
    }
}

impl Parse for ServiceMethod {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let sig = input.parse()?;

        let lookahead = input.lookahead1();
        let (brace_token, stmts, semi_token) = if lookahead.peek(token::Brace) {
            let content;
            let brace_token = braced!(content in input);
            let stmts = content.call(Block::parse_within)?;

            (Some(brace_token), stmts, None)
        } else if lookahead.peek(Token![;]) {
            let semi_token: Token![;] = input.parse()?;
            (None, Vec::new(), Some(semi_token))
        } else {
            return Err(lookahead.error());
        };

        Ok(Self {
            attrs,
            sig,
            semi_token,
            default: brace_token.map(|brace_token| Block { brace_token, stmts }),
        })
    }
}

impl ToTokens for ServiceMethod {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        attribute_tokens(&self.attrs).to_tokens(tokens);

        self.sig.to_tokens(tokens);
        self.default.to_tokens(tokens);
        self.semi_token.to_tokens(tokens);
    }
}

#[derive(Debug)]
pub struct ServiceMethodSignature {
    pub asyncness: Option<Token![async]>,
    pub fn_token: Token![fn],
    pub ident: Ident,
    pub paren_token: token::Paren,
    pub self_kind: SelfKind,
    pub comma_token: Option<Token![,]>,
    pub args: Punctuated<ServiceMethodArgument, Token![,]>,
    pub output: ReturnType,
}

impl ServiceMethodSignature {
    pub fn return_type(&self) -> Type {
        match &self.output {
            ReturnType::Default => syn::parse2(quote! { () }).unwrap(),
            ReturnType::Type(_, ty) => ty.as_ref().clone(),
        }
    }
}

impl Parse for ServiceMethodSignature {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let asyncness = input.parse()?;
        let fn_token = input.parse()?;
        let ident = input.parse()?;

        let content;
        let paren_token = parenthesized!(content in input);
        let self_kind = content.parse()?;
        let (comma_token, args) = if let Some(comma_token) = content.parse::<Option<Token![,]>>()? {
            (Some(comma_token), Punctuated::parse_terminated(&content)?)
        } else {
            (None, Punctuated::new())
        };

        let output = input.parse()?;

        Ok(Self {
            asyncness,
            fn_token,
            ident,
            paren_token,
            self_kind,
            comma_token,
            args,
            output,
        })
    }
}

impl ToTokens for ServiceMethodSignature {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        self.asyncness.to_tokens(tokens);
        self.fn_token.to_tokens(tokens);
        self.ident.to_tokens(tokens);
        self.paren_token.surround(tokens, |tokens| {
            self.self_kind.to_tokens(tokens);
            self.comma_token.to_tokens(tokens);
            self.args.to_tokens(tokens);
        });
        self.output.to_tokens(tokens);
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SelfKind {
    Ref {
        and_token: Token![&],
        self_token: Token![self],
    },
    RefMut {
        and_token: Token![&],
        mut_token: Token![mut],
        self_token: Token![self],
    },
}

impl Parse for SelfKind {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        if input.peek(Token![&]) {
            let and_token = input.parse()?;
            if input.peek(Token![mut]) {
                Ok(Self::RefMut {
                    and_token,
                    mut_token: input.parse()?,
                    self_token: input.parse()?,
                })
            } else {
                Ok(Self::Ref {
                    and_token,
                    self_token: input.parse()?,
                })
            }
        } else {
            Ok(Self::Ref {
                and_token: input.parse()?,
                self_token: input.parse()?,
            })
        }
    }
}

impl ToTokens for SelfKind {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        match self {
            Self::Ref {
                and_token,
                self_token,
            } => {
                and_token.to_tokens(tokens);
                self_token.to_tokens(tokens);
            }
            Self::RefMut {
                and_token,
                mut_token,
                self_token,
            } => {
                and_token.to_tokens(tokens);
                mut_token.to_tokens(tokens);
                self_token.to_tokens(tokens);
            }
        }
    }
}

#[derive(Debug)]
pub struct ServiceMethodArgument {
    pub attrs: Vec<Attribute>,
    pub ident: Ident,
    pub colon_token: Token![:],
    pub ty: Type,
}

impl Parse for ServiceMethodArgument {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let ident = input.parse()?;
        let colon_token = input.parse()?;
        let ty = input.parse()?;

        Ok(Self {
            attrs,
            ident,
            colon_token,
            ty,
        })
    }
}

impl ToTokens for ServiceMethodArgument {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        attribute_tokens(&self.attrs).to_tokens(tokens);

        self.ident.to_tokens(tokens);
        self.colon_token.to_tokens(tokens);
        self.ty.to_tokens(tokens);
    }
}