Add discriminant() and from_discriminant(num) methods

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Unreleased
 
+# Version 1.2.0 (2024-07-02)
+
+- Add `discriminant()` and `from_discriminant(num)` methods
+
 # Version 1.1.0 (2024-06-27)
 
 - Add `values()` method and document all methods in UnitEnum macro

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "unit-enum"
-version = "1.1.0"
+version = "1.2.0"
 authors = ["Renaud Denis <[email protected]>"]
 description = "A procedural macro for deriving ordinal methods in unit-like enums for Rust."
 license = "MIT OR Apache-2.0"

README.md

@@ -1,13 +1,14 @@
 # UnitEnum Crate Documentation
 
 The `unit-enum` crate provides a procedural macro `UnitEnum` designed to enhance enums in Rust, particularly those
-consisting solely of unit variants. This macro simplifies working with such enums by providing methods to convert
-between enum variants and their ordinal positions, along with utility methods to count the number of variants and iterate over them.
+consisting solely of unit variants. This macro simplifies working with such enums by providing useful utility methods.
 
 ## Features
 
 - `ordinal`: Retrieve the ordinal of an enum variant, starting from 0.
 - `from_ordinal`: Convert an ordinal back to an enum variant, if possible.
+- `discriminant`: Retrieve the discriminant of an enum variant.
+- `from_discriminant`: Convert a discriminant back to an enum variant.
 - `len`: Get the total number of variants in the enum.
 - `values`: Returns an iterator over all variants of the enum, allowing for easy iteration and handling of each variant.
 
@@ -22,7 +23,7 @@ Add the following to your `Cargo.toml`:
 
 ```toml
 [dependencies]
-unit-enum = "1.1.0"
+unit-enum = "1.2.0"
 ```
 
 ## Quick Start
