zd,zv: Value: Add support for enums without explicit values

dbus2 · Feb 25, 2024 · 4d68aae · 4d68aae
1 parent e0a6344
commit 4d68aae
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 37 deletions.
diff --git a/zvariant/src/lib.rs b/zvariant/src/lib.rs
@@ -142,8 +142,8 @@ mod tests {
     use crate::Fd;
     use crate::{
         serialized::{Context, Format},
-        Array, Basic, DeserializeDict, DeserializeValue, Dict, Error, ObjectPath, Result,
-        SerializeDict, SerializeValue, Signature, Str, Structure, Type, Value, BE, LE,
+        Array, Basic, DeserializeDict, DeserializeValue, Dict, Error, ObjectPath, OwnedValue,
+        Result, SerializeDict, SerializeValue, Signature, Str, Structure, Type, Value, BE, LE,
         NATIVE_ENDIAN,
     };
 
@@ -1579,6 +1579,40 @@ mod tests {
         assert_eq!(decoded, s);
     }
 
+    #[test]
+    fn derive_value() {
+        #[repr(u16)]
+        #[derive(Value, OwnedValue, Type, Debug, PartialEq)]
+        enum Enum {
+            Variant1,
+            Variant2,
+            Variant3,
+        }
+
+        assert_eq!(Enum::signature(), u16::signature());
+        assert_eq!(Value::from(Enum::Variant1), Value::U16(0));
+        assert_eq!(Enum::try_from(Value::U16(2)), Ok(Enum::Variant3));
+        assert_eq!(Enum::try_from(Value::U16(4)), Err(Error::IncorrectType));
+
+        #[derive(Value, OwnedValue, Type, Debug, PartialEq)]
+        enum Enum2 {
+            Variant1,
+            Variant2,
+            Variant3,
+        }
+
+        assert_eq!(Enum2::signature(), u32::signature());
+        assert_eq!(Value::from(Enum2::Variant1), Value::U32(0));
+        assert_eq!(Value::from(Enum2::Variant2), Value::U32(1));
+        assert_eq!(Enum2::try_from(Value::U32(1)), Ok(Enum2::Variant2));
+        assert_eq!(Enum2::try_from(Value::U32(2)), Ok(Enum2::Variant3));
+        assert_eq!(Enum2::try_from(Value::U32(3)), Err(Error::IncorrectType));
+        assert_eq!(
+            Enum::try_from(Value::Str("Variant1".into())),
+            Err(Error::IncorrectType)
+        );
+    }
+
     #[test]
     fn serialized_size() {
         let ctxt = Context::new_dbus(LE, 0);

diff --git a/zvariant_derive/src/value.rs b/zvariant_derive/src/value.rs
@@ -238,27 +238,11 @@ fn impl_enum(
     };
 
     let mut variant_names = vec![];
-    let mut variant_values = vec![];
     for variant in &data.variants {
         // Ensure all variants of the enum are unit type
         match variant.fields {
             Fields::Unit => {
                 variant_names.push(&variant.ident);
-                let value = match &variant
-                    .discriminant
-                    .as_ref()
-                    .ok_or_else(|| Error::new(variant.span(), "expected `Name = Value` variants"))?
-                    .1
-                {
-                    Expr::Lit(lit_exp) => &lit_exp.lit,
-                    _ => {
-                        return Err(Error::new(
-                            variant.span(),
-                            "expected `Name = Value` variants",
-                        ))
-                    }
-                };
-                variant_values.push(value);
             }
             _ => return Err(Error::new(variant.span(), "must be a unit variant")),
         }
@@ -271,13 +255,7 @@ fn impl_enum(
                 impl ::std::convert::From<#name> for #zv::Value<'_> {
                     #[inline]
                     fn from(e: #name) -> Self {
-                        let u: #repr = match e {
-                            #(
-                                #name::#variant_names => #variant_values
-                            ),*
-                        };
-
-                        <#zv::Value as ::std::convert::From<_>>::from(u).into()
+                        <#zv::Value as ::std::convert::From<_>>::from(e as #repr).into()
                     }
                 }
             },
@@ -290,14 +268,8 @@ fn impl_enum(
 
                     #[inline]
                     fn try_from(e: #name) -> #zv::Result<Self> {
-                        let u: #repr = match e {
-                            #(
-                                #name::#variant_names => #variant_values
-                            ),*
-                        };
-
                         <#zv::OwnedValue as ::std::convert::TryFrom<_>>::try_from(
-                            <#zv::Value as ::std::convert::From<_>>::from(u)
+                            <#zv::Value as ::std::convert::From<_>>::from(e as #repr)
                         )
                     }
                 }
@@ -313,12 +285,15 @@ fn impl_enum(
             fn try_from(value: #value_type) -> #zv::Result<Self> {
                 let v: #repr = ::std::convert::TryInto::try_into(value)?;
 
-                ::std::result::Result::Ok(match v {
+                ::std::result::Result::Ok(
                     #(
-                        #variant_values => #name::#variant_names
-                     ),*,
-                    _ => return ::std::result::Result::Err(#zv::Error::IncorrectType),
-                })
+                        if v == #name::#variant_names as #repr {
+                            #name::#variant_names
+                        } else
+                    )* {
+                        return ::std::result::Result::Err(#zv::Error::IncorrectType);
+                    }
+                )
             }
         }