scroll.server/crates/titan-derive/src/codec.rs

use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use syn::{Data, DeriveInput, Error, Fields, GenericArgument, PathArguments, Result, Type};
#[derive(Clone, Copy, PartialEq, Eq)]
enum Leaf {
    Auto,
    Int,
    VInt,
    Byte,
    Short,
    Bool,
    Str,
    StrRef,
    Long,
    VLong,
    Bytes,
    Nested,
}
struct FieldSpec {
    leaf: Leaf,
    split: bool,
    null_list: bool,
    skip: bool,
    stop_if_eof: bool,
}
impl Default for FieldSpec {
    fn default() -> Self {
        Self {
            leaf: Leaf::Auto,
            split: false,
            null_list: false,
            skip: false,
            stop_if_eof: false,
        }
    }
}
pub fn expand_payload(input: &DeriveInput) -> Result<TokenStream> {
    let ident = &input.ident;
    if has_raw(input)? {
        return expand_raw(input);
    }
    let partial = has_partial(input)?;
    let fields = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(named) => named.named.iter().collect::<Vec<_>>(),
            Fields::Unit => Vec::new(),
            Fields::Unnamed(unnamed) => {
                return Err(Error::new_spanned(
                    unnamed,
                    "tuple structs are not supported, use named fields",
                ))
            }
        },
        Data::Enum(data) => {
            return Err(Error::new_spanned(
                data.enum_token,
                "enums are not supported",
            ))
        }
        Data::Union(data) => {
            return Err(Error::new_spanned(
                data.union_token,
                "unions are not supported",
            ))
        }
    };
    let mut encode_body = Vec::new();
    let mut decode_body = Vec::new();
    let mut struct_init = Vec::new();
    for field in &fields {
        let spec = parse_field_spec(&field.attrs)?;
        let name = field.ident.as_ref().expect("named field");
        let ty = &field.ty;
        if spec.skip {
            if !partial {
                struct_init.push(quote!(#name: ::core::default::Default::default()));
            }
            continue;
        }
        let codec = codec_for_spec(ty, &spec)?;
        encode_body.push(quote! {
            <#codec as ::titan::codec::Codec<#ty>>::write(writer, &self.#name)?;
        });
        if partial {
            if spec.stop_if_eof {
                decode_body.push(quote! {
                    if reader.is_at_end() {
                        return ::core::result::Result::Ok(decoded);
                    }
                });
            }
            decode_body.push(quote! {
                decoded.#name = <#codec as ::titan::codec::Codec<#ty>>::read(reader)?;
            });
        } else {
            let binding = format_ident!("field_{}", name);
            decode_body.push(quote! {
                let #binding = <#codec as ::titan::codec::Codec<#ty>>::read(reader)?;
            });
            struct_init.push(quote!(#name: #binding));
        }
    }
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    let decode_impl = if partial {
        quote! {
            fn decode(reader: &mut ::titan::io::ByteStreamReader<'_>) -> ::titan::error::Result<Self> {
                let mut decoded = <Self as ::core::default::Default>::default();
                #(#decode_body)*
                ::core::result::Result::Ok(decoded)
            }
        }
    } else {
        quote! {
            fn decode(reader: &mut ::titan::io::ByteStreamReader<'_>) -> ::titan::error::Result<Self> {
                #(#decode_body)*
                ::core::result::Result::Ok(Self { #(#struct_init),* })
            }
        }
    };
    Ok(quote! {
        impl #impl_generics ::titan::message::Payload for #ident #ty_generics #where_clause {
            fn encode(&self, writer: &mut ::titan::io::ByteStreamWriter) -> ::titan::error::Result<()> {
                #(#encode_body)*
                ::core::result::Result::Ok(())
            }
            #decode_impl
        }
    })
}
fn expand_raw(input: &DeriveInput) -> Result<TokenStream> {
    let ident = &input.ident;
    let Data::Struct(data) = &input.data else {
        return Err(Error::new_spanned(ident, "#[codec(raw)] needs a struct"));
    };
    let Fields::Named(named) = &data.fields else {
        return Err(Error::new_spanned(
            ident,
            "#[codec(raw)] needs one named field",
        ));
    };
    let fields: Vec<_> = named.named.iter().collect();
    if fields.len() != 1 {
        return Err(Error::new_spanned(
            ident,
            "#[codec(raw)] needs exactly one field",
        ));
    }
    let name = fields[0].ident.as_ref().expect("named field");
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    Ok(quote! {
        impl #impl_generics ::titan::message::Payload for #ident #ty_generics #where_clause {
            fn encode(&self, writer: &mut ::titan::io::ByteStreamWriter) -> ::titan::error::Result<()> {
                writer.write_raw(&self.#name);
                ::core::result::Result::Ok(())
            }
            fn decode(reader: &mut ::titan::io::ByteStreamReader<'_>) -> ::titan::error::Result<Self> {
                ::core::result::Result::Ok(Self { #name: reader.read_remaining() })
            }
        }
    })
}
fn has_partial(input: &DeriveInput) -> Result<bool> {
    container_flag(input, "partial")
}
fn has_raw(input: &DeriveInput) -> Result<bool> {
    container_flag(input, "raw")
}
fn container_flag(input: &DeriveInput, wanted: &str) -> Result<bool> {
    let mut found = false;
    for attr in &input.attrs {
        if !attr.path().is_ident("codec") {
            continue;
        }
        attr.parse_nested_meta(|meta| {
            if meta.path.is_ident("partial") || meta.path.is_ident("raw") {
                if meta.path.is_ident(wanted) {
                    found = true;
                }
                return Ok(());
            }
            Err(meta.error("unsupported container level #[codec(...)] key"))
        })?;
    }
    Ok(found)
}
fn parse_field_spec(attrs: &[syn::Attribute]) -> Result<FieldSpec> {
    let mut spec = FieldSpec::default();
    for attr in attrs {
        if !attr.path().is_ident("codec") {
            continue;
        }
        attr.parse_nested_meta(|meta| {
            let ident = meta
                .path
                .get_ident()
                .ok_or_else(|| meta.error("expected an identifier"))?
                .to_string();
            let leaf = match ident.as_str() {
                "int" => Some(Leaf::Int),
                "vint" => Some(Leaf::VInt),
                "byte" => Some(Leaf::Byte),
                "short" => Some(Leaf::Short),
                "bool" => Some(Leaf::Bool),
                "string" => Some(Leaf::Str),
                "string_ref" => Some(Leaf::StrRef),
                "long" => Some(Leaf::Long),
                "vlong" => Some(Leaf::VLong),
                "bytes" => Some(Leaf::Bytes),
                "nested" => Some(Leaf::Nested),
                "null_list" => {
                    spec.null_list = true;
                    None
                }
                "split" => {
                    spec.split = true;
                    None
                }
                "skip" => {
                    spec.skip = true;
                    None
                }
                "stop_if_eof" => {
                    spec.stop_if_eof = true;
                    None
                }
                _ => return Err(meta.error("unsupported #[codec(...)] key")),
            };
            if let Some(leaf) = leaf {
                spec.leaf = leaf;
            }
            Ok(())
        })?;
    }
    Ok(spec)
}
fn codec_for_spec(ty: &Type, spec: &FieldSpec) -> Result<TokenStream> {
    if spec.null_list {
        let inner = option_inner(ty)
            .and_then(vec_inner)
            .ok_or_else(|| Error::new_spanned(ty, "#[codec(null_list)] needs Option<Vec<T>>"))?;
        let inner_codec = codec_for(inner, spec.leaf, false)?;
        return Ok(quote!(::titan::codec::NullList<#inner_codec>));
    }
    codec_for(ty, spec.leaf, spec.split)
}
fn vec_inner(ty: &Type) -> Option<&Type> {
    let Type::Path(path) = ty else {
        return None;
    };
    let segment = path.path.segments.last()?;
    if segment.ident != "Vec" {
        return None;
    }
    generic_argument(segment)
}
fn codec_for(ty: &Type, leaf: Leaf, split: bool) -> Result<TokenStream> {
    match ty {
        Type::Array(array) => {
            let len = &array.len;
            let element = array.elem.as_ref();
            if split {
                let inner = option_inner(element).ok_or_else(|| {
                    Error::new_spanned(
                        element,
                        "#[codec(split)] requires an array of Option<T> elements",
                    )
                })?;
                let inner_codec = codec_for(inner, leaf, false)?;
                return Ok(quote!(::titan::codec::SplitArr<#inner_codec, { #len }>));
            }
            let inner_codec = codec_for(element, leaf, false)?;
            Ok(quote!(::titan::codec::Arr<#inner_codec, { #len }>))
        }
        Type::Path(path) => {
            let segment = path
                .path
                .segments
                .last()
                .ok_or_else(|| Error::new_spanned(ty, "unsupported empty type path"))?;
            let name = segment.ident.to_string();
            if name == "Option" {
                match leaf {
                    Leaf::Str => return Ok(quote!(::titan::codec::Str)),
                    Leaf::Bytes => return Ok(quote!(::titan::codec::Bytes)),
                    _ => {}
                }
                let inner = generic_argument(segment)
                    .ok_or_else(|| Error::new_spanned(ty, "Option requires a type argument"))?;
                let inner_codec = codec_for(inner, leaf, false)?;
                return Ok(quote!(::titan::codec::Opt<#inner_codec>));
            }
            if name == "Vec" {
                if leaf == Leaf::Bytes {
                    return Ok(quote!(::titan::codec::Bytes));
                }
                let inner = generic_argument(segment)
                    .ok_or_else(|| Error::new_spanned(ty, "Vec requires a type argument"))?;
                let inner_codec = codec_for(inner, leaf, false)?;
                return Ok(quote!(::titan::codec::List<#inner_codec>));
            }
            Ok(leaf_codec(leaf, &name))
        }
        other => Err(Error::new_spanned(other, "unsupported field type")),
    }
}
fn leaf_codec(leaf: Leaf, type_name: &str) -> TokenStream {
    let resolved = match leaf {
        Leaf::Auto => match type_name {
            "i32" => Leaf::VInt,
            "u8" => Leaf::Byte,
            "i16" => Leaf::Short,
            "bool" => Leaf::Bool,
            "String" => Leaf::StrRef,
            "LogicLong" => Leaf::VLong,
            _ => Leaf::Nested,
        },
        explicit => explicit,
    };
    match resolved {
        Leaf::Int => quote!(::titan::codec::Int),
        Leaf::VInt => quote!(::titan::codec::VInt),
        Leaf::Byte => quote!(::titan::codec::Byte),
        Leaf::Short => quote!(::titan::codec::Short),
        Leaf::Bool => quote!(::titan::codec::Bool),
        Leaf::Str => quote!(::titan::codec::Str),
        Leaf::StrRef => quote!(::titan::codec::StrRef),
        Leaf::Long => quote!(::titan::codec::Long),
        Leaf::VLong => quote!(::titan::codec::VLong),
        Leaf::Bytes => quote!(::titan::codec::Bytes),
        Leaf::Auto | Leaf::Nested => quote!(::titan::codec::Nested),
    }
}
fn option_inner(ty: &Type) -> Option<&Type> {
    let Type::Path(path) = ty else {
        return None;
    };
    let segment = path.path.segments.last()?;
    if segment.ident != "Option" {
        return None;
    }
    generic_argument(segment)
}
fn generic_argument(segment: &syn::PathSegment) -> Option<&Type> {
    let PathArguments::AngleBracketed(args) = &segment.arguments else {
        return None;
    };
    args.args.iter().find_map(|arg| match arg {
        GenericArgument::Type(ty) => Some(ty),
        _ => None,
    })
}