enumset_derive/src/lib.rs

   1 #![recursion_limit="256"]
   2
   3 extern crate proc_macro;
   4
   5 use darling::*;
   6 use proc_macro::TokenStream;
   7 use proc_macro2::{TokenStream as SynTokenStream, Literal, Span};
   8 use std::{collections::HashSet, fmt::Display};
   9 use syn::{*, Result, Error};
  10 use syn::spanned::Spanned;
  11 use quote::*;
  12
  13 /// Helper function for emitting compile errors.
  14 fn error<T>(span: Span, message: impl Display) -> Result<T> {
  15     Err(Error::new(span, message))
  16 }
  17
  18 /// Decodes the custom attributes for our custom derive.
  19 #[derive(FromDeriveInput, Default)]
  20 #[darling(attributes(enumset), default)]
  21 struct EnumsetAttrs {
  22     no_ops: bool,
  23     no_super_impls: bool,
  24     #[darling(default)]
  25     repr: Option<String>,
  26     serialize_as_list: bool,
  27     serialize_deny_unknown: bool,
  28     #[darling(default)]
  29     serialize_repr: Option<String>,
  30     #[darling(default)]
  31     crate_name: Option<String>,
  32 }
  33
  34 /// An variant in the enum set type.
  35 struct EnumSetValue {
  36     /// The name of the variant.
  37     name: Ident,
  38     /// The discriminant of the variant.
  39     variant_repr: u32,
  40 }
  41
  42 /// Stores information about the enum set type.
  43 #[allow(dead_code)]
  44 struct EnumSetInfo {
  45     /// The name of the enum.
  46     name: Ident,
  47     /// The crate name to use.
  48     crate_name: Option<Ident>,
  49     /// The numeric type to represent the `EnumSet` as in memory.
  50     explicit_mem_repr: Option<Ident>,
  51     /// The numeric type to serialize the enum as.
  52     explicit_serde_repr: Option<Ident>,
  53     /// Whether the underlying repr of the enum supports negative values.
  54     has_signed_repr: bool,
  55     /// Whether the underlying repr of the enum supports values higher than 2^32.
  56     has_large_repr: bool,
  57     /// A list of variants in the enum.
  58     variants: Vec<EnumSetValue>,
  59
  60     /// The highest encountered variant discriminant.
  61     max_discrim: u32,
  62     /// The current variant discriminant. Used to track, e.g. `A=10,B,C`.
  63     cur_discrim: u32,
  64     /// A list of variant names that are already in use.
  65     used_variant_names: HashSet<String>,
  66     /// A list of variant discriminants that are already in use.
  67     used_discriminants: HashSet<u32>,
  68
  69     /// Avoid generating operator overloads on the enum type.
  70     no_ops: bool,
  71     /// Avoid generating implementations for `Clone`, `Copy`, `Eq`, and `PartialEq`.
  72     no_super_impls: bool,
  73     /// Serialize the enum as a list.
  74     serialize_as_list: bool,
  75     /// Disallow unknown bits while deserializing the enum.
  76     serialize_deny_unknown: bool,
  77 }
  78 impl EnumSetInfo {
  79     fn new(input: &DeriveInput, attrs: EnumsetAttrs) -> EnumSetInfo {
  80         EnumSetInfo {
  81             name: input.ident.clone(),
  82             crate_name: attrs.crate_name.map(|x| Ident::new(&x, Span::call_site())),
  83             explicit_mem_repr: attrs.repr.map(|x| Ident::new(&x, Span::call_site())),
  84             explicit_serde_repr: attrs.serialize_repr.map(|x| Ident::new(&x, Span::call_site())),
  85             has_signed_repr: false,
  86             has_large_repr: false,
  87             variants: Vec::new(),
  88             max_discrim: 0,
  89             cur_discrim: 0,
  90             used_variant_names: HashSet::new(),
  91             used_discriminants: HashSet::new(),
  92             no_ops: attrs.no_ops,
  93             no_super_impls: attrs.no_super_impls,
  94             serialize_as_list: attrs.serialize_as_list,
  95             serialize_deny_unknown: attrs.serialize_deny_unknown
  96         }
  97     }
  98
  99     /// Sets an explicit repr for the enumset.
 100     fn push_explicit_repr(&mut self, attr_span: Span, repr: &str) -> Result<()> {
 101         // Check whether the repr is supported, and if so, set some flags for better error
 102         // messages later on.
 103         match repr {
 104             "Rust" | "C" | "u8" | "u16" | "u32" => Ok(()),
 105             "usize" | "u64" | "u128" => {
 106                 self.has_large_repr = true;
 107                 Ok(())
 108             }
 109             "i8" | "i16" | "i32" => {
 110                 self.has_signed_repr = true;
 111                 Ok(())
 112             }
 113             "isize" | "i64" | "i128" => {
 114                 self.has_signed_repr = true;
 115                 self.has_large_repr = true;
 116                 Ok(())
 117             }
 118             _ => error(attr_span, "Unsupported repr.")
 119         }
 120     }
 121     /// Adds a variant to the enumset.
 122     fn push_variant(&mut self, variant: &Variant) -> Result<()> {
 123         if self.used_variant_names.contains(&variant.ident.to_string()) {
 124             error(variant.span(), "Duplicated variant name.")
 125         } else if let Fields::Unit = variant.fields {
 126             // Parse the discriminant.
 127             if let Some((_, expr)) = &variant.discriminant {
 128                 let discriminant_fail_message = format!(
 129                     "Enum set discriminants must be `u32`s.{}",
 130                     if self.has_signed_repr || self.has_large_repr {
 131                         format!(
 132                             " ({} discrimiants are still unsupported with reprs that allow them.)",
 133                             if self.has_large_repr {
 134                                 "larger"
 135                             } else if self.has_signed_repr {
 136                                 "negative"
 137                             } else {
 138                                 "larger or negative"
 139                             }
 140                         )
 141                     } else {
 142                         String::new()
 143                     },
 144                 );
 145                 if let Expr::Lit(ExprLit { lit: Lit::Int(i), .. }) = expr {
 146                     match i.base10_parse() {
 147                         Ok(val) => self.cur_discrim = val,
 148                         Err(_) => error(expr.span(), &discriminant_fail_message)?,
 149                     }
 150                 } else {
 151                     error(variant.span(), &discriminant_fail_message)?;
 152                 }
 153             }
 154
 155             // Validate the discriminant.
 156             let discriminant = self.cur_discrim;
 157             if discriminant >= 128 {
 158                 let message = if self.variants.len() <= 127 {
 159                     "`#[derive(EnumSetType)]` currently only supports discriminants up to 127."
 160                 } else {
 161                     "`#[derive(EnumSetType)]` currently only supports enums up to 128 variants."
 162                 };
 163                 error(variant.span(), message)?;
 164             }
 165             if self.used_discriminants.contains(&discriminant) {
 166                 error(variant.span(), "Duplicated enum discriminant.")?;
 167             }
 168
 169             // Add the variant to the info.
 170             self.cur_discrim += 1;
 171             if discriminant > self.max_discrim {
 172                 self.max_discrim = discriminant;
 173             }
 174             self.variants.push(EnumSetValue {
 175                 name: variant.ident.clone(),
 176                 variant_repr: discriminant,
 177             });
 178             self.used_variant_names.insert(variant.ident.to_string());
 179             self.used_discriminants.insert(discriminant);
 180
 181             Ok(())
 182         } else {
 183             error(variant.span(), "`#[derive(EnumSetType)]` can only be used on fieldless enums.")
 184         }
 185     }
 186     /// Validate the enumset type.
 187     fn validate(&self) -> Result<()> {
 188         // Check if all bits of the bitset can fit in the serialization representation.
 189         if let Some(explicit_serde_repr) = &self.explicit_serde_repr {
 190             let is_overflowed = match explicit_serde_repr.to_string().as_str() {
 191                 "u8" => self.max_discrim >= 8,
 192                 "u16" => self.max_discrim >= 16,
 193                 "u32" => self.max_discrim >= 32,
 194                 "u64" => self.max_discrim >= 64,
 195                 "u128" => self.max_discrim >= 128,
 196                 _ => error(
 197                     Span::call_site(),
 198                     "Only `u8`, `u16`, `u32`, `u64` and `u128` are supported for serde_repr."
 199                 )?,
 200             };
 201             if is_overflowed {
 202                 error(Span::call_site(), "serialize_repr cannot be smaller than bitset.")?;
 203             }
 204         }
 205         // Check if all bits of the bitset can fit in the memory representation, if one was given.
 206         if let Some(explicit_mem_repr) = &self.explicit_mem_repr {
 207             let is_overflowed = match explicit_mem_repr.to_string().as_str() {
 208                 "u8" => self.max_discrim >= 8,
 209                 "u16" => self.max_discrim >= 16,
 210                 "u32" => self.max_discrim >= 32,
 211                 "u64" => self.max_discrim >= 64,
 212                 "u128" => self.max_discrim >= 128,
 213                 _ => error(
 214                     Span::call_site(),
 215                     "Only `u8`, `u16`, `u32`, `u64` and `u128` are supported for repr."
 216                 )?,
 217             };
 218             if is_overflowed {
 219                 error(Span::call_site(), "repr cannot be smaller than bitset.")?;
 220             }
 221         }
 222         Ok(())
 223     }
 224
 225     /// Computes the underlying type used to store the enumset.
 226     fn enumset_repr(&self) -> SynTokenStream {
 227         if let Some(explicit_mem_repr) = &self.explicit_mem_repr {
 228             explicit_mem_repr.to_token_stream()
 229         } else if self.max_discrim <= 7 {
 230             quote! { u8 }
 231         } else if self.max_discrim <= 15 {
 232             quote! { u16 }
 233         } else if self.max_discrim <= 31 {
 234             quote! { u32 }
 235         } else if self.max_discrim <= 63 {
 236             quote! { u64 }
 237         } else if self.max_discrim <= 127 {
 238             quote! { u128 }
 239         } else {
 240             panic!("max_variant > 127?")
 241         }
 242     }
 243     /// Computes the underlying type used to serialize the enumset.
 244     #[cfg(feature = "serde")]
 245     fn serde_repr(&self) -> SynTokenStream {
 246         if let Some(serde_repr) = &self.explicit_serde_repr {
 247             quote! { #serde_repr }
 248         } else {
 249             self.enumset_repr()
 250         }
 251     }
 252
 253     /// Returns a bitmask of all variants in the set.
 254     fn all_variants(&self) -> u128 {
 255         let mut accum = 0u128;
 256         for variant in &self.variants {
 257             assert!(variant.variant_repr <= 127);
 258             accum |= 1u128 << variant.variant_repr as u128;
 259         }
 260         accum
 261     }
 262 }
 263
 264 /// Generates the actual `EnumSetType` impl.
 265 fn enum_set_type_impl(info: EnumSetInfo) -> SynTokenStream {
 266     let name = &info.name;
 267
 268     let enumset = match &info.crate_name {
 269         Some(crate_name) => quote!(::#crate_name),
 270         None => {
 271             #[cfg(feature = "proc-macro-crate")]
 272             {
 273                 use proc_macro_crate::FoundCrate;
 274
 275                 let crate_name = proc_macro_crate::crate_name("enumset");
 276                 match crate_name {
 277                     Ok(FoundCrate::Name(name)) => {
 278                         let ident = Ident::new(&name, Span::call_site());
 279                         quote!(::#ident)
 280                     }
 281                     _ => quote!(::enumset),
 282                 }
 283             }
 284
 285             #[cfg(not(feature = "proc-macro-crate"))]
 286             {
 287                 quote!(::enumset)
 288             }
 289         },
 290     };
 291     let typed_enumset = quote!(#enumset::EnumSet<#name>);
 292     let core = quote!(#enumset::__internal::core_export);
 293
 294     let repr = info.enumset_repr();
 295     let all_variants = Literal::u128_unsuffixed(info.all_variants());
 296
 297     let ops = if info.no_ops {
 298         quote! {}
 299     } else {
 300         quote! {
 301             impl <O : Into<#typed_enumset>> #core::ops::Sub<O> for #name {
 302                 type Output = #typed_enumset;
 303                 fn sub(self, other: O) -> Self::Output {
 304                     #enumset::EnumSet::only(self) - other.into()
 305                 }
 306             }
 307             impl <O : Into<#typed_enumset>> #core::ops::BitAnd<O> for #name {
 308                 type Output = #typed_enumset;
 309                 fn bitand(self, other: O) -> Self::Output {
 310                     #enumset::EnumSet::only(self) & other.into()
 311                 }
 312             }
 313             impl <O : Into<#typed_enumset>> #core::ops::BitOr<O> for #name {
 314                 type Output = #typed_enumset;
 315                 fn bitor(self, other: O) -> Self::Output {
 316                     #enumset::EnumSet::only(self) | other.into()
 317                 }
 318             }
 319             impl <O : Into<#typed_enumset>> #core::ops::BitXor<O> for #name {
 320                 type Output = #typed_enumset;
 321                 fn bitxor(self, other: O) -> Self::Output {
 322                     #enumset::EnumSet::only(self) ^ other.into()
 323                 }
 324             }
 325             impl #core::ops::Not for #name {
 326                 type Output = #typed_enumset;
 327                 fn not(self) -> Self::Output {
 328                     !#enumset::EnumSet::only(self)
 329                 }
 330             }
 331             impl #core::cmp::PartialEq<#typed_enumset> for #name {
 332                 fn eq(&self, other: &#typed_enumset) -> bool {
 333                     #enumset::EnumSet::only(*self) == *other
 334                 }
 335             }
 336         }
 337     };
 338
 339
 340     #[cfg(feature = "serde")]
 341     let serde = quote!(#enumset::__internal::serde);
 342
 343     #[cfg(feature = "serde")]
 344     let serde_ops = if info.serialize_as_list {
 345         let expecting_str = format!("a list of {}", name);
 346         quote! {
 347             fn serialize<S: #serde::Serializer>(
 348                 set: #enumset::EnumSet<#name>, ser: S,
 349             ) -> #core::result::Result<S::Ok, S::Error> {
 350                 use #serde::ser::SerializeSeq;
 351                 let mut seq = ser.serialize_seq(#core::prelude::v1::Some(set.len()))?;
 352                 for bit in set {
 353                     seq.serialize_element(&bit)?;
 354                 }
 355                 seq.end()
 356             }
 357             fn deserialize<'de, D: #serde::Deserializer<'de>>(
 358                 de: D,
 359             ) -> #core::result::Result<#enumset::EnumSet<#name>, D::Error> {
 360                 struct Visitor;
 361                 impl <'de> #serde::de::Visitor<'de> for Visitor {
 362                     type Value = #enumset::EnumSet<#name>;
 363                     fn expecting(
 364                         &self, formatter: &mut #core::fmt::Formatter,
 365                     ) -> #core::fmt::Result {
 366                         write!(formatter, #expecting_str)
 367                     }
 368                     fn visit_seq<A>(
 369                         mut self, mut seq: A,
 370                     ) -> #core::result::Result<Self::Value, A::Error> where
 371                         A: #serde::de::SeqAccess<'de>
 372                     {
 373                         let mut accum = #enumset::EnumSet::<#name>::new();
 374                         while let #core::prelude::v1::Some(val) = seq.next_element::<#name>()? {
 375                             accum |= val;
 376                         }
 377                         #core::prelude::v1::Ok(accum)
 378                     }
 379                 }
 380                 de.deserialize_seq(Visitor)
 381             }
 382         }
 383     } else {
 384         let serialize_repr = info.serde_repr();
 385         let check_unknown = if info.serialize_deny_unknown {
 386             quote! {
 387                 if value & !#all_variants != 0 {
 388                     use #serde::de::Error;
 389                     return #core::prelude::v1::Err(
 390                         D::Error::custom("enumset contains unknown bits")
 391                     )
 392                 }
 393             }
 394         } else {
 395             quote! { }
 396         };
 397         quote! {
 398             fn serialize<S: #serde::Serializer>(
 399                 set: #enumset::EnumSet<#name>, ser: S,
 400             ) -> #core::result::Result<S::Ok, S::Error> {
 401                 #serde::Serialize::serialize(&(set.__priv_repr as #serialize_repr), ser)
 402             }
 403             fn deserialize<'de, D: #serde::Deserializer<'de>>(
 404                 de: D,
 405             ) -> #core::result::Result<#enumset::EnumSet<#name>, D::Error> {
 406                 let value = <#serialize_repr as #serde::Deserialize>::deserialize(de)?;
 407                 #check_unknown
 408                 #core::prelude::v1::Ok(#enumset::EnumSet {
 409                     __priv_repr: (value & #all_variants) as #repr,
 410                 })
 411             }
 412         }
 413     };
 414
 415     #[cfg(not(feature = "serde"))]
 416     let serde_ops = quote! { };
 417
 418     let is_uninhabited = info.variants.is_empty();
 419     let is_zst = info.variants.len() == 1;
 420     let into_impl = if is_uninhabited {
 421         quote! {
 422             fn enum_into_u32(self) -> u32 {
 423                 panic!(concat!(stringify!(#name), " is uninhabited."))
 424             }
 425             unsafe fn enum_from_u32(val: u32) -> Self {
 426                 panic!(concat!(stringify!(#name), " is uninhabited."))
 427             }
 428         }
 429     } else if is_zst {
 430         let variant = &info.variants[0].name;
 431         quote! {
 432             fn enum_into_u32(self) -> u32 {
 433                 self as u32
 434             }
 435             unsafe fn enum_from_u32(val: u32) -> Self {
 436                 #name::#variant
 437             }
 438         }
 439     } else {
 440         let variant_name: Vec<_> = info.variants.iter().map(|x| &x.name).collect();
 441         let variant_value: Vec<_> = info.variants.iter().map(|x| x.variant_repr).collect();
 442
 443         let const_field: Vec<_> = ["IS_U8", "IS_U16", "IS_U32", "IS_U64", "IS_U128"]
 444             .iter().map(|x| Ident::new(x, Span::call_site())).collect();
 445         let int_type: Vec<_> = ["u8", "u16", "u32", "u64", "u128"]
 446             .iter().map(|x| Ident::new(x, Span::call_site())).collect();
 447
 448         quote! {
 449             fn enum_into_u32(self) -> u32 {
 450                 self as u32
 451             }
 452             unsafe fn enum_from_u32(val: u32) -> Self {
 453                 // We put these in const fields so the branches they guard aren't generated even
 454                 // on -O0
 455                 #(const #const_field: bool =
 456                     #core::mem::size_of::<#name>() == #core::mem::size_of::<#int_type>();)*
 457                 match val {
 458                     // Every valid variant value has an explicit branch. If they get optimized out,
 459                     // great. If the representation has changed somehow, and they don't, oh well,
 460                     // there's still no UB.
 461                     #(#variant_value => #name::#variant_name,)*
 462                     // Helps hint to the LLVM that this is a transmute. Note that this branch is
 463                     // still unreachable.
 464                     #(x if #const_field => {
 465                         let x = x as #int_type;
 466                         *(&x as *const _ as *const #name)
 467                     })*
 468                     // Default case. Sometimes causes LLVM to generate a table instead of a simple
 469                     // transmute, but, oh well.
 470                     _ => #core::hint::unreachable_unchecked(),
 471                 }
 472             }
 473         }
 474     };
 475
 476     let eq_impl = if is_uninhabited {
 477         quote!(panic!(concat!(stringify!(#name), " is uninhabited.")))
 478     } else {
 479         quote!((*self as u32) == (*other as u32))
 480     };
 481
 482     // used in the enum_set! macro `const fn`s.
 483     let self_as_repr_mask = if is_uninhabited {
 484         quote! { 0 } // impossible anyway
 485     } else {
 486         quote! { 1 << self as #repr }
 487     };
 488
 489     let super_impls = if info.no_super_impls {
 490         quote! {}
 491     } else {
 492         quote! {
 493             impl #core::cmp::PartialEq for #name {
 494                 fn eq(&self, other: &Self) -> bool {
 495                     #eq_impl
 496                 }
 497             }
 498             impl #core::cmp::Eq for #name { }
 499             #[allow(clippy::expl_impl_clone_on_copy)]
 500             impl #core::clone::Clone for #name {
 501                 fn clone(&self) -> Self {
 502                     *self
 503                 }
 504             }
 505             impl #core::marker::Copy for #name { }
 506         }
 507     };
 508
 509     let impl_with_repr = if info.explicit_mem_repr.is_some() {
 510         quote! {
 511             unsafe impl #enumset::EnumSetTypeWithRepr for #name {
 512                 type Repr = #repr;
 513             }
 514         }
 515     } else {
 516         quote! {}
 517     };
 518
 519     quote! {
 520         unsafe impl #enumset::__internal::EnumSetTypePrivate for #name {
 521             type Repr = #repr;
 522             const ALL_BITS: Self::Repr = #all_variants;
 523             #into_impl
 524             #serde_ops
 525         }
 526
 527         unsafe impl #enumset::EnumSetType for #name { }
 528
 529         #impl_with_repr
 530         #super_impls
 531
 532         impl #name {
 533             /// Creates a new enumset with only this variant.
 534             #[deprecated(note = "This method is an internal implementation detail generated by \
 535                                  the `enumset` crate's procedural macro. It should not be used \
 536                                  directly. Use `EnumSet::only` instead.")]
 537             #[doc(hidden)]
 538             pub const fn __impl_enumset_internal__const_only(self) -> #enumset::EnumSet<#name> {
 539                 #enumset::EnumSet { __priv_repr: #self_as_repr_mask }
 540             }
 541
 542             /// Creates a new enumset with this variant added.
 543             #[deprecated(note = "This method is an internal implementation detail generated by \
 544                                  the `enumset` crate's procedural macro. It should not be used \
 545                                  directly. Use the `|` operator instead.")]
 546             #[doc(hidden)]
 547             pub const fn __impl_enumset_internal__const_merge(
 548                 self, chain: #enumset::EnumSet<#name>,
 549             ) -> #enumset::EnumSet<#name> {
 550                 #enumset::EnumSet { __priv_repr: chain.__priv_repr | #self_as_repr_mask }
 551             }
 552         }
 553
 554         #ops
 555     }
 556 }
 557
 558 #[proc_macro_derive(EnumSetType, attributes(enumset))]
 559 pub fn derive_enum_set_type(input: TokenStream) -> TokenStream {
 560     let input: DeriveInput = parse_macro_input!(input);
 561     let attrs: EnumsetAttrs = match EnumsetAttrs::from_derive_input(&input) {
 562         Ok(attrs) => attrs,
 563         Err(e) => return e.write_errors().into(),
 564     };
 565     match derive_enum_set_type_0(input, attrs) {
 566         Ok(v) => v,
 567         Err(e) => e.to_compile_error().into(),
 568     }
 569 }
 570 fn derive_enum_set_type_0(input: DeriveInput, attrs: EnumsetAttrs) -> Result<TokenStream> {
 571     if !input.generics.params.is_empty() {
 572         error(
 573             input.generics.span(),
 574             "`#[derive(EnumSetType)]` cannot be used on enums with type parameters.",
 575         )
 576     } else if let Data::Enum(data) = &input.data {
 577         let mut info = EnumSetInfo::new(&input, attrs);
 578         for attr in &input.attrs {
 579             if attr.path.is_ident(&Ident::new("repr", Span::call_site())) {
 580                 let meta: Ident = attr.parse_args()?;
 581                 info.push_explicit_repr(attr.span(), meta.to_string().as_str())?;
 582             }
 583         }
 584         for variant in &data.variants {
 585             info.push_variant(variant)?;
 586         }
 587         info.validate()?;
 588         Ok(enum_set_type_impl(info).into())
 589     } else {
 590         error(input.span(), "`#[derive(EnumSetType)]` may only be used on enums")
 591     }
 592 }