libpso/psopacket/src/lib.rs

// TODO: try less hard, just mem::transmute all this junk away
// TODO: PSOPacketData derive

#![recursion_limit="256"]

extern crate proc_macro;

use proc_macro::TokenStream;
use syn::{parse_macro_input, ItemStruct, NestedMeta, DeriveInput, Field};
use syn::punctuated::Iter;
use quote::quote;

#[derive(Debug, PartialEq)]
enum AttrMeta {
    None,
    Utf8,
    Utf16,
    NoDebug,
}

#[derive(Debug)]
enum AttrType {
    Value(syn::TypePath, syn::Ident, AttrMeta),
    Array(syn::TypePath, syn::Ident, usize, AttrMeta)
}

fn generate_struct_def(name: syn::Ident, attrs: &Vec<AttrType>) -> proc_macro2::TokenStream {
    let mut struct_def = Vec::new();
    for attr in attrs {
        let element = match attr {
            AttrType::Value(ty, name, _) => {
                quote!(pub #name: #ty)
            },
            AttrType::Array(ty, name, len, _) => {
                quote!(pub #name: [#ty; #len])
            }
        };
        struct_def.push(element);
    }

    quote! {
        pub struct #name {
            #(#struct_def),*
        }
    }
}


fn generate_from_bytes(attrs: &Vec<AttrType>) -> Vec<proc_macro2::TokenStream> {
    let mut from_bytes = Vec::new();
    for attr in attrs {
        let element = match attr {
            AttrType::Value(ty, name, _) => {
                let type_str = ty.path.segments[0].ident.to_string();
                if type_str == "Vec" {
                    let vec_type = match &ty.path.segments[0].arguments {
                        syn::PathArguments::AngleBracketed(arg) => {
                            match &arg.args[0] {
                                syn::GenericArgument::Type(typ) => {
                                    match &typ {
                                        syn::Type::Path(path) => {
                                            Some(path.path.segments[0].ident.clone())
                                        }
                                        _ => None
                                    }
                                }
                                _ => None,
                            }
                        }
                        _ => None
                    }.unwrap();

                    quote! {
                        #name: {
                            let mut tmp = Vec::new();

                            for _ in 0..flag  {
                                tmp.push(<#vec_type as PSOPacketData>::from_bytes(&mut cur)?);
                            }
                            tmp
                        }
                    }
                }
                else {
                    quote! {
                        #name: <#ty as PSOPacketData>::from_bytes(&mut cur)?,
                    }
                }
            },
            AttrType::Array(ty, name, len, _) => {
                quote! {
                    #name: {
                        let mut arr = [#ty::default(); #len];
                        for e in arr.iter_mut() {
                            *e = #ty::from_bytes(&mut cur)?
                        }
                        arr
                    },
                }
            }
        };
        from_bytes.push(element);
    }
    from_bytes
}

fn generate_as_bytes(attrs: &Vec<AttrType>) -> Vec<proc_macro2::TokenStream> {
    let mut as_bytes = Vec::new();
    for attr in attrs {
        let element = match attr {
            AttrType::Value(ty, name, _) => {
                let type_str = ty.path.segments[0].ident.to_string();
                if type_str == "Vec" {
                    quote! {
                        flag = self.#name.len() as u32;
                        for i in self.#name.iter()  {
                            buf.extend_from_slice(&PSOPacketData::as_bytes(i));
                        }
                    }
                }
                else {
                    quote! {
                        buf.extend_from_slice(&PSOPacketData::as_bytes(&self.#name));
                    }
                }
            },
            AttrType::Array(_ty, name, len, _) => {
                quote! {
                    for i in 0..#len {
                        buf.extend_from_slice(&self.#name[i].as_bytes());
                    }
                }
            }
        };
        as_bytes.push(element);
    }
    as_bytes
}


fn generate_psopacket_impl(pkt_cmd: u16, name: syn::Ident, attrs: &Vec<AttrType>, include_flag: bool) -> proc_macro2::TokenStream {
    let from_bytes = generate_from_bytes(&attrs);
    let as_bytes = generate_as_bytes(&attrs);

    quote! {
        impl PSOPacket for #name {
            fn from_bytes(data: &[u8]) -> Result<#name, PacketParseError> {
                let mut cur = std::io::Cursor::new(data);
                let mut b: [u8; 2] = [0; 2];
                cur.read(&mut b).unwrap();
                let len = u16::from_le_bytes(b);
                cur.read(&mut b).unwrap();
                let cmd = u16::from_le_bytes(b);
                let mut f: [u8; 4] = [0; 4];

                let flag = if #include_flag {
                    cur.read(&mut f).unwrap();
                    u32::from_le_bytes(f)
                }
                else { 0 };

                if cmd != #pkt_cmd {
                    return Err(PacketParseError::WrongPacketCommand {expected: #pkt_cmd, got: cmd});
                }

                if len as usize != data.len() {
                    return Err(PacketParseError::WrongPacketSize(len, data.len()));
                }

                let result = Ok(#name {
                    #(#from_bytes)*
                });

                if cur.position() as usize != data.len() {
                    return Err(PacketParseError::DataStructNotLargeEnough(cur.position(), data.len()));
                }

                result
            }

            fn as_bytes(&self) -> Vec<u8> {
                let mut buf = Vec::new();
                let mut flag = 0;
                #(#as_bytes)*

                while buf.len() % 4 != 0 {
                    buf.push(0);
                }

                let pkt_len = (buf.len() + if #include_flag { 8 } else { 4 }) as u16;
                let mut prebuf: Vec<u8> = Vec::new();

                prebuf.extend_from_slice(&u16::to_le_bytes(pkt_len));
                prebuf.extend_from_slice(&u16::to_le_bytes(#pkt_cmd));
                if #include_flag {
                    prebuf.extend_from_slice(&u32::to_le_bytes(flag));
                }
                prebuf.append(&mut buf);
                prebuf
            }
        }
    }
}

fn generate_debug_impl(name: syn::Ident, attrs: &Vec<AttrType>) -> proc_macro2::TokenStream {
    let mut dbg_write = Vec::new();
    for attr in attrs {
        let element = match attr {
            AttrType::Value(ty, name, meta) => {
                let ident_str = name.to_string();
                let type_str = ty.path.segments[0].ident.to_string();
                match meta {
                    AttrMeta::NoDebug => quote! {
                        write!(f, "    {} {}: [...]\n", #ident_str, #type_str)?;
                    },
                    _ => quote! {
                        write!(f, "    {} {}: {:?}\n", #ident_str, #type_str, self.#name)?;
                    }
                }
            },
            AttrType::Array(ty, name, len, meta) => {
                let ident_str = name.to_string();
                let type_str = ty.path.segments[0].ident.to_string();
                match meta {
                    AttrMeta::Utf8 => quote! {
                        match std::str::from_utf8(&self.#name) {
                            Ok(v) => write!(f, "    {} [utf8; {}]: {:?}\n", #ident_str, #len, v)?,
                            Err(_) => write!(f, "    {} [{}; {}]: {:?}\n", #ident_str, #type_str, #len, self.#name.to_vec())?,
                        };
                    },
                    AttrMeta::Utf16 => quote! {
                        match String::from_utf16(&self.#name) {
                            Ok(v) => write!(f, "    {} [utf16; {}]: {:?}\n", #ident_str, #len, v)?,
                            Err(_) => write!(f, "    {} [{}; {}]: {:?}\n", #ident_str, #type_str, #len, self.#name.to_vec())?,
                        };
                    },
                    AttrMeta::NoDebug => quote! {
                        write!(f, "    {} [{}; {}]: [...]\n", #ident_str, #type_str, #len)?;
                    },
                    AttrMeta::None => quote! {
                        write!(f, "    {} [{}; {}]: {:?}\n", #ident_str, #type_str, #len, self.#name.to_vec())?;
                    }
                }
            }
        };
        dbg_write.push(element);
    }
    let name_str = name.to_string();
    quote! {
        impl std::fmt::Debug for #name {
            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                write!(f, "packet {} {{\n", #name_str)?;
                #(#dbg_write)*
                write!(f, "}}")
            }
        }
    }
}


fn generate_partialeq_impl(name: syn::Ident, attrs: &Vec<AttrType>) -> proc_macro2::TokenStream {
    let mut partialeq = Vec::new();
    for attr in attrs {
        let element = match attr {
            AttrType::Value(_, name, _) => {
                quote! {
                    if self.#name!= other.#name {
                        return false;
                    }
                }
            },
            AttrType::Array(_, name, _, _) => {
                quote! {
                    if self.#name[..] != other.#name[..] {
                        return false;
                    }
                }

            }
        };
        partialeq.push(element);
    }
    quote! {
        impl std::cmp::PartialEq for #name {
            fn eq(&self, other: &Self) -> bool {
                #(#partialeq)*
                true
            }
        }
    }

}


fn get_struct_fields(fields: Iter<Field>) -> Result<Vec<AttrType>, TokenStream> {
    let mut attrs = Vec::new();

    let mut must_be_last = false;
    for field in fields {
        if must_be_last {
            return Err(syn::Error::new(field.ident.as_ref().unwrap().span(), "variables can not follow Vec or String").to_compile_error().into());
        }
        let mut attr_meta = AttrMeta::None;
        for attr in &field.attrs {
            attr_meta = match attr.path.segments[0].ident.to_string().as_str() {
                "utf8" => AttrMeta::Utf8,
                "utf16" => AttrMeta::Utf16,
                "nodebug" => AttrMeta::NoDebug,
                _ => AttrMeta::None
            }
        }
        match &field.ty {
            syn::Type::Array(ty) => {
                if let (syn::Type::Path(ref ty), syn::Expr::Lit(ref lit)) = (&*ty.elem, &ty.len) {
                    if let syn::Lit::Int(ref int) = lit.lit {
                        attrs.push(AttrType::Array(ty.clone(),
                                                   field.ident.as_ref().unwrap().clone(),
                                                   int.base10_parse().unwrap(),
                                                   attr_meta
                        ))
                    }
                }
            },
            syn::Type::Path(ty) => {
                let type_str = ty.path.segments[0].ident.to_string();
                if type_str == "String" || type_str == "Vec"{
                    must_be_last = true;
                }
                attrs.push(AttrType::Value(ty.clone(),
                                           field.ident.as_ref().unwrap().clone(),
                                           attr_meta))
            },
            _ => {}
        }
    }
    Ok(attrs)
}


#[proc_macro_attribute]
pub fn pso_packet(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as syn::AttributeArgs);
    let mut cmd = 0;
    let mut flag = true;
    let mut manual_flag = false;
    for a in args {
        match &a {
            NestedMeta::Lit(lit) => {
                if let syn::Lit::Int(litint) = lit {
                    cmd = litint.base10_parse().unwrap();
                }
            },
            NestedMeta::Meta(k) => {
                if let syn::Meta::Path(syn::Path {segments, ..}) = k {
                    match segments[0].ident.to_string().as_str() {
                        "no_flag" => flag = false,
                        "manual_flag" => {
                            flag = false;
                            manual_flag = true;
                        },
                        _ => {
                            return syn::Error::new(segments[0].ident.span(), "unknown macro param").to_compile_error().into();
                        }
                    }
                }
            },
        }
    }

    let pkt_struct = parse_macro_input!(item as ItemStruct);
    let attrs = match get_struct_fields(pkt_struct.fields.iter()) {
        Ok(a) => a,
        Err(err) => return err
    };

    if manual_flag {
        match &attrs[0] {
            AttrType::Array(_, ident, _, _) => {
                if ident.to_string() != "flag" {
                    return syn::Error::new(pkt_struct.ident.span(), "struct must have flag as the first field if manual_flag is set").to_compile_error().into();
                }
            },
            AttrType::Value(_, ident, _) => {
                if ident.to_string() != "flag" {
                    return syn::Error::new(pkt_struct.ident.span(), "struct must have flag as the first field if manual_flag is set").to_compile_error().into();
                }
            }
        }
    }

    let struct_def = generate_struct_def(pkt_struct.ident.clone(), &attrs);
    let psopacket_impl = generate_psopacket_impl(cmd, pkt_struct.ident.clone(), &attrs, flag);
    let debug_impl =  generate_debug_impl(pkt_struct.ident.clone(), &attrs);
    let partialeq_impl =  generate_partialeq_impl(pkt_struct.ident.clone(), &attrs);

    let q = quote!{
        #[derive(Clone)]
        #struct_def
        #psopacket_impl
        #debug_impl
        #partialeq_impl
    };

    q.into()
}

fn generate_psomessage_impl(msg_cmd: u8, name: syn::Ident, attrs: &Vec<AttrType>) -> proc_macro2::TokenStream {
    let from_bytes = generate_from_bytes(&attrs);
    let as_bytes = generate_as_bytes(&attrs);

    quote! {
        impl PSOMessage for #name {
            const CMD: u8 = #msg_cmd;
            fn from_bytes<R: std::io::Read + std::io::Seek >(mut cur: &mut R) -> Result<#name, PacketParseError> {
                let mut buf1 = [0u8; 1];
                cur.read(&mut buf1).unwrap();
                let cmd = buf1[0];
                cur.read(&mut buf1).unwrap();
                let size = buf1[0];

                let mut subbuf = vec![0u8; size as usize * 4 - 2];
                let len = cur.read(&mut subbuf).unwrap();

                if cmd != #msg_cmd {
                    return Err(PacketParseError::WrongMessageCommand {expected: #msg_cmd, got: cmd});
                }

                if len != size as usize * 4 - 2 {
                    return Err(PacketParseError::WrongPacketSize(size as u16 * 4, len));
                }

                let mut cur = std::io::Cursor::new(subbuf);
                let result = Ok(#name {
                    #(#from_bytes)*
                });

                result
            }

            fn as_bytes(&self) -> Vec<u8> {
                let mut buf = Vec::new();
                #(#as_bytes)*

                while buf.len() % 4 != 2 {
                    buf.push(0);
                }

                let mut fullbuf = Vec::new();
                fullbuf.push(#msg_cmd);
                fullbuf.push((buf.len() as u8 + 2) / 4);
                fullbuf.extend_from_slice(&mut buf);

                fullbuf
            }
        }
    }
}

#[proc_macro_attribute]
pub fn pso_message(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as syn::AttributeArgs);
    let mut cmd = 0;
    for a in args {
        if let NestedMeta::Lit(lit) = a {
            if let syn::Lit::Int(litint) = lit {
                cmd = litint.base10_parse().unwrap();
            }
        }
    }

    let pkt_struct = parse_macro_input!(item as ItemStruct);
    let mut attrs = match get_struct_fields(pkt_struct.fields.iter()) {
        Ok(a) => a,
        Err(err) => return err
    };

    // this is a lot of work to make a `u8` token, surely this can be easier?
    let mut punctuated: syn::punctuated::Punctuated<syn::PathSegment, syn::Token![::]> = syn::punctuated::Punctuated::new();
    punctuated.push_value(syn::PathSegment {
        ident: syn::Ident::new("u8", proc_macro2::Span::call_site()),
        arguments: syn::PathArguments::None,
    });
    let u8tpath = syn::TypePath {
        qself: None,
        path: syn::Path {
            leading_colon: None,
            segments: punctuated
        }
    };
    attrs.insert(0, AttrType::Value(u8tpath.clone(), syn::Ident::new("target", proc_macro2::Span::call_site()), AttrMeta::None));
    attrs.insert(0, AttrType::Value(u8tpath, syn::Ident::new("client", proc_macro2::Span::call_site()), AttrMeta::None));

    let struct_def = generate_struct_def(pkt_struct.ident.clone(), &attrs);
    let psopacket_impl = generate_psomessage_impl(cmd, pkt_struct.ident.clone(), &attrs);
    let debug_impl =  generate_debug_impl(pkt_struct.ident.clone(), &attrs);
    let partialeq_impl =  generate_partialeq_impl(pkt_struct.ident.clone(), &attrs);

    let q = quote!{
        #[derive(Clone)]
        #struct_def
        #psopacket_impl
        #debug_impl
        #partialeq_impl
    };

    q.into()
}

#[proc_macro_derive(PSOPacketData)]
pub fn pso_packet_data(input: TokenStream) -> TokenStream {
    let derive = parse_macro_input!(input as DeriveInput);

    let name = derive.ident;

    let fields = if let syn::Data::Struct(strct) = derive.data {
        strct.fields
    }
    else {
        return syn::Error::new(name.span(), "PSOPacketData only works on structs").to_compile_error().into();
    };

    let attrs = match get_struct_fields(fields.iter()) {
        Ok(a) => a,
        Err(err) => return err
    };

    let from_bytes = generate_from_bytes(&attrs);
    let as_bytes = generate_as_bytes(&attrs);

    let impl_pso_data_packet = quote! {
        impl PSOPacketData for #name {
            fn from_bytes<R: std::io::Read + std::io::Seek>(mut cur: &mut R) -> Result<Self, PacketParseError> {
                Ok(#name {
                    #(#from_bytes)*
                })
            }

            fn as_bytes(&self) -> Vec<u8> {
                let mut buf = Vec::new();
                #(#as_bytes)*
                buf
            }
        }
    };

    let partialeq = generate_partialeq_impl(name.clone(), &attrs);
    let debug = generate_debug_impl(name, &attrs);

    let q = quote! {
        #impl_pso_data_packet
        #partialeq
        #debug
    };

    q.into()
}