Add Enum Map Derive Macro

README.md

@@ -42,6 +42,7 @@ Strum has implemented the following macros:
 | [IntoStaticStr] | Implements `From<MyEnum> for &'static str` on an enum |
 | [EnumVariantNames] | Adds an associated `VARIANTS` constant which is an array of discriminant names |
 | [EnumIter] | Creates a new type that iterates of the variants of an enum. |
+| [EnumMap] | Creates a new type that stores an item of a specified type for each variant of the enum. |
 | [EnumProperty] | Add custom properties to enum variants. |
 | [EnumMessage] | Add a verbose message to an enum variant. |
 | [EnumDiscriminants] | Generate a new type with only the discriminant names. |

strum_macros/src/lib.rs

@@ -409,6 +409,44 @@ pub fn enum_is(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     toks.into()
 }
 
+/// Creates a new type that maps all the variants of an enum to another generic value.
+///
+/// This macro does not support any additional data on your variants.
+/// The macro creates a new type called `YourEnumMap<T>`.
+/// The map has a field of type `T` for each variant of `YourEnum`. The map automatically implements `Index<T>` and `IndexMut<T>`.
+/// ```
+/// use strum_macros::EnumMap;
+///
+/// #[derive(EnumMap)]
+/// enum Color {
+///     Red,
+///     Yellow,
+///     Green,
+///     Blue,
+/// }
+///
+/// assert_eq!(ColorMap::default(), ColorMap::new(0, 0, 0, 0));
+/// assert_eq!(ColorMap::filled(2), ColorMap::new(2, 2, 2, 2));
+/// assert_eq!(ColorMap::from_closure(|_| 3), ColorMap::new(3, 3, 3, 3));
+/// assert_eq!(ColorMap::default().transform(|_, val| val + 2), ColorMap::new(2, 2, 2, 2));
+///
+/// let mut complex_map = ColorMap::from_closure(|color| match color {
+///     Color::Red => 0,
+///     _ => 3
+/// });
+/// complex_map[Color::Green] = complex_map[Color::Red];
+/// assert_eq!(complex_map, ColorMap::new(0, 3, 0, 3));
+///
+/// ```
+#[proc_macro_derive(EnumMap, attributes(strum))]
+pub fn enum_map(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let ast = syn::parse_macro_input!(input as DeriveInput);
+
+    let toks = macros::enum_map::enum_map_inner(&ast).unwrap_or_else(|err| err.to_compile_error());
+    debug_print_generated(&ast, &toks);
+    toks.into()
+}
+
 /// Add a function to enum that allows accessing variants by its discriminant
 ///
 /// This macro adds a standalone function to obtain an enum variant by its discriminant. The macro adds

strum_macros/src/macros/enum_map.rs

