Repository: Peternator7/strum
Pull request #297 (2023-09-21 21:08:35)

Add `StaticVariantsArray` trait and derive
## About this PR

This PR introduces a new trait `StaticVariantsArray` which defines an associated constant `ALL_VARIANTS: &'static [Self]`. Accompained to it is a derive that will create this array from an `enum` composed of only unit-type variants.

## Motivation

While working in a UI project with [ImGui](https://github.com/imgui-rs/imgui-rs), I stumbled upon [this method, Ui::combo](https://docs.rs/imgui/latest/imgui/struct.Ui.html#method.combo), which expects a slice and a mapping function to build the combobox's options. My options, as you might've guessed, are defined in an enumerator, which looks a little like this:

```rust
enum Operation {
  Add,
  Sub,
  Mul,
  Div,
  ...
}
```

The user, however, wouldn't see those names and, instead, I implement `Into<char>` for the respective characters you can imagine.

Presented with that problem, I could go for one of two solutions:
1. Use `EnumVariantNames`'s `<enum>::VARIANTS`, which is a static slice of strings. The problem with that approach is that it would generate an array with the names of the variants, which could be desired in many cases, but is not on mine as you can see above;
2. Use `EnumIter` and collect the result into a `Vec`. While the application would work fine with this approach, it seems overkill to heap-allocate a vector whose elements we know at compile time.

Considering those two points, I've decided to introduce this trait (and derive).

## Concerns

My only concern with this implementation is the name of both the trait/derive and the associated constant. I believe `VARIANTS` is a more fitting name, but it is already taken. I thought about renaming that to `VARIANT_NAMES`, but it feels like an unecessary breaking change, although I'd like to hear more opinions on this.

## Closing thoughts

This is my first public macro contribution and, while I think it's quite simple and small, I'd love to get feedback on the implementation! 


Patch:
diff --git a/README.md b/README.md
index 491c24f..c4ae3db 100644
--- a/README.md
+++ b/README.md
@@ -46,6 +46,7 @@ Strum has implemented the following macros:
 | [EnumMessage] | Add a verbose message to an enum variant. |
 | [EnumDiscriminants] | Generate a new type with only the discriminant names. |
 | [EnumCount] | Add a constant `usize` equal to the number of variants. |
+| [StaticVariantsArray] | Adds an associated `ALL_VARIANTS` constant which is an array of all enum discriminants |
 
 # Contributing
 
@@ -80,3 +81,4 @@ Strumming is also a very whimsical motion, much like writing Rust code.
 [EnumDiscriminants]: https://docs.rs/strum_macros/0.25/strum_macros/derive.EnumDiscriminants.html
 [EnumCount]: https://docs.rs/strum_macros/0.25/strum_macros/derive.EnumCount.html
 [FromRepr]: https://docs.rs/strum_macros/0.25/strum_macros/derive.FromRepr.html
+[StaticVariantsArray]: https://docs.rs/strum_macros/0.25/strum_macros/derive.StaticVariantsArray.html
diff --git a/strum/src/lib.rs b/strum/src/lib.rs
index 7fbe5e2..31d7626 100644
--- a/strum/src/lib.rs
+++ b/strum/src/lib.rs
@@ -211,6 +211,16 @@ pub trait VariantNames {
     const VARIANTS: &'static [&'static str];
 }
 
+/// A trait for retrieving a static array containing all the variants in an Enum.
+/// This trait can be autoderived by `strum_macros`. For derived usage, all the
+/// variants in the enumerator need to be unit-types, which means you can't autoderive
+/// enums with inner data in one or more variants. Consider using it alongside
+/// [`EnumDiscriminants`] if you require inner data but still want to have an
+/// static array of variants.
+pub trait StaticVariantsArray: std::marker::Sized + 'static {
+    const ALL_VARIANTS: &'static [Self];
+}
+
 #[cfg(feature = "derive")]
 pub use strum_macros::*;
 
@@ -241,5 +251,6 @@ DocumentMacroRexports! {
     EnumVariantNames,
     FromRepr,
     IntoStaticStr,
-    ToString
+    ToString,
+    StaticVariantsArray
 }
