derive/src/lib.rs

   1 use convert_case::{Case, Casing};
   2 use darling::{FromDeriveInput, FromField, FromMeta, FromVariant};
   3 use proc_macro::TokenStream;
   4 use proc_macro2::TokenStream as TokStr;
   5 use quote::{quote, ToTokens};
   6 use syn::{parse_macro_input, parse_quote};
   7
   8 #[derive(Debug, FromMeta, Copy, Clone, Eq, PartialEq)]
   9 #[darling(rename_all = "snake_case")]
  10 enum To {
  11     Clt,
  12     Srv,
  13 }
  14
  15 #[derive(Debug, FromMeta)]
  16 struct MacroArgs {
  17     to: To,
  18     repr: Option<syn::Type>,
  19     tag: Option<String>,
  20     content: Option<String>,
  21     #[darling(default)]
  22     custom: bool,
  23     #[darling(default)]
  24     enumset: bool,
  25 }
  26
  27 fn wrap_attr(attr: &mut syn::Attribute) {
  28     let path = attr.path.clone();
  29     let tokens = attr.tokens.clone();
  30
  31     match attr.path.get_ident().map(|i| i.to_string()).as_deref() {
  32         Some("mt") => {
  33             *attr = parse_quote! {
  34                 #[cfg_attr(any(feature = "client", feature = "server"), #path #tokens)]
  35             };
  36         }
  37         Some("serde") => {
  38             *attr = parse_quote! {
  39                 #[cfg_attr(feature = "serde", #path #tokens)]
  40             };
  41         }
  42         _ => {}
  43     }
  44 }
  45
  46 #[proc_macro_attribute]
  47 pub fn mt_derive(attr: TokenStream, item: TokenStream) -> TokenStream {
  48     let attr_args = parse_macro_input!(attr as syn::AttributeArgs);
  49     let mut input = parse_macro_input!(item as syn::Item);
  50
  51     let args = match MacroArgs::from_list(&attr_args) {
  52         Ok(v) => v,
  53         Err(e) => {
  54             return TokenStream::from(e.write_errors());
  55         }
  56     };
  57
  58     let (serializer, deserializer) = match args.to {
  59         To::Clt => ("server", "client"),
  60         To::Srv => ("client", "server"),
  61     };
  62
  63     let mut out = quote! {
  64         #[derive(Debug)]
  65         #[cfg_attr(feature = "random", derive(GenerateRandom))]
  66         #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
  67     };
  68
  69     macro_rules! iter {
  70         ($t:expr, $f:expr) => {
  71             $t.iter_mut().for_each($f)
  72         };
  73     }
  74
  75     match &mut input {
  76         syn::Item::Enum(e) => {
  77             iter!(e.attrs, wrap_attr);
  78             iter!(e.variants, |v| {
  79                 iter!(v.attrs, wrap_attr);
  80                 iter!(v.fields, |f| iter!(f.attrs, wrap_attr));
  81             });
  82
  83             if args.enumset {
  84                 out.extend(quote! {
  85                     #[derive(EnumSetType)]
  86                     #[enumset(serialize_as_map)]
  87                 });
  88
  89                 if let Some(repr) = args.repr {
  90                     let repr_str = repr.to_token_stream().to_string();
  91
  92                     out.extend(quote! {
  93                         #[enumset(repr = #repr_str)]
  94                     });
  95                 } else if !args.custom {
  96                     panic!("missing repr for enum");
  97                 }
  98             } else {
  99                 let has_payload = e
 100                     .variants
 101                     .iter()
 102                     .find_map(|v| if v.fields.is_empty() { None } else { Some(()) })
 103                     .is_some();
 104
 105                 if has_payload {
 106                     let tag = args.tag.expect("missing tag for enum with payload");
 107
 108                     out.extend(quote! {
 109                         #[cfg_attr(feature = "serde", serde(tag = #tag))]
 110                     });
 111
 112                     if let Some(content) = args.content {
 113                         out.extend(quote! {
 114                             #[cfg_attr(feature = "serde", serde(content = #content))]
 115                         });
 116                     }
 117                 } else {
 118                     out.extend(quote! {
 119                         #[derive(Copy, Eq)]
 120                     });
 121                 }
 122
 123                 out.extend(quote! {
 124                     #[derive(Clone, PartialEq)]
 125                 });
 126
 127                 if !args.custom {
 128                     out.extend(quote! {
 129                         #[cfg_attr(feature = #serializer, derive(MtSerialize))]
 130                         #[cfg_attr(feature = #deserializer, derive(MtDeserialize))]
 131                     });
 132                 }
 133
 134                 if let Some(repr) = args.repr {
 135                     if repr == parse_quote! { str } {
 136                         out.extend(quote! {
 137                                                         #[cfg_attr(any(feature = "client", feature = "server"), mt(string_repr))]
 138                                                 });
 139                     } else {
 140                         out.extend(quote! {
 141                             #[repr(#repr)]
 142                         });
 143                     }
 144                 } else if !args.custom {
 145                     panic!("missing repr for enum");
 146                 }
 147             }
 148
 149             out.extend(quote! {
 150                 #[cfg_attr(feature = "serde", serde(rename_all = "snake_case"))]
 151             });
 152         }
 153         syn::Item::Struct(s) => {
 154             iter!(s.attrs, wrap_attr);
 155             iter!(s.fields, |f| iter!(f.attrs, wrap_attr));
 156
 157             out.extend(quote! {
 158                 #[derive(Clone, PartialEq)]
 159             });
 160
 161             if !args.custom {
 162                 out.extend(quote! {
 163                     #[cfg_attr(feature = #serializer, derive(MtSerialize))]
 164                     #[cfg_attr(feature = #deserializer, derive(MtDeserialize))]
 165                 });
 166             }
 167         }
 168         _ => panic!("only enum and struct supported"),
 169     }
 170
 171     out.extend(input.to_token_stream());
 172     out.into()
 173 }
 174
 175 #[derive(Debug, Default, FromDeriveInput, FromVariant, FromField)]
 176 #[darling(attributes(mt))]
 177 #[darling(default)]
 178 struct MtArgs {
 179     #[darling(multiple)]
 180     const_before: Vec<syn::Expr>, // must implement MtSerialize + MtDeserialize + PartialEq
 181     #[darling(multiple)]
 182     const_after: Vec<syn::Expr>, // must implement MtSerialize + MtDeserialize + PartialEq
 183     size: Option<syn::Type>, // must implement MtCfg
 184     len: Option<syn::Type>,  // must implement MtCfg
 185     default: bool,           // type must implement Default
 186     string_repr: bool,       // for enums
 187     zlib: bool,
 188     zstd: bool,
 189     typename: Option<syn::Ident>, // remote derive
 190 }
 191
 192 type Fields<'a> = Vec<(TokStr, &'a syn::Field)>;
 193
 194 fn get_fields(fields: &syn::Fields, ident: impl Fn(TokStr) -> TokStr) -> Fields {
 195     match fields {
 196         syn::Fields::Named(fs) => fs
 197             .named
 198             .iter()
 199             .map(|f| (ident(f.ident.as_ref().unwrap().to_token_stream()), f))
 200             .collect(),
 201         syn::Fields::Unnamed(fs) => fs
 202             .unnamed
 203             .iter()
 204             .enumerate()
 205             .map(|(i, f)| (ident(syn::Index::from(i).to_token_stream()), f))
 206             .collect(),
 207         syn::Fields::Unit => Vec::new(),
 208     }
 209 }
 210
 211 fn serialize_args(args: &MtArgs, code: &mut TokStr) {
 212     macro_rules! impl_compress {
 213         ($create:expr) => {
 214             *code = quote! {
 215                 let mut __writer = {
 216                     let mut __stream = $create;
 217                     let __writer = &mut __stream;
 218                     #code
 219                     __stream.finish()?
 220                 };
 221             };
 222         };
 223     }
 224
 225     if args.zlib {
 226         impl_compress!(mt_ser::flate2::write::ZlibEncoder::new(
 227             __writer,
 228             mt_ser::flate2::Compression::default()
 229         ));
 230     }
 231
 232     if args.zstd {
 233         impl_compress!(mt_ser::zstd::stream::write::Encoder::new(__writer, 0)?);
 234     }
 235
 236     if let Some(size) = &args.size {
 237         *code = quote! {
 238             mt_ser::MtSerialize::mt_serialize::<#size>(&{
 239                 let mut __buf = Vec::new();
 240                 let __writer = &mut __buf;
 241                 #code
 242                 __buf
 243             }, __writer)?;
 244         };
 245     }
 246
 247     for x in args.const_before.iter().rev() {
 248         *code = quote! {
 249             #x.mt_serialize::<mt_ser::DefCfg>(__writer)?;
 250             #code
 251         }
 252     }
 253
 254     for x in args.const_after.iter() {
 255         *code = quote! {
 256             #code
 257             #x.mt_serialize::<mt_ser::DefCfg>(__writer)?;
 258         }
 259     }
 260 }
 261
 262 fn deserialize_args(args: &MtArgs, code: &mut TokStr) {
 263     macro_rules! impl_compress {
 264         ($create:expr) => {
 265             *code = quote! {
 266                 {
 267                     let mut __owned_reader = $create;
 268                     let __reader = &mut __owned_reader;
 269
 270                     #code
 271                 }
 272             }
 273         };
 274     }
 275
 276     if args.zlib {
 277         impl_compress!(mt_ser::flate2::read::ZlibDecoder::new(mt_ser::WrapRead(
 278             __reader
 279         )));
 280     }
 281
 282     if args.zstd {
 283         impl_compress!(mt_ser::zstd::stream::read::Decoder::new(mt_ser::WrapRead(
 284             __reader
 285         ))?);
 286     }
 287
 288     if let Some(size) = &args.size {
 289         *code = quote! {
 290             #size::mt_deserialize::<DefCfg>(__reader).and_then(|size| {
 291                 let mut __owned_reader = std::io::Read::take(
 292                     mt_ser::WrapRead(__reader), size as u64);
 293                 let __reader = &mut __owned_reader;
 294
 295                 #code
 296             })
 297         };
 298     }
 299
 300     let impl_const = |value: &TokStr| {
 301         quote! {
 302             {
 303                 fn deserialize_same_type<T: MtDeserialize>(
 304                     _: &T,
 305                     reader: &mut impl std::io::Read
 306                 ) -> Result<T, mt_ser::DeserializeError> {
 307                     T::mt_deserialize::<mt_ser::DefCfg>(reader)
 308                 }
 309
 310                 deserialize_same_type(&want, __reader)
 311                     .and_then(|got| {
 312                         if want == got {
 313                             #value
 314                         } else {
 315                             Err(mt_ser::DeserializeError::InvalidConst(
 316                                 Box::new(want), Box::new(got)
 317                             ))
 318                         }
 319                     })
 320             }
 321         }
 322     };
 323
 324     for want in args.const_before.iter().rev() {
 325         let imp = impl_const(&code);
 326         *code = quote! {
 327             {
 328                 let want = #want;
 329                 #imp
 330             }
 331         };
 332     }
 333
 334     for want in args.const_after.iter() {
 335         let imp = impl_const(&quote! { Ok(value) });
 336         *code = quote! {
 337             {
 338                 let want = #want;
 339                 #code.and_then(|value| { #imp })
 340             }
 341         };
 342     }
 343 }
 344
 345 fn serialize_fields(fields: &Fields) -> TokStr {
 346     fields
 347         .iter()
 348         .map(|(ident, field)| {
 349             let args = MtArgs::from_field(field).unwrap();
 350             let def = parse_quote! { mt_ser::DefCfg };
 351             let len = args.len.as_ref().unwrap_or(&def);
 352
 353             let mut code = quote! { mt_ser::MtSerialize::mt_serialize::<#len>(#ident, __writer)?; };
 354             serialize_args(&args, &mut code);
 355
 356             code
 357         })
 358         .collect()
 359 }
 360
 361 fn deserialize_fields(fields: &Fields) -> TokStr {
 362     fields
 363         .iter()
 364         .map(|(ident, field)| {
 365             let args = MtArgs::from_field(field).unwrap();
 366
 367             let def = parse_quote! { mt_ser::DefCfg };
 368             let len = args.len.as_ref().unwrap_or(&def);
 369             let mut code = quote! { mt_ser::MtDeserialize::mt_deserialize::<#len>(__reader) };
 370
 371             if args.default {
 372                 code = quote! {
 373                     mt_ser::OrDefault::or_default(#code)
 374                 };
 375             }
 376
 377             deserialize_args(&args, &mut code);
 378
 379             quote! {
 380                 let #ident = #code?;
 381             }
 382         })
 383         .collect()
 384 }
 385
 386 fn get_fields_struct(input: &syn::Fields) -> (Fields, TokStr) {
 387     let ident_fn = match input {
 388         syn::Fields::Unnamed(_) => |f| {
 389             quote! {
 390                 mt_ser::paste::paste! { [<field_ #f>] }
 391             }
 392         },
 393         _ => |f| quote! { #f },
 394     };
 395
 396     let fields = get_fields(input, ident_fn);
 397     let fields_comma: TokStr = fields
 398         .iter()
 399         .rfold(TokStr::new(), |after, (ident, _)| quote! { #ident, #after });
 400
 401     let fields_struct = match input {
 402         syn::Fields::Named(_) => quote! { { #fields_comma } },
 403         syn::Fields::Unnamed(_) => quote! { ( #fields_comma ) },
 404         syn::Fields::Unit => TokStr::new(),
 405     };
 406
 407     (fields, fields_struct)
 408 }
 409
 410 fn get_repr(input: &syn::DeriveInput, args: &MtArgs) -> syn::Type {
 411     if args.string_repr {
 412         parse_quote! { &str }
 413     } else {
 414         input
 415             .attrs
 416             .iter()
 417             .find(|a| a.path.is_ident("repr"))
 418             .expect("missing repr")
 419             .parse_args()
 420             .expect("invalid repr")
 421     }
 422 }
 423
 424 fn iter_variants(e: &syn::DataEnum, args: &MtArgs, mut f: impl FnMut(&syn::Variant, &syn::Expr)) {
 425     let mut discr = parse_quote! { 0 };
 426
 427     for v in e.variants.iter() {
 428         discr = if args.string_repr {
 429             let lit = v.ident.to_string().to_case(Case::Snake);
 430             parse_quote! { #lit }
 431         } else {
 432             v.discriminant.clone().map(|x| x.1).unwrap_or(discr)
 433         };
 434
 435         f(v, &discr);
 436
 437         discr = parse_quote! { 1 + #discr };
 438     }
 439 }
 440
 441 #[proc_macro_derive(MtSerialize, attributes(mt))]
 442 pub fn derive_serialize(input: TokenStream) -> TokenStream {
 443     let input = parse_macro_input!(input as syn::DeriveInput);
 444     let args = MtArgs::from_derive_input(&input).unwrap();
 445     let typename = args.typename.as_ref().unwrap_or(&input.ident);
 446     let generics = &input.generics;
 447     let mut generics_bounded = generics.clone();
 448     for t in generics_bounded.params.iter_mut() {
 449         match t {
 450             syn::GenericParam::Type(x) => *t = parse_quote! { #x: mt_ser::MtSerialize },
 451             _ => {}
 452         }
 453     }
 454
 455     let mut code = match &input.data {
 456         syn::Data::Enum(e) => {
 457             let repr = get_repr(&input, &args);
 458             let mut variants = TokStr::new();
 459
 460             iter_variants(e, &args, |v, discr| {
 461                 let args = MtArgs::from_variant(v).unwrap();
 462
 463                 let (fields, fields_struct) = get_fields_struct(&v.fields);
 464
 465                 let mut code = serialize_fields(&fields);
 466                 serialize_args(&args, &mut code);
 467
 468                 let ident = &v.ident;
 469
 470                 variants.extend(quote! {
 471                                         #typename::#ident #fields_struct => {
 472                                                 mt_ser::MtSerialize::mt_serialize::<mt_ser::DefCfg>(&((#discr) as #repr), __writer)?;
 473                                                 #code
 474                                         }
 475                                 });
 476             });
 477
 478             quote! {
 479                 match self {
 480                     #variants
 481                 }
 482             }
 483         }
 484         syn::Data::Struct(s) => serialize_fields(&get_fields(&s.fields, |f| quote! { &self.#f })),
 485         _ => {
 486             panic!("only enum and struct supported");
 487         }
 488     };
 489
 490     serialize_args(&args, &mut code);
 491
 492     quote! {
 493                 #[automatically_derived]
 494                 impl #generics_bounded mt_ser::MtSerialize for #typename #generics {
 495                         fn mt_serialize<C: mt_ser::MtCfg>(&self, __writer: &mut impl std::io::Write) -> Result<(), mt_ser::SerializeError> {
 496                                 #code
 497
 498                                 Ok(())
 499                         }
 500                 }
 501         }.into()
 502 }
 503
 504 #[proc_macro_derive(MtDeserialize, attributes(mt))]
 505 pub fn derive_deserialize(input: TokenStream) -> TokenStream {
 506     let input = parse_macro_input!(input as syn::DeriveInput);
 507     let args = MtArgs::from_derive_input(&input).unwrap();
 508     let typename = args.typename.as_ref().unwrap_or(&input.ident);
 509     let generics = &input.generics;
 510     let mut generics_bounded = generics.clone();
 511     for t in generics_bounded.params.iter_mut() {
 512         match t {
 513             syn::GenericParam::Type(x) => *t = parse_quote! { #x: mt_ser::MtDeserialize },
 514             _ => {}
 515         }
 516     }
 517
 518     let mut code = match &input.data {
 519         syn::Data::Enum(e) => {
 520             let repr = get_repr(&input, &args);
 521
 522             let mut consts = TokStr::new();
 523             let mut arms = TokStr::new();
 524
 525             iter_variants(e, &args, |v, discr| {
 526                 let args = MtArgs::from_variant(v).unwrap();
 527
 528                 let ident = &v.ident;
 529                 let (fields, fields_struct) = get_fields_struct(&v.fields);
 530                 let mut code = deserialize_fields(&fields);
 531                 code = quote! {
 532                     #code
 533                     Ok(Self::#ident #fields_struct)
 534                 };
 535
 536                 deserialize_args(&args, &mut code);
 537
 538                 consts.extend(quote! {
 539                     const #ident: #repr = #discr;
 540                 });
 541
 542                 arms.extend(quote! {
 543                     #ident => { #code }
 544                 });
 545             });
 546
 547             let type_str = typename.to_string();
 548             let discr_match = if args.string_repr {
 549                 quote! {
 550                     let __discr = String::mt_deserialize::<DefCfg>(__reader)?;
 551                     match __discr.as_str()
 552                 }
 553             } else {
 554                 quote! {
 555                     let __discr = mt_ser::MtDeserialize::mt_deserialize::<DefCfg>(__reader)?;
 556                     match __discr
 557                 }
 558             };
 559
 560             quote! {
 561                 #consts
 562
 563                 #discr_match {
 564                     #arms
 565                     _ => Err(mt_ser::DeserializeError::InvalidEnum(#type_str, Box::new(__discr)))
 566                 }
 567             }
 568         }
 569         syn::Data::Struct(s) => {
 570             let (fields, fields_struct) = get_fields_struct(&s.fields);
 571             let code = deserialize_fields(&fields);
 572
 573             quote! {
 574                 #code
 575                 Ok(Self #fields_struct)
 576             }
 577         }
 578         _ => {
 579             panic!("only enum and struct supported");
 580         }
 581     };
 582
 583     deserialize_args(&args, &mut code);
 584
 585     quote! {
 586                 #[automatically_derived]
 587                 impl #generics_bounded mt_ser::MtDeserialize for #typename #generics {
 588                         #[allow(non_upper_case_globals)]
 589                         fn mt_deserialize<C: mt_ser::MtCfg>(__reader: &mut impl std::io::Read) -> Result<Self, mt_ser::DeserializeError> {
 590                                 #code
 591                         }
 592                 }
 593         }.into()
 594 }