enumset_derive/src/lib.rs

   1 #![recursion_limit="128"]
   2 #![cfg_attr(feature = "nightly", feature(proc_macro_diagnostic))]
   3
   4 extern crate syn;
   5 extern crate proc_macro;
   6 extern crate proc_macro2;
   7 extern crate quote;
   8
   9 use self::proc_macro::{TokenStream, TokenTree, Literal};
  10
  11 use proc_macro2::{TokenStream as SynTokenStream};
  12 use syn::*;
  13 use syn::export::Span;
  14 use syn::spanned::Spanned;
  15 use quote::*;
  16
  17 #[cfg(feature = "nightly")]
  18 fn error(span: Span, data: &str) -> TokenStream {
  19     span.unstable().error(data).emit();
  20     TokenStream::new()
  21 }
  22
  23 #[cfg(not(feature = "nightly"))]
  24 fn error(_: Span, data: &str) -> TokenStream {
  25     panic!("{}", data)
  26 }
  27
  28 fn enum_set_type_impl(name: &Ident, all_variants: u128, repr: Ident) -> SynTokenStream {
  29     let typed_enumset = quote!(::enumset::EnumSet<#name>);
  30     let core = quote!(::enumset::internal::core);
  31
  32     // proc_macro2 does not support creating u128 literals.
  33     let all_variants_tt = TokenTree::Literal(Literal::u128_unsuffixed(all_variants));
  34     let all_variants_tt = SynTokenStream::from(TokenStream::from(all_variants_tt));
  35
  36     quote! {
  37         unsafe impl ::enumset::EnumSetType for #name {
  38             type Repr = #repr;
  39             const ALL_BITS: Self::Repr = #all_variants_tt;
  40
  41             fn enum_into_u8(self) -> u8 {
  42                 self as u8
  43             }
  44             unsafe fn enum_from_u8(val: u8) -> Self {
  45                 #core::mem::transmute(val)
  46             }
  47         }
  48
  49         impl <O : Into<#typed_enumset>> #core::ops::Sub<O> for #name {
  50             type Output = #typed_enumset;
  51             fn sub(self, other: O) -> Self::Output {
  52                 ::enumset::EnumSet::only(self) - other.into()
  53             }
  54         }
  55         impl <O : Into<#typed_enumset>> #core::ops::BitAnd<O> for #name {
  56             type Output = #typed_enumset;
  57             fn bitand(self, other: O) -> Self::Output {
  58                 ::enumset::EnumSet::only(self) & other.into()
  59             }
  60         }
  61         impl <O : Into<#typed_enumset>> #core::ops::BitOr<O> for #name {
  62             type Output = #typed_enumset;
  63             fn bitor(self, other: O) -> Self::Output {
  64                 ::enumset::EnumSet::only(self) | other.into()
  65             }
  66         }
  67         impl <O : Into<#typed_enumset>> #core::ops::BitXor<O> for #name {
  68             type Output = #typed_enumset;
  69             fn bitxor(self, other: O) -> Self::Output {
  70                 ::enumset::EnumSet::only(self) ^ other.into()
  71             }
  72         }
  73         impl #core::ops::Not for #name {
  74             type Output = #typed_enumset;
  75             fn not(self) -> Self::Output {
  76                 !::enumset::EnumSet::only(self)
  77             }
  78         }
  79         impl #core::cmp::PartialEq<#typed_enumset> for #name {
  80             fn eq(&self, other: &#typed_enumset) -> bool {
  81                 ::enumset::EnumSet::only(*self) == *other
  82             }
  83         }
  84     }
  85 }
  86
  87 #[proc_macro_derive(EnumSetType)]
  88 pub fn derive_enum_set_type(input: TokenStream) -> TokenStream {
  89     let input: DeriveInput = parse_macro_input!(input);
  90     if let Data::Enum(data) = input.data {
  91         if !input.generics.params.is_empty() {
  92             error(input.generics.span(),
  93                   "`#[derive(EnumSetType)]` cannot be used on enums with type parameters.")
  94         } else {
  95             let mut all_variants = 0u128;
  96             let mut max_variant = 0;
  97             let mut current_variant = 0;
  98             let mut has_manual_discriminant = false;
  99
 100             for variant in &data.variants {
 101                 if let Fields::Unit = variant.fields {
 102                     if let Some((_, expr)) = &variant.discriminant {
 103                         if let Expr::Lit(ExprLit { lit: Lit::Int(i), .. }) = expr {
 104                             current_variant = i.value();
 105                             has_manual_discriminant = true;
 106                         } else {
 107                             return error(variant.span(), "Unrecognized discriminant for variant.")
 108                         }
 109                     }
 110
 111                     if current_variant >= 128 {
 112                         let message = if has_manual_discriminant {
 113                             "`#[derive(EnumSetType)]` only supports enum discriminants up to 127."
 114                         } else {
 115                             "`#[derive(EnumSetType)]` only supports enums up to 128 variants."
 116                         };
 117                         return error(variant.span(), message)
 118                     }
 119
 120                     if all_variants & (1 << current_variant) != 0 {
 121                         return error(variant.span(),
 122                                      &format!("Duplicate enum discriminant: {}", current_variant))
 123                     }
 124                     all_variants |= 1 << current_variant;
 125                     if current_variant > max_variant {
 126                         max_variant = current_variant
 127                     }
 128
 129                     current_variant += 1;
 130                 } else {
 131                     return error(variant.span(),
 132                                  "`#[derive(EnumSetType)]` can only be used on C-like enums.")
 133                 }
 134             }
 135
 136             let repr = Ident::new(if max_variant <= 8 {
 137                 "u8"
 138             } else if max_variant <= 16 {
 139                 "u16"
 140             } else if max_variant <= 32 {
 141                 "u32"
 142             } else if max_variant <= 64 {
 143                 "u64"
 144             } else if max_variant <= 128 {
 145                 "u128"
 146             } else {
 147                 panic!("max_variant > 128?")
 148             }, Span::call_site());
 149
 150             enum_set_type_impl(&input.ident, all_variants, repr).into()
 151         }
 152     } else {
 153         error(input.span(), "`#[derive(EnumSetType)]` may only be used on enums")
 154     }
 155 }