From a66c79ccb39e8320786bac75b2c2a54edfe3291a Mon Sep 17 00:00:00 2001
From: Tom Dohrmann <t.dohrmann@enlyze.com>
Date: Tue, 28 May 2024 19:53:32 +0200
Subject: [PATCH] support deriving Zeroable for fieldless enums (#233)

* support deriving Zeroable for fieldless enums

* add test for deriving Zeroable on enum
---
 bytemuck_derive/src/traits.rs  | 42 ++++++++++++++++++++++++++++------
 bytemuck_derive/tests/basic.rs |  8 +++++++
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/bytemuck_derive/src/traits.rs b/bytemuck_derive/src/traits.rs
index 7924f59..49c0333 100644
--- a/bytemuck_derive/src/traits.rs
+++ b/bytemuck_derive/src/traits.rs
@@ -127,15 +127,43 @@ impl Derivable for Zeroable {
     Ok(syn::parse_quote!(#crate_name::Zeroable))
   }
 
-  fn asserts(
-    input: &DeriveInput, crate_name: &TokenStream,
-  ) -> Result<TokenStream> {
+  fn check_attributes(ty: &Data, attributes: &[Attribute]) -> Result<()> {
+    let repr = get_repr(attributes)?;
+    match ty {
+      Data::Struct(_) => Ok(()),
+      Data::Enum(DataEnum { variants,.. }) => {
+        if !repr.repr.is_integer() {
+          bail!("Zeroable requires the enum to be an explicit #[repr(Int)]")
+        }
+
+        if variants.iter().any(|variant| !variant.fields.is_empty()) {
+          bail!("Only fieldless enums are supported for Zeroable")
+        }
+
+        let iter = VariantDiscriminantIterator::new(variants.iter());
+        let mut has_zero_variant = false;
+        for res in iter {
+          let discriminant = res?;
+          if discriminant == 0 {
+            has_zero_variant = true;
+            break;
+          }
+        }
+        if !has_zero_variant {
+          bail!("No variant's discriminant is 0")
+        }
+
+        Ok(())
+      },
+      Data::Union(_) => Ok(())
+    }
+  }
+
+  fn asserts(input: &DeriveInput, crate_name: &TokenStream) -> Result<TokenStream> {
     match &input.data {
       Data::Union(_) => Ok(quote!()), // unions are always `Zeroable`
-      Data::Struct(_) => {
-        generate_fields_are_trait(input, Self::ident(input, crate_name)?)
-      }
-      Data::Enum(_) => bail!("Deriving Zeroable is not supported for enums"),
+      Data::Struct(_) => generate_fields_are_trait(input, Self::ident(input, crate_name)?),
+      Data::Enum(_) => Ok(quote!()),
     }
   }
 
diff --git a/bytemuck_derive/tests/basic.rs b/bytemuck_derive/tests/basic.rs
index b7960be..73ad59d 100644
--- a/bytemuck_derive/tests/basic.rs
+++ b/bytemuck_derive/tests/basic.rs
@@ -50,6 +50,14 @@ struct ZeroGeneric<T: bytemuck::Zeroable> {
   a: T,
 }
 
+#[derive(Zeroable)]
+#[repr(u8)]
+enum ZeroEnum {
+  A = 0,
+  B = 1,
+  C = 2,
+}
+
 #[derive(TransparentWrapper)]
 #[repr(transparent)]
 struct TransparentSingle {