diff --git a/strum_macros/Cargo.toml b/strum_macros/Cargo.toml
index c1b9762..3c51382 100644
--- a/strum_macros/Cargo.toml
+++ b/strum_macros/Cargo.toml
@@ -26,4 +26,4 @@ rustversion = "1.0"
 syn = { version = "2.0", features = ["parsing", "extra-traits"] }
 
 [dev-dependencies]
-strum = "0.25"
+strum = { path = "../strum" }
diff --git a/strum_macros/src/helpers/mod.rs b/strum_macros/src/helpers/mod.rs
index 142ea0b..006194d 100644
--- a/strum_macros/src/helpers/mod.rs
+++ b/strum_macros/src/helpers/mod.rs
@@ -15,6 +15,14 @@ pub fn non_enum_error() -> syn::Error {
     syn::Error::new(Span::call_site(), "This macro only supports enums.")
 }
 
+pub fn non_unit_variant_error() -> syn::Error {
+    syn::Error::new(
+        Span::call_site(),
+        "This macro only supports enums of strictly unit variants. Consider \
+        using it in conjunction with [`EnumDiscriminants`]"
+    )
+}
+
 pub fn strum_discriminants_passthrough_error(span: &impl Spanned) -> syn::Error {
     syn::Error::new(
         span.span(),
diff --git a/strum_macros/src/lib.rs b/strum_macros/src/lib.rs
index 82db12a..4ec505c 100644
--- a/strum_macros/src/lib.rs
+++ b/strum_macros/src/lib.rs
@@ -194,6 +194,37 @@ pub fn variant_names(input: proc_macro::TokenStream) -> proc_macro::TokenStream
     toks.into()
 }
 
+/// Adds a static array with all of the Enum's variants.
+/// 
+/// Implements `Strum::StaticVariantsArray` which adds an associated constant `ALL_VARIANTS`.
+/// This constant contains an array with all the variants of the enumerator.
+/// 
+/// This trait can only be autoderived if the enumerator is composed only of unit-type variants,
+/// meaning that the variants must not have any data.
+/// 
+/// ```
+/// use strum::StaticVariantsArray;
+/// 
+/// #[derive(StaticVariantsArray, Debug, PartialEq, Eq)]
+/// enum Op {
+///     Add,
+///     Sub,
+///     Mul,
+///     Div,
+/// }
+/// 
+/// assert_eq!(Op::ALL_VARIANTS, &[Op::Add, Op::Sub, Op::Mul, Op::Div]);
+/// ```
+#[proc_macro_derive(StaticVariantsArray, attributes(strum))]
+pub fn static_variants_array(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse_macro_input!(input as DeriveInput);
+
+    let toks = macros::static_variants_array::static_variants_array_inner(&ast)
+        .unwrap_or_else(|err| err.to_compile_error());
+    debug_print_generated(&ast, &toks);
+    toks.into()
+}
+
 #[proc_macro_derive(AsStaticStr, attributes(strum))]
 #[deprecated(
     since = "0.22.0",
@@ -690,7 +721,7 @@ pub fn enum_properties(input: proc_macro::TokenStream) -> proc_macro::TokenStrea
 /// // Bring trait into scope
 /// use std::str::FromStr;
 /// use strum::{IntoEnumIterator, EnumMessage};
-/// use strum_macros::{EnumDiscriminants, EnumIter, EnumString, EnumMessage};
+/// use strum_macros::{EnumDiscriminants, EnumIter, EnumString};
 ///
 /// #[derive(Debug)]
 /// struct NonDefault;
diff --git a/strum_macros/src/macros/mod.rs b/strum_macros/src/macros/mod.rs
index 8df8cd6..9486e01 100644
--- a/strum_macros/src/macros/mod.rs
+++ b/strum_macros/src/macros/mod.rs
@@ -7,6 +7,7 @@ pub mod enum_properties;
 pub mod enum_try_as;
 pub mod enum_variant_names;
 pub mod from_repr;
+pub mod static_variants_array;
 
 mod strings;
 
diff --git a/strum_macros/src/macros/static_variants_array.rs b/strum_macros/src/macros/static_variants_array.rs
new file mode 100644
index 0000000..2b4130c
--- /dev/null
+++ b/strum_macros/src/macros/static_variants_array.rs
@@ -0,0 +1,36 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{Data, DeriveInput, Fields};
+
+use crate::helpers::{non_enum_error, HasTypeProperties, non_unit_variant_error};
+
+pub fn static_variants_array_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
+    let name = &ast.ident;
+    let gen = &ast.generics;
+    let (impl_generics, ty_generics, where_clause) = gen.split_for_impl();
+
+    let variants = match &ast.data {
+        Data::Enum(v) => &v.variants,
+        _ => return Err(non_enum_error()),
+    };
+
+    let type_properties = ast.get_type_properties()?;
+    let strum_module_path = type_properties.crate_module_path();
+
+    let idents = variants
+        .iter()
+        .cloned()
+        .map(|v| {
+            match v.fields {
+                Fields::Unit => Ok(v.ident),
+                _ => Err(non_unit_variant_error())
+            }
+        })
+        .collect::<syn::Result<Vec<_>>>()?;
+
+    Ok(quote! {
+        impl #impl_generics #strum_module_path::StaticVariantsArray for #name #ty_generics #where_clause {
+            const ALL_VARIANTS: &'static [Self] = &[ #(#name::#idents),* ];
+        }
+    })
+}


