derive/src/lib.rs

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