add proc macros ; benchmark ; example

1 files changed, 508 insertions, 0 deletions

diff --git a/macros/src/decode.rs b/macros/src/decode.rs
new file mode 100644
index 0000000..c289cce
--- /dev/null
+++ b/macros/src/decode.rs
@@ -0,0 +1,508 @@
+use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::quote;
+use syn::{Data, DataEnum, DataStruct, DeriveInput, Generics, Ident, PathArguments, Type};
+
+use crate::encode::{next_group, next_ident, parse_variant_attrs, skip_commas};
+
+// ── Helper: extract element type from collection ────────────────────
+
+fn extract_element_type(ty: &Type) -> Option<Type> {
+    if let Type::Path(type_path) = ty
+        && let Some(segment) = type_path.path.segments.last()
+    {
+        match segment.ident.to_string().as_str() {
+            "Vec" | "Option" | "HashSet" | "BTreeSet" | "BinaryHeap" | "LinkedList"
+            | "VecDeque" => {
+                if let PathArguments::AngleBracketed(args) = &segment.arguments
+                    && let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first()
+                {
+                    return Some(inner_ty.clone());
+                }
+            }
+            "HashMap" | "BTreeMap" => {
+                if let PathArguments::AngleBracketed(args) = &segment.arguments {
+                    let mut types = args.args.iter().filter_map(|arg| {
+                        if let syn::GenericArgument::Type(t) = arg {
+                            Some(t.clone())
+                        } else {
+                            None
+                        }
+                    });
+                    if let (Some(k), Some(v)) = (types.next(), types.next()) {
+                        return Some(syn::parse_quote! { (#k, #v) });
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+    None
+}
+
+// ── Field attribute parsing ─────────────────────────────────────────
+
+enum DecodeAttr {
+    Normal,
+    Count(syn::Type),
+    Len(syn::Type),
+    Custom(syn::Path),
+}
+
+struct FieldAttrs {
+    decode: DecodeAttr,
+    element_type: Option<Type>,
+    condition: Option<syn::Expr>,
+}
+
+fn parse_field_attrs(field: &syn::Field) -> FieldAttrs {
+    let mut decode: Option<DecodeAttr> = None;
+    let mut condition: Option<syn::Expr> = None;
+
+    for attr in &field.attrs {
+        if !attr.path().is_ident("codec") {
+            continue;
+        }
+
+        let syn::Meta::List(list) = &attr.meta else {
+            continue;
+        };
+
+        let mut iter = list.tokens.clone().into_iter();
+
+        loop {
+            skip_commas(&mut iter);
+
+            let ident = match next_ident(&mut iter) {
+                Some(i) => i,
+                None => break,
+            };
+
+            match ident.to_string().as_str() {
+                "skip" => {
+                    assert!(decode.is_none(), "multiple codec attributes on one field");
+                    decode = Some(DecodeAttr::Normal);
+                }
+                "count" => {
+                    assert!(decode.is_none(), "multiple codec attributes on one field");
+                    let group = next_group(&mut iter).expect("expected count(Type)");
+                    let len_ty: syn::Type =
+                        syn::parse2(group.stream()).expect("expected type in count(...)");
+                    let element_type = extract_element_type(&field.ty);
+                    decode = Some(DecodeAttr::Count(len_ty));
+                    return FieldAttrs {
+                        decode: decode.unwrap(),
+                        element_type,
+                        condition,
+                    };
+                }
+                "len" => {
+                    assert!(decode.is_none(), "multiple codec attributes on one field");
+                    let group = next_group(&mut iter).expect("expected len(Type)");
+                    let ty: syn::Type =
+                        syn::parse2(group.stream()).expect("expected type in len(...)");
+                    decode = Some(DecodeAttr::Len(ty));
+                }
+                "with" => {
+                    assert!(decode.is_none(), "multiple codec attributes on one field");
+                    let group = next_group(&mut iter).expect("expected with(path)");
+                    let path: syn::Path =
+                        syn::parse2(group.stream()).expect("expected path in with(...)");
+                    decode = Some(DecodeAttr::Custom(path));
+                }
+                "if" => {
+                    assert!(condition.is_none(), "multiple codec(if(...)) on one field");
+                    let group = next_group(&mut iter).expect("expected if(expr)");
+                    let expr: syn::Expr =
+                        syn::parse2(group.stream()).expect("expected expression in if(...)");
+                    condition = Some(expr);
+                }
+                other => panic!("unknown codec helper on decode: `{other}`"),
+            }
+        }
+    }
+
+    let element_type = if matches!(decode, Some(DecodeAttr::Count(_))) {
+        extract_element_type(&field.ty)
+    } else {
+        None
+    };
+
+    FieldAttrs {
+        decode: decode.unwrap_or(DecodeAttr::Normal),
+        element_type,
+        condition,
+    }
+}
+
+// ── Context type extraction ─────────────────────────────────────────
+
+fn extract_context_ty(attrs: &[syn::Attribute]) -> Option<Type> {
+    attrs.iter().find_map(|attr| {
+        if attr.path().is_ident("context") {
+            let syn::Meta::List(list) = &attr.meta else {
+                return None;
+            };
+            syn::parse2::<Type>(list.tokens.clone()).ok()
+        } else {
+            None
+        }
+    })
+}
+
+// ── Build impl generics ─────────────────────────────────────────────
+
+fn build_decode_impl_generics(
+    params: &syn::punctuated::Punctuated<syn::GenericParam, syn::Token![,]>,
+    where_clause: &Option<syn::WhereClause>,
+    context_ty: &Option<Type>,
+) -> (TokenStream2, TokenStream2, TokenStream2, TokenStream2) {
+    let params = params.iter().cloned().collect::<Vec<_>>();
+    let has_params = !params.is_empty();
+
+    if let Some(ctx_ty) = context_ty {
+        let impl_generics = if has_params {
+            quote! { <#(#params),*> }
+        } else {
+            quote! {}
+        };
+        let ty_generics = if has_params {
+            quote! { <#(#params),*> }
+        } else {
+            quote! {}
+        };
+        let data_param = quote! { #ctx_ty };
+        let where_clause_tokens = where_clause
+            .as_ref()
+            .map(|wc| quote! { #wc })
+            .unwrap_or_default();
+        (impl_generics, ty_generics, where_clause_tokens, data_param)
+    } else {
+        let mut all_params = vec![syn::parse_quote! { Data }];
+        all_params.extend(params.clone());
+        let impl_generics = quote! { <#(#all_params),*> };
+        let ty_generics = if has_params {
+            quote! { <#(#params),*> }
+        } else {
+            quote! {}
+        };
+        let data_param = quote! { Data };
+        let where_clause_tokens = where_clause
+            .as_ref()
+            .map(|wc| quote! { #wc })
+            .unwrap_or_default();
+        (impl_generics, ty_generics, where_clause_tokens, data_param)
+    }
+}
+
+// ── Decode code generation ──────────────────────────────────────────
+
+pub fn derive_decode(
+    DeriveInput {
+        ident,
+        generics,
+        data,
+        attrs,
+        ..
+    }: DeriveInput,
+) -> TokenStream {
+    let context_ty = extract_context_ty(&attrs);
+    match data {
+        Data::Struct(data_struct) => impl_decode_struct(ident, generics, data_struct, context_ty),
+        Data::Enum(data_enum) => impl_decode_enum(ident, generics, data_enum, context_ty),
+        Data::Union(_) => panic!("Not implemented for Union"),
+    }
+}
+
+fn gen_decode_field(
+    field_type: &Type,
+    attrs: &FieldAttrs,
+    data_param: &TokenStream2,
+) -> TokenStream2 {
+    let inner = match &attrs.decode {
+        DecodeAttr::Normal => quote! {
+            <#field_type as zr_protocol::codec::decode::Decode<#data_param>>::decode(buf, ctx)?
+        },
+        DecodeAttr::Count(len_type) => {
+            let element_type = attrs
+                .element_type
+                .clone()
+                .unwrap_or_else(|| field_type.clone());
+            quote! {
+                {
+                    let count: #len_type = <#len_type as zr_protocol::codec::decode::Decode<#data_param>>::decode(buf, ctx)?;
+                    let mut items = Vec::with_capacity(count as usize);
+                    for _ in 0..count {
+                        items.push(
+                            <#element_type as zr_protocol::codec::decode::Decode<#data_param>>::decode(buf, ctx)?
+                        );
+                    }
+                    items
+                }
+            }
+        }
+        DecodeAttr::Len(len_type) => quote! {
+            {
+                let byte_len: #len_type = <#len_type as zr_protocol::codec::decode::Decode<#data_param>>::decode(buf, ctx)?;
+                let byte_len = byte_len as usize;
+                if buf.len() < byte_len {
+                    return Err(zr_protocol::codec::error::CodecError::IoError(std::io::Error::new(
+                        std::io::ErrorKind::UnexpectedEof,
+                        "insufficient bytes for len-prefixed field",
+                    )));
+                }
+                let (mut data, rest) = buf.split_at(byte_len);
+                *buf = rest;
+                <#field_type as zr_protocol::codec::decode::Decode<#data_param>>::decode(&mut data, ctx)?
+            }
+        },
+        DecodeAttr::Custom(fn_path) => quote! {
+            #fn_path(buf, ctx)?
+        },
+    };
+
+    match &attrs.condition {
+        Some(expr) => quote! {
+            if #expr {
+                #inner
+            } else {
+                Default::default()
+            }
+        },
+        None => inner,
+    }
+}
+
+fn impl_decode_struct(
+    ident: Ident,
+    generics: Generics,
+    data_struct: DataStruct,
+    context_ty: Option<Type>,
+) -> TokenStream {
+    let (impl_generics, ty_generics, where_clause, data_param) =
+        build_decode_impl_generics(&generics.params, &generics.where_clause, &context_ty);
+
+    let decode_fields = match &data_struct.fields {
+        syn::Fields::Named(fields) => {
+            let field_decodes: Vec<_> = fields
+                .named
+                .iter()
+                .map(|field| {
+                    let name = field.ident.as_ref().unwrap();
+                    let ty = &field.ty;
+                    let attrs = parse_field_attrs(field);
+                    let decoded = gen_decode_field(ty, &attrs, &data_param);
+                    quote! { #name: #decoded }
+                })
+                .collect();
+            quote! { Ok(Self { #(#field_decodes),* }) }
+        }
+        syn::Fields::Unnamed(fields) => {
+            let field_decodes: Vec<_> = fields
+                .unnamed
+                .iter()
+                .map(|field| {
+                    let ty = &field.ty;
+                    let attrs = parse_field_attrs(field);
+                    gen_decode_field(ty, &attrs, &data_param)
+                })
+                .collect();
+            quote! { Ok(Self(#(#field_decodes),*)) }
+        }
+        syn::Fields::Unit => quote! { Ok(Self) },
+    };
+
+    quote! {
+        impl #impl_generics zr_protocol::codec::decode::Decode<#data_param> for #ident #ty_generics #where_clause {
+            fn decode(buf: &mut &[u8], ctx: &zr_protocol::context::Context<#data_param>) -> zr_protocol::codec::error::Result<Self> {
+                #decode_fields
+            }
+        }
+    }
+    .into()
+}
+
+fn impl_decode_enum(
+    ident: Ident,
+    generics: Generics,
+    data_enum: DataEnum,
+    context_ty: Option<Type>,
+) -> TokenStream {
+    let (_, _, _, data_param) =
+        build_decode_impl_generics(&generics.params, &generics.where_clause, &context_ty);
+
+    // Parse all variant attributes
+    let parsed: Vec<_> = data_enum.variants.iter().map(parse_variant_attrs).collect();
+
+    // Build explicit ID set and auto-ID counter
+    let mut explicit_ids = std::collections::HashSet::new();
+    for attrs in &parsed {
+        if let Some(id) = attrs.id {
+            explicit_ids.insert(id);
+        }
+    }
+
+    // Assign discriminants (same logic as encode)
+    let mut auto_counter: u8 = 0;
+    let mut discriminants: Vec<u8> = Vec::with_capacity(data_enum.variants.len());
+    for attrs in &parsed {
+        if let Some(id) = attrs.id {
+            discriminants.push(id);
+        } else {
+            while explicit_ids.contains(&auto_counter) {
+                auto_counter = auto_counter.wrapping_add(1);
+            }
+            let d = auto_counter;
+            auto_counter = auto_counter.wrapping_add(1);
+            discriminants.push(d);
+        }
+    }
+
+    // Group variants by discriminant ID
+    use std::collections::BTreeMap;
+    type VariantEntry<'a> = (usize, &'a Ident, &'a syn::Fields, &'a Option<syn::Expr>);
+    let mut groups: BTreeMap<u8, Vec<VariantEntry>> = BTreeMap::new();
+
+    for (i, (variant, attrs)) in data_enum.variants.iter().zip(parsed.iter()).enumerate() {
+        let disc = discriminants[i];
+        if attrs.skip {
+            continue;
+        }
+        groups.entry(disc).or_default().push((
+            i,
+            &variant.ident,
+            &variant.fields,
+            &attrs.condition,
+        ));
+    }
+
+    // Generate decode match arms
+    let match_arms: Vec<_> = groups
+        .iter()
+        .map(|(disc, entries)| {
+            if entries.len() == 1 {
+                let (_, variant_ident, fields, condition) = &entries[0];
+                let field_decode = gen_variant_decode(&ident, variant_ident, fields, &data_param);
+
+                match condition {
+                    Some(expr) => quote! {
+                        #disc => {
+                            if #expr {
+                                #field_decode
+                            } else {
+                                Err(zr_protocol::codec::error::CodecError::Custom(
+                                    concat!("variant ", stringify!(#variant_ident), " not valid in current context").into()
+                                ))
+                            }
+                        }
+                    },
+                    None => quote! {
+                        #disc => { #field_decode }
+                    },
+                }
+            } else {
+                let mut arms: Vec<TokenStream2> = Vec::new();
+                let mut has_unconditional = false;
+
+                for (_, variant_ident, fields, condition) in entries {
+                    let field_decode = gen_variant_decode(&ident, variant_ident, fields, &data_param);
+
+                    match condition {
+                        Some(expr) => {
+                            arms.push(quote! {
+                                if #expr {
+                                    #field_decode
+                                }
+                            });
+                        }
+                        None => {
+                            has_unconditional = true;
+                            arms.push(field_decode);
+                        }
+                    }
+                }
+
+                if has_unconditional {
+                    let last = arms.pop().unwrap();
+                    let chain = arms.into_iter().rev().fold(last, |acc, arm| {
+                        quote! { #arm else { #acc } }
+                    });
+                    quote! { #disc => { #chain } }
+                } else {
+                    let chain = arms.into_iter().rev().fold(
+                        quote! {
+                            Err(zr_protocol::codec::error::CodecError::Custom(
+                                format!("no variant for ID {:#04x} in current context", #disc).into()
+                            ))
+                        },
+                        |acc, arm| {
+                            quote! { #arm else { #acc } }
+                        },
+                    );
+                    quote! { #disc => { #chain } }
+                }
+            }
+        })
+        .collect();
+
+    let (impl_generics, ty_generics, where_clause, _) =
+        build_decode_impl_generics(&generics.params, &generics.where_clause, &context_ty);
+
+    quote! {
+        impl #impl_generics zr_protocol::codec::decode::Decode<#data_param> for #ident #ty_generics #where_clause {
+            fn decode(buf: &mut &[u8], ctx: &zr_protocol::context::Context<#data_param>) -> zr_protocol::codec::error::Result<Self> {
+                let id = <u8 as zr_protocol::codec::decode::Decode<#data_param>>::decode(buf, ctx)?;
+                match id {
+                    #(#match_arms),*
+                    other => Err(zr_protocol::codec::error::CodecError::Custom(
+                        format!("unknown discriminant: {other:#04x}").into()
+                    )),
+                }
+            }
+        }
+    }
+    .into()
+}
+
+fn gen_variant_decode(
+    enum_ident: &Ident,
+    variant_ident: &Ident,
+    fields: &syn::Fields,
+    data_param: &TokenStream2,
+) -> TokenStream2 {
+    match fields {
+        syn::Fields::Named(fields) => {
+            let field_decodes: Vec<_> = fields
+                .named
+                .iter()
+                .map(|field| {
+                    let name = field.ident.as_ref().unwrap();
+                    let ty = &field.ty;
+                    let attrs = parse_field_attrs(field);
+                    let decoded = gen_decode_field(ty, &attrs, data_param);
+                    quote! { #name: #decoded }
+                })
+                .collect();
+            quote! {
+                Ok(#enum_ident::#variant_ident { #(#field_decodes),* })
+            }
+        }
+        syn::Fields::Unnamed(fields) => {
+            let field_decodes: Vec<_> = fields
+                .unnamed
+                .iter()
+                .map(|field| {
+                    let ty = &field.ty;
+                    let attrs = parse_field_attrs(field);
+                    gen_decode_field(ty, &attrs, data_param)
+                })
+                .collect();
+            quote! {
+                Ok(#enum_ident::#variant_ident(#(#field_decodes),*))
+            }
+        }
+        syn::Fields::Unit => quote! {
+            Ok(#enum_ident::#variant_ident)
+        },
+    }
+}