From e040bf4050dd6d8bdae152a90a38384aff86066a Mon Sep 17 00:00:00 2001
From: Evgeniy Dushistov <dushistov@mail.ru>
Date: Mon, 5 Nov 2018 21:36:09 +0300
Subject: [PATCH] feature: add enum count macros (#37)

* feature: add enum count macros

fix #15

* add docs for EnumCount
---
 README.md                       |  5 +++++
 strum/src/lib.rs                |  5 +++++
 strum_macros/src/enum_count.rs  | 30 ++++++++++++++++++++++++++++++
 strum_macros/src/lib.rs         |  9 +++++++++
 strum_tests/tests/enum_count.rs | 23 +++++++++++++++++++++++
 5 files changed, 72 insertions(+)
 create mode 100644 strum_macros/src/enum_count.rs
 create mode 100644 strum_tests/tests/enum_count.rs

diff --git a/README.md b/README.md
index 614b810..8ae042d 100644
--- a/README.md
+++ b/README.md
@@ -380,6 +380,11 @@ Strum has implemented the following macros:
     }
     ```
 
+8. `EnumCount`: for a given enum generates implementation of `strum::EnumCount`,
+    which returns number of variants via `strum::EnumCount::count` method,
+	also for given `enum MyEnum` generates `const MYENUM_COUNT: usize`
+	which gives the same value as `strum::EnumCount` (which is usefull for array sizes, etc.).
+
 # Additional Attributes
 
 Strum supports several custom attributes to modify the generated code. At the enum level, the
diff --git a/strum/src/lib.rs b/strum/src/lib.rs
index 81bcaf6..34450ad 100644
--- a/strum/src/lib.rs
+++ b/strum/src/lib.rs
@@ -669,3 +669,8 @@ where
 {
     fn as_static(&self) -> &'static T;
 }
+
+/// Number of variants in Enum
+pub trait EnumCount {
+    fn count() -> usize;
+}
diff --git a/strum_macros/src/enum_count.rs b/strum_macros/src/enum_count.rs
new file mode 100644
index 0000000..712d508
--- /dev/null
+++ b/strum_macros/src/enum_count.rs
@@ -0,0 +1,30 @@
+use proc_macro2::{Span, TokenStream};
+use syn;
+
+pub(crate) fn enum_count_inner(ast: &syn::DeriveInput) -> TokenStream {
+    let n = match ast.data {
+        syn::Data::Enum(ref v) => v.variants.len(),
+        _ => panic!("EnumCount can only be used with enums"),
+    };
+    // Used in the quasi-quotation below as `#name`
+    let name = &ast.ident;
+    let const_name = &syn::Ident::new(
+        &format!("{}_COUNT", name.to_string().to_uppercase()),
+        Span::call_site(),
+    );
+
+    // Helper is provided for handling complex generic types correctly and effortlessly
+    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
+
+    quote! {
+            // The generated impl
+            impl #impl_generics ::strum::EnumCount for #name #ty_generics #where_clause {
+                fn count() -> usize {
+                    #n
+                }
+            }
+
+            #[allow(dead_code)]
+            pub const #const_name: usize = #n;
+    }
+}
diff --git a/strum_macros/src/lib.rs b/strum_macros/src/lib.rs
index 3003a64..30d4f5b 100644
--- a/strum_macros/src/lib.rs
+++ b/strum_macros/src/lib.rs
@@ -20,6 +20,7 @@ extern crate proc_macro2;
 mod as_ref_str;
 mod case_style;
 mod display;
+mod enum_count;
 mod enum_discriminants;
 mod enum_iter;
 mod enum_messages;
@@ -124,3 +125,11 @@ pub fn enum_discriminants(input: proc_macro::TokenStream) -> proc_macro::TokenSt
     debug_print_generated(&ast, &toks);
     toks.into()
 }
+
+#[proc_macro_derive(EnumCount, attributes(strum))]
+pub fn enum_count(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse(input).unwrap();
+    let toks = enum_count::enum_count_inner(&ast);
+    debug_print_generated(&ast, &toks);
+    toks.into()
+}
diff --git a/strum_tests/tests/enum_count.rs b/strum_tests/tests/enum_count.rs
new file mode 100644
index 0000000..bdb5b74
--- /dev/null
+++ b/strum_tests/tests/enum_count.rs
@@ -0,0 +1,23 @@
+extern crate strum;
+#[macro_use]
+extern crate strum_macros;
+
+use strum::{EnumCount, IntoEnumIterator};
+
+#[derive(Debug, EnumCount, EnumIter)]
+enum Week {
+    Sunday,
+    Monday,
+    Tuesday,
+    Wednesday,
+    Thursday,
+    Friday,
+    Saturday,
+}
+
+#[test]
+fn simple_test() {
+    assert_eq!(7, Week::count());
+    assert_eq!(Week::count(), WEEK_COUNT);
+    assert_eq!(Week::iter().count(), WEEK_COUNT);
+}