@@ -0,0 +1,217 @@
+use proc_macro2::{Span, TokenStream};
+use quote::{format_ident, quote};
+use syn::{spanned::Spanned, Data, DeriveInput, Fields, Ident};
+
+use crate::helpers::{non_enum_error, HasStrumVariantProperties};
+
+pub fn enum_map_inner(ast: &DeriveInput) -> syn::Result<TokenStream> {
+    let name = &ast.ident;
+    let gen = &ast.generics;
+    let vis = &ast.vis;
+    let mut doc_comment = format!("A map over the variants of [{}]", name);
+
+    if gen.lifetimes().count() > 0 {
+        return Err(syn::Error::new(
+            Span::call_site(),
+            "This macro doesn't support enums with lifetimes.",
+        ));
+    }
+
+    let Data::Enum(data_enum) = &ast.data else {
+        return Err(non_enum_error())
+    };
+    let map_name = syn::parse_str::<Ident>(&format!("{}Map", name)).unwrap();
+
+    let variants = &data_enum.variants;
+
+    // the identifiers of each variant, in PascalCase
+    let mut pascal_idents = Vec::new();
+    // the identifiers of each struct field, in snake_case
+    let mut snake_idents = Vec::new();
+    // match arms in the form `MyEnumMap::Variant => &self.variant,`
+    let mut get_matches = Vec::new();
+    // match arms in the form `MyEnumMap::Variant => &mut self.variant,`
+    let mut get_matches_mut = Vec::new();
+    // match arms in the form `MyEnumMap::Variant => self.variant = new_value`
+    let mut set_matches = Vec::new();
+    // struct fields of the form `variant: func(MyEnum::Variant),*
+    let mut closure_fields = Vec::new();
+    // struct fields of the form `variant: func(MyEnum::Variant, self.variant),`
+    let mut transform_fields = Vec::new();
+
+    // identifiers for disabled variants
+    let mut disabled_variants = Vec::new();
+    // match arms for disabled variants
+    let mut disabled_matches = Vec::new();
+
+    for variant in variants {
+        // skip disabled variants
+        if variant.get_variant_properties()?.disabled.is_some() {
+            let disabled_ident = &variant.ident;
+            let panic_message =
+                format!("Can't use `{disabled_ident}` with `{map_name}` - variant is disabled for Strum features");
+            disabled_variants.push(disabled_ident);
+            disabled_matches.push(quote!(#name::#disabled_ident => panic!(#panic_message),));
+            continue;
+        }
+
+        // Error on fields with data
+        let Fields::Unit = &variant.fields else {
+            return Err(syn::Error::new(
+                variant.fields.span(),
+                "This macro doesn't support enums with non-unit variants",
+            ))
+        };
+
+        let pascal_case = &variant.ident;
+        // switch PascalCase to snake_case. This naively assumes they use PascalCase
+        let snake_case = format_ident!(
+            "{}",
+            pascal_case
+                .to_string()
+                .chars()
+                .enumerate()
+                .fold(String::new(), |mut s, (i, c)| {
+                    if c.is_uppercase() && i > 0 {
+                        s.push('-');
+                    }
+                    s.push(c.to_ascii_lowercase());
+                    s
+                })
+        );
+
+        get_matches.push(quote! {#name::#pascal_case => &self.#snake_case,});
+        get_matches_mut.push(quote! {#name::#pascal_case => &mut self.#snake_case,});
+        set_matches.push(quote! {#name::#pascal_case => self.#snake_case = new_value,});
+        closure_fields.push(quote! {#snake_case: func(#name::#pascal_case),});
+        transform_fields.push(quote! {#snake_case: func(#name::#pascal_case, &self.#snake_case),});
+        pascal_idents.push(pascal_case);
+        snake_idents.push(snake_case);
+    }
+
+    // if the index operation can panic, add that to the documentation
+    if !disabled_variants.is_empty() {
+        doc_comment.push_str(&format!(
+            "\n# Panics\nIndexing `{map_name}` with any of the following variants will cause a panic:"
+        ));
+        for variant in disabled_variants {
+            doc_comment.push_str(&format!("\n\n- `{name}::{variant}`"));
+        }
+    }
+
+    let doc_new = format!("Create a new {map_name} with a value for each variant of {name}");
+    let doc_closure =
+        format!("Create a new {map_name} by running a function on each variant of `{name}`");
+    let doc_transform = format!("Create a new `{map_name}` by running a function on each variant of `{name}` and the corresponding value in the current `{map_name}`");
+    let doc_filled = format!("Create a new `{map_name}` with the same value in each field.");
+    let doc_option_all = format!("Converts `{map_name}<Option<T>>` into `Option<{map_name}<T>>`. Returns `Some` if all fields are `Some`, otherwise returns `None`.");
+    let doc_result_all_ok = format!("Converts `{map_name}<Result<T, E>>` into `Option<{map_name}>`. Returns `Some` if all fields are `Ok`, otherwise returns `None`.");
+
+    Ok(quote! {
+        #[doc = #doc_comment]
+        #[allow(
+            missing_copy_implementations,
+        )]
+        #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
+        #vis struct #map_name<T> {
+            #(#snake_idents: T,)*
+        }
+
+        impl<T: Clone> #map_name<T> {
+            #[doc = #doc_filled]
+            #vis fn filled(value: T) -> #map_name<T> {
+                #map_name {
+                    #(#snake_idents: value.clone(),)*
+                }
+            }
+        }
+
+        impl<T> #map_name<T> {
+            #[doc = #doc_new]
+            #vis fn new(
+                #(#snake_idents: T,)*
+            ) -> #map_name<T> {
+                #map_name {
+                    #(#snake_idents,)*
+                }
+            }
+
+            #[doc = #doc_closure]
+            #vis fn from_closure<F: Fn(#name)->T>(func: F) -> #map_name<T> {
+              #map_name {
+                #(#closure_fields)*
+              }
+            }
+
+            #[doc = #doc_transform]
+            #vis fn transform<U, F: Fn(#name, &T)->U>(&self, func: F) -> #map_name<U> {
+              #map_name {
+                #(#transform_fields)*
+              }
+            }
+
+            // // E.g. so that if you're using EnumIter as well, these functions work nicely
+            // fn get(&self, variant: #name) -> &T {
+            //     match variant {
+            //         #(#get_matches)*
+            //     }
+            // }
+
+            // fn set(&mut self, variant: #name, new_value: T) {
+            //     match variant {
+            //         #(#set_matches)*
+            //     }
+            // }
+        }
+
+        impl<T> core::ops::Index<#name> for #map_name<T> {
+            type Output = T;
+
+            fn index(&self, idx: #name) -> &T {
+                match idx {
+                    #(#get_matches)*
+                    #(#disabled_matches)*
+                }
+            }
+        }
+
+        impl<T> core::ops::IndexMut<#name> for #map_name<T> {
+            fn index_mut(&mut self, idx: #name) -> &mut T {
+                match idx {
+                    #(#get_matches_mut)*
+                    #(#disabled_matches)*
+                }
+            }
+        }
+
+        impl<T> #map_name<Option<T>> {
+            #[doc = #doc_option_all]
+            #vis fn all(self) -> Option<#map_name<T>> {
+                if let #map_name {
+                    #(#snake_idents: Some(#snake_idents),)*
+                } = self {
+                    Some(#map_name {
+                        #(#snake_idents,)*
+                    })
+                } else {
+                    None
+                }
+            }
+        }
+
+        impl<T, E> #map_name<Result<T, E>> {
+            #[doc = #doc_result_all_ok]
+            #vis fn all_ok(self) -> Option<#map_name<T>> {
+                if let #map_name {
+                    #(#snake_idents: Ok(#snake_idents),)*
+                } = self {
+                    Some(#map_name {
+                        #(#snake_idents,)*
+                    })
+                } else {
+                    None
+                }
+            }
+        }
+    })
+}