@@ -32,29 +33,44 @@ use unit_enum::UnitEnum;
 
 #[derive(Debug, Clone, Copy, PartialEq, UnitEnum)]
 enum Color {
-  Red,
-  Green,
-  Blue
+    Red = 10,
+    Green,
+    Blue = 45654
 }
 
 fn main() {
-  println!("Ordinal of Green: {:?}", Color::Green.ordinal());
-  // Ordinal of Green: 1
+    println!("Ordinal of Green: {:?}", Color::Green.ordinal());
+    // Ordinal of Green: 1
 
-  println!("Value of ordinal 2: {:?}", Color::from_ordinal(2));
-  // Value of ordinal 2: Some(Blue)
+    println!("Value of ordinal 2: {:?}", Color::from_ordinal(2));
+    // Value of ordinal 2: Some(Blue)
 
-  println!("Number of Color variants: {:?}", Color::len());
-  // Number of Color variants: 3
+    println!("Value of ordinal 4: {:?}", Color::from_ordinal(4));
+    // Value of ordinal 4: None
 
-  println!("List of Color variants: {:?}", Color::values().collect::<Vec<_>>());
-  // List of Color variants: [Red, Green, Blue]
+    println!("Discriminant of Blue: {:?}", Color::Blue.discriminant());
+    // Discriminant of Blue: 45654
+
+    println!("Discriminant of Green: {:?}", Color::Green.discriminant());
+    // Discriminant of Green: 11
+
+    println!("Value of discriminant 10: {:?}", Color::from_discriminant(10));
+    // Value of discriminant 10: Some(Red)
+
+    println!("Value of discriminant 0: {:?}", Color::from_discriminant(0));
+    // Value of discriminant 0: None
+
+    println!("Number of Color variants: {:?}", Color::len());
+    // Number of Color variants: 3
+
+    println!("List of Color variants: {:?}", Color::values().collect::<Vec<_>>());
+    // List of Color variants: [Red, Green, Blue]
 }
 ```
 
 ## Contributing
 
-Contributions are welcome! Please feel free to submit pull requests or open issues on our [GitHub repository](#).
+Contributions are welcome! Please feel free to submit pull requests or open issues on our [GitHub repository](https://github.com/tylium/unit-enum).
 
 ## License
 

examples/example.rs

@@ -2,9 +2,9 @@ use unit_enum::UnitEnum;
 
 #[derive(Debug, Clone, Copy, PartialEq, UnitEnum)]
 enum Color {
-    Red,
+    Red = 10,
     Green,
-    Blue
+    Blue = 45654
 }
 
 fn main() {
@@ -14,6 +14,21 @@ fn main() {
     println!("Value of ordinal 2: {:?}", Color::from_ordinal(2));
     // Value of ordinal 2: Some(Blue)
 
+    println!("Value of ordinal 4: {:?}", Color::from_ordinal(4));
+    // Value of ordinal 4: None
+
+    println!("Discriminant of Blue: {:?}", Color::Blue.discriminant());
+    // Discriminant of Blue: 45654
+
+    println!("Discriminant of Green: {:?}", Color::Green.discriminant());
+    // Discriminant of Green: 11
+
+    println!("Value of discriminant 10: {:?}", Color::from_discriminant(10));
+    // Value of discriminant 10: Some(Red)
+
+    println!("Value of discriminant 0: {:?}", Color::from_discriminant(0));
+    // Value of discriminant 0: None
+
     println!("Number of Color variants: {:?}", Color::len());
     // Number of Color variants: 3
 

src/lib.rs

@@ -1,7 +1,7 @@
 use proc_macro::TokenStream;
 
 use quote::quote;
-use syn::{Data, DeriveInput, Fields, parse_macro_input};
+use syn::{Data, DeriveInput, Expr, Fields, parse_macro_input, punctuated::Punctuated, token::Comma, Variant};
 
 #[proc_macro_derive(UnitEnum)]
 pub fn unit_enum_derive(input: TokenStream) -> TokenStream {
@@ -22,37 +22,54 @@ fn impl_unit_enum(ast: &DeriveInput) -> TokenStream {
     let ordinal_match_arms = variants.iter().enumerate().map(|(index, variant)| {
         let variant_name = &variant.ident;
         match &variant.fields {
-            Fields::Unit => {
-                quote! { #name::#variant_name => #index, }
-            }
+            Fields::Unit => quote! { #name::#variant_name => #index },
             _ => panic!("UnitEnum only supports unit variants (no fields)"),
         }
     });
 
     let from_ordinal_match_arms = variants.iter().enumerate().map(|(index, variant)| {
         let variant_name = &variant.ident;
-        quote! { #index => Some(#name::#variant_name), }
+        quote! { #index => Some(#name::#variant_name) }
     });
 
     let values_arms = (0..num_variants).map(|index| {
         quote! { #name::from_ordinal(#index).unwrap() }
     });
 
+    let (discriminant_match_arms, from_discriminant_match_arms) = generate_discriminant_arms(variants, name);
+
     let gen = quote! {
         impl #name {
             /// Returns the zero-based ordinal of the enum variant.
             pub fn ordinal(&self) -> usize {
                 match self {
-                    #( #ordinal_match_arms )*
+                    #(#ordinal_match_arms,)*
                 }
             }
 
             /// Converts a zero-based ordinal to an enum variant, if possible.
             ///
             /// Returns `None` if the ordinal is out of range.
-             pub fn from_ordinal(ord: usize) -> Option<Self> {
+            pub fn from_ordinal(ord: usize) -> Option<Self> {
                 match ord {
-                    #( #from_ordinal_match_arms )*
+                    #(#from_ordinal_match_arms,)*
+                    _ => None,
+                }
+            }
+
+            /// Returns the discriminant value of the enum variant.
+            pub fn discriminant(&self) -> i32 {
+                match self {
+                    #(#discriminant_match_arms,)*
+                }
+            }
+
+            /// Converts a discriminant value back to an enum variant, if possible.
+            ///
+            /// Returns `None` if the discriminant does not correspond to any variant.
+            pub fn from_discriminant(discr: i32) -> Option<Self> {
+                match discr {
+                    #(#from_discriminant_match_arms,)*
                     _ => None,
                 }
             }
@@ -67,9 +84,31 @@ fn impl_unit_enum(ast: &DeriveInput) -> TokenStream {
             /// This method creates an iterator that yields each variant in
             /// definition order, starting from the first variant.
             pub fn values() -> impl Iterator<Item = Self> {
-                vec![ #( #values_arms ),* ].into_iter()
+                vec![#(#values_arms,)*].into_iter()
             }
         }
     };
     gen.into()
-}
+}
+
+fn generate_discriminant_arms(variants: &Punctuated<Variant, Comma>, name: &syn::Ident) -> (Vec<proc_macro2::TokenStream>, Vec<proc_macro2::TokenStream>) {
+    let mut last_discriminant: Option<Expr> = None;
+    let discriminant_match_arms = variants.iter().map(|variant| {
+        let variant_name = &variant.ident;
+        let discriminant = variant.discriminant.as_ref().map(|(_, expr)| expr.clone())
+            .or_else(|| last_discriminant.clone()).unwrap_or_else(|| syn::parse_quote! { 0 });
+        last_discriminant = Some(syn::parse_quote! { #discriminant + 1 });
+
+        quote! { #name::#variant_name => #discriminant as i32 }
+    }).collect::<Vec<_>>();
+
+    let from_discriminant_match_arms = variants.iter().map(|variant| {
+        let variant_name = &variant.ident;
+        let discriminant = variant.discriminant.as_ref().map(|(_, expr)| expr.clone())
+            .or_else(|| last_discriminant.clone()).unwrap_or_else(|| syn::parse_quote! { 0 });
+
+        quote! { x if x == #discriminant as i32 => Some(#name::#variant_name) }
+    }).collect::<Vec<_>>();
+
+    (discriminant_match_arms, from_discriminant_match_arms)
+}