From a630404e1848ae178d80bbfeb2c41e48d36cbfc3 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
---
 derive/src/traits.rs  | 34 +++++++++++++++++++++++++++++++++-
 derive/tests/basic.rs |  8 ++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/derive/src/traits.rs b/derive/src/traits.rs
index abe27f1..65aca40 100644
--- a/derive/src/traits.rs
+++ b/derive/src/traits.rs
@@ -124,11 +124,43 @@ impl Derivable for Zeroable {
     Ok(syn::parse_quote!(#crate_name::Zeroable))
   }
 
+  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::Enum(_) => Ok(quote!()),
     }
   }
 
diff --git a/derive/tests/basic.rs b/derive/tests/basic.rs
index d3d589d..dfb1ff6 100644
--- a/derive/tests/basic.rs
+++ b/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 {