strum_macros/src/macros/mod.rs

@@ -2,6 +2,7 @@ pub mod enum_count;
 pub mod enum_discriminants;
 pub mod enum_is;
 pub mod enum_iter;
+pub mod enum_map;
 pub mod enum_messages;
 pub mod enum_properties;
 pub mod enum_variant_names;

strum_tests/tests/enum_map.rs

@@ -0,0 +1,87 @@
+use strum::EnumMap;
+
+#[derive(EnumMap)]
+enum Color {
+    Red,
+    Yellow,
+    Green,
+    #[strum(disabled)]
+    Teal,
+    Blue,
+    #[strum(disabled)]
+    Indigo,
+}
+
+#[test]
+fn default() {
+    assert_eq!(ColorMap::default(), ColorMap::new(0, 0, 0, 0));
+}
+
+#[test]
+#[should_panic]
+fn disabled() {
+    let _ = ColorMap::<u8>::default()[Color::Indigo];
+}
+
+#[test]
+fn filled() {
+    assert_eq!(ColorMap::filled(42), ColorMap::new(42, 42, 42, 42));
+}
+
+#[test]
+fn from_closure() {
+    assert_eq!(
+        ColorMap::from_closure(|color| match color {
+            Color::Red => 1,
+            _ => 2,
+        }),
+        ColorMap::new(1, 2, 2, 2)
+    );
+}
+
+#[test]
+fn index() {
+    let map = ColorMap::new(18, 25, 7, 2);
+    assert_eq!(map[Color::Red], 18);
+    assert_eq!(map[Color::Yellow], 25);
+    assert_eq!(map[Color::Green], 7);
+    assert_eq!(map[Color::Blue], 2);
+}
+
+#[test]
+fn index_mut() {
+    let mut map = ColorMap::new(18, 25, 7, 2);
+    map[Color::Green] = 5;
+    map[Color::Red] *= 4;
+    assert_eq!(map[Color::Green], 5);
+    assert_eq!(map[Color::Red], 72);
+}
+
+#[test]
+fn option_all() {
+    let mut map: ColorMap<Option<u8>> = ColorMap::filled(None);
+    map[Color::Red] = Some(64);
+    map[Color::Green] = Some(32);
+    map[Color::Blue] = Some(16);
+
+    assert_eq!(map.clone().all(), None);
+
+    map[Color::Yellow] = Some(8);
+    assert_eq!(map.all(), Some(ColorMap::new(64, 8, 32, 16)));
+}
+
+#[test]
+fn result_all_ok() {
+    let mut map: ColorMap<Result<u8, u8>> = ColorMap::filled(Ok(4));
+    assert_eq!(map.clone().all_ok(), Some(ColorMap::filled(4)));
+
+    map[Color::Red] = Err(22);
+
+    assert_eq!(map.all_ok(), None);
+}
+
+#[test]
+fn transform() {
+    let all_two = ColorMap::filled(2);
+    assert_eq!(all_two.transform(|_, n| *n * 2), ColorMap::filled(4));
+}