Potato/potato-protocol-derive/src/lib.rs

#![recursion_limit = "128"]

use darling::{FromDeriveInput, FromVariant};
use quote::ToTokens;
use syn::Path;

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

#[derive(FromDeriveInput)]
#[darling(attributes(packet))]
struct PacketOpts {
    handshake_id: Option<Path>,
    status_id: Option<Path>,
    login_id: Option<Path>,
    configuration_id: Option<Path>,
    play_id: Option<Path>,
}

#[proc_macro_derive(Packet, attributes(packet))]
pub fn packet_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast = syn::parse(input).unwrap();
    let opts = PacketOpts::from_derive_input(&ast).expect("Wrong options");
    let packet_encodable_impl = generate_packet_encodable_impl(&ast);
    let packet_impl = generate_packet_impl(&ast, opts);

    quote! {
        #packet_encodable_impl
        #packet_impl
    }
    .into()
}

#[proc_macro_derive(PacketEncodable, attributes(packet))]
pub fn packet_encodable_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast = syn::parse(input).unwrap();
    let generated = generate_packet_encodable_impl(&ast);

    generated.into()
}

#[proc_macro_derive(JsonString)]
pub fn json_string_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast: syn::DeriveInput = syn::parse(input).unwrap();
    let ident = &ast.ident;

    quote! {
        impl crate::packet_encodable::JsonString for #ident {}
    }
    .into()
}

fn generate_packet_impl(
    ast: &syn::DeriveInput,
    packet_opts: PacketOpts,
) -> proc_macro2::TokenStream {
    let ident = &ast.ident;

    let handshake_id = packet_opts
        .handshake_id
        .map(ToTokens::into_token_stream)
        .unwrap_or(quote! {-1});

    let status_id = packet_opts
        .status_id
        .map(ToTokens::into_token_stream)
        .unwrap_or(quote! {-1});

    let login_id = packet_opts
        .login_id
        .map(ToTokens::into_token_stream)
        .unwrap_or(quote! {-1});

    let configuration_id = packet_opts
        .configuration_id
        .map(ToTokens::into_token_stream)
        .unwrap_or(quote! {-1});

    let play_id = packet_opts
        .play_id
        .map(ToTokens::into_token_stream)
        .unwrap_or(quote! {-1});

    quote! {
        impl crate::packet::Packet for #ident {
            const HANDSHAKE_ID: i32 = #handshake_id;
            const STATUS_ID: i32 = #status_id;
            const LOGIN_ID: i32 = #login_id;
            const CONFIGURATION_ID: i32 = #configuration_id;
            const PLAY_ID: i32 = #play_id;
        }
    }
}

fn generate_packet_encodable_impl(ast: &syn::DeriveInput) -> proc_macro2::TokenStream {
    match ast.data {
        syn::Data::Enum(syn::DataEnum { ref variants, .. }) => {
            generate_packet_encodable_enum_impl(ast, variants)
        }
        syn::Data::Struct(syn::DataStruct { ref fields, .. }) => {
            generate_packet_encodable_struct_impl(ast, fields)
        }

        syn::Data::Union(_) => panic!("PacketEncodable can not be implemented for unions"),
    }
}

fn generate_packet_encodable_enum_impl(
    ast: &syn::DeriveInput,
    variants: &syn::punctuated::Punctuated<syn::Variant, syn::token::Comma>,
) -> proc_macro2::TokenStream {
    let ident = &ast.ident;
    let encode_arms: Vec<_> = variants
        .iter()
        .map(generate_enum_variant_encode_packet_arms)
        .collect();
    let decode_arms: Vec<_> = variants
        .iter()
        .map(generate_enum_variant_decode_packet_arms)
        .collect();
    quote! {
        impl crate::packet_encodable::PacketEncodable for #ident {
            fn encode_packet(&self, buf: &mut impl bytes::BufMut) -> Result<(), crate::packet_encodable::PacketEncodeError> {
                let value: crate::datatypes::var_int::VarInt = match self {
                #(#encode_arms)*
                }.into();

                value.encode_packet(buf)
            }

            fn decode_packet(buf: &mut impl bytes::Buf) -> Result<Self, crate::packet_encodable::PacketDecodeError> {
                let value: i32 = crate::datatypes::var_int::VarInt::decode_packet(buf)?.into();
                match value {
                    #(#decode_arms)*
                    _ => Err(crate::packet_encodable::PacketDecodeError::UnkownEnumVariant {
                        name: stringify!(#ident),
                        value,
                    }),
                }
            }
        }
    }
}

fn generate_packet_encodable_struct_impl(
    ast: &syn::DeriveInput,
    fields: &syn::Fields,
) -> proc_macro2::TokenStream {
    let ident = &ast.ident;
    let fields_encode_packet_calls: Vec<_> = fields
        .iter()
        .map(generate_struct_field_encode_packet_calls)
        .collect();
    let fields_decode_packet_calls: Vec<_> = fields
        .iter()
        .map(generate_struct_field_decode_packet_calls)
        .collect();

    quote! {
        impl crate::packet_encodable::PacketEncodable for #ident {
            fn encode_packet(&self, buf: &mut impl bytes::BufMut) -> Result<(), crate::packet_encodable::PacketEncodeError> {
                #(#fields_encode_packet_calls)*
                Ok(())
            }

            fn decode_packet(buf: &mut impl bytes::Buf) -> Result<Self, crate::packet_encodable::PacketDecodeError> {
                Ok(Self {
                    #(#fields_decode_packet_calls)*
                })
            }
        }
    }
}

fn generate_struct_field_decode_packet_calls(field: &syn::Field) -> proc_macro2::TokenStream {
    let field_ident = &field.ident;
    let field_ty = &field.ty;

    quote! {
        #field_ident: <#field_ty>::decode_packet(buf)?,
    }
}

fn generate_struct_field_encode_packet_calls(field: &syn::Field) -> proc_macro2::TokenStream {
    let field_ident = &field.ident;

    quote! {
        self.#field_ident.encode_packet(buf)?;
    }
}

#[derive(FromVariant)]
#[darling(attributes(packet))]
struct EnumVariant {
    id: i32,
}

fn generate_enum_variant_decode_packet_arms(variant: &syn::Variant) -> proc_macro2::TokenStream {
    let ident = &variant.ident;
    let enum_id = EnumVariant::from_variant(variant).unwrap().id;

    quote! {
        #enum_id => Ok(Self::#ident),
    }
}

fn generate_enum_variant_encode_packet_arms(variant: &syn::Variant) -> proc_macro2::TokenStream {
    let ident = &variant.ident;
    let enum_id = EnumVariant::from_variant(variant).unwrap().id;

    quote! {
        Self::#ident => #enum_id,
    }
}