Test patch:
diff --git a/strum_tests/tests/static_variants_array.rs b/strum_tests/tests/static_variants_array.rs
new file mode 100644
index 0000000..f5f72c1
--- /dev/null
+++ b/strum_tests/tests/static_variants_array.rs
@@ -0,0 +1,87 @@
+use strum::{StaticVariantsArray, EnumDiscriminants};
+
+mod core {} // ensure macros call `::core`
+
+#[test]
+fn simple() {
+    #[derive(StaticVariantsArray, PartialEq, Eq, Debug)]
+    enum Operation {
+        Add,
+        Sub,
+        Mul,
+        Div,
+    }
+
+    assert_eq!(
+        Operation::ALL_VARIANTS,
+        &[
+            Operation::Add,
+            Operation::Sub,
+            Operation::Mul,
+            Operation::Div,
+        ]
+    );
+}
+
+#[test]
+fn in_enum_discriminants() {
+    #[allow(dead_code)]
+    #[derive(EnumDiscriminants)]
+    #[strum_discriminants(derive(StaticVariantsArray))]
+    #[strum_discriminants(name(GeometricShapeDiscriminants))]
+    enum GeometricShape {
+        Point,
+        Circle(i32),
+        Rectangle {
+            width: i32,
+            height: i32,
+        }
+    }
+
+    assert_eq!(
+        GeometricShapeDiscriminants::ALL_VARIANTS,
+        &[
+            GeometricShapeDiscriminants::Point,
+            GeometricShapeDiscriminants::Circle,
+            GeometricShapeDiscriminants::Rectangle,
+        ]
+    );
+}
+
+#[test]
+fn empty_enum() {
+    #[derive(StaticVariantsArray, PartialEq, Eq, Debug)]
+    enum Empty {}
+
+    assert_eq!(
+        Empty::ALL_VARIANTS,
+        &[],
+    );
+}
+
+#[test]
+fn variants_with_values() {
+    #[derive(StaticVariantsArray, PartialEq, Eq, Debug)]
+    enum WeekDay {
+        Sunday = 0,
+        Monday = 1,
+        Tuesday = 2,
+        Wednesday = 3,
+        Thursday = 4,
+        Friday = 5,
+        Saturday = 6,
+    }
+
+    assert_eq!(
+        WeekDay::ALL_VARIANTS,
+        &[
+            WeekDay::Sunday,
+            WeekDay::Monday,
+            WeekDay::Tuesday,
+            WeekDay::Wednesday,
+            WeekDay::Thursday,
+            WeekDay::Friday,
+            WeekDay::Saturday,
+        ],
+    );
+}
\ No newline at end of file