Repository: sharksforarms/deku
Pull request #139 (2020-11-26 20:32:03)

Enum improvements: Default to discriminant value
Addresses some thoughts from: https://github.com/sharksforarms/deku/issues/83


Patch:
diff --git a/deku-derive/src/lib.rs b/deku-derive/src/lib.rs
index 13c38b9..710400f 100644
--- a/deku-derive/src/lib.rs
+++ b/deku-derive/src/lib.rs
@@ -400,6 +400,15 @@ impl VariantData {
             ));
         }
 
+        if let Some(id) = &receiver.id {
+            if id.to_string() == "_" {
+                return Err((
+                    receiver.ident.span(),
+                    "error: `id_pat` should be used for `_`",
+                ));
+            }
+        }
+
         Ok(())
     }
 }
diff --git a/deku-derive/src/macros/deku_read.rs b/deku-derive/src/macros/deku_read.rs
index 5c781ad..bb78038 100644
--- a/deku-derive/src/macros/deku_read.rs
+++ b/deku-derive/src/macros/deku_read.rs
@@ -133,12 +133,10 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
 
     let magic_read = emit_magic_read(input)?;
 
-    let mut variant_matches = vec![];
     let mut has_default_match = false;
+    let mut pre_match_tokens = vec![];
+    let mut variant_matches = vec![];
 
-    /*
-    FIXME: The loop body is too big
-    */
     for variant in variants {
         // check if the first field has an ident, if not, it's a unnamed struct
         let variant_is_named = variant
@@ -148,16 +146,28 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
             .and_then(|v| v.ident.as_ref())
             .is_some();
 
-        let variant_id = if let Some(variant_id) = &variant.id {
-            variant_id.clone()
+        let (consume_id, variant_id) = if let Some(variant_id) = &variant.id {
+            (true, variant_id.clone())
         } else if let Some(variant_id_pat) = &variant.id_pat {
-            variant_id_pat.clone()
+            (false, variant_id_pat.clone())
+        } else if variant.fields.style.is_unit() {
+            let ident = &variant.ident;
+            let internal_ident = gen_internal_field_ident(&quote!(#ident));
+            pre_match_tokens.push(quote! {
+                let #internal_ident = #id_type::try_from(Self::#ident as isize)?;
+            });
+            (true, quote! { _ if variant_id == #internal_ident })
         } else {
-            // id attribute not provided, treat it as a catch-all default
-            has_default_match = true;
-            quote! { _ }
+            return Err(syn::Error::new(
+                variant.ident.span(),
+                "DekuRead: `id` must be specified on non-unit variants",
+            ));
         };
 
+        if variant_id.to_string() == "_" {
+            has_default_match = true;
+        }
+
         let variant_ident = &variant.ident;
         let variant_reader = &variant.reader;
 
@@ -178,7 +188,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
                 super::gen_enum_init(variant_is_named, variant_ident, internal_fields);
 
             // if we're consuming an id, set the rest to new_rest before reading the variant
-            let new_rest = if variant.id.is_some() {
+            let new_rest = if consume_id {
                 quote! {
                     rest = new_rest;
                 }
@@ -206,7 +216,13 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
     if !has_default_match {
         variant_matches.push(quote! {
             _ => {
-                return Err(DekuError::Parse(format!("Could not match enum variant id = {:?} on enum `{}`", variant_id, #ident_as_string)));
+                return Err(DekuError::Parse(
+                            format!(
+                                "Could not match enum variant id = {:?} on enum `{}`",
+                                variant_id,
+                                #ident_as_string
+                            )
+                        ));
             }
         });
     }
@@ -227,6 +243,8 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
     let variant_read = quote! {
         #variant_id_read
 
+        #(#pre_match_tokens)*
+
         let value = match variant_id {
             #(#variant_matches),*
         };
@@ -274,6 +292,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
     };
 
     tokens.extend(quote! {
+        #[allow(non_snake_case)]
         impl #imp DekuRead<#ctx_types> for #ident #wher {
             fn read<'a>(input: &'a BitSlice<Msb0, u8>, #ctx_arg) -> Result<(&'a BitSlice<Msb0, u8>, Self), DekuError> {
                 #read_body
@@ -285,6 +304,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
         let read_body = wrap_default_ctx(read_body, &input.ctx, &input.ctx_default);
 
         tokens.extend(quote! {
+            #[allow(non_snake_case)]
             impl #imp DekuRead for #ident #wher {
                 fn read<'a>(input: &'a BitSlice<Msb0, u8>, _: ()) -> Result<(&'a BitSlice<Msb0, u8>, Self), DekuError> {
                     #read_body
@@ -463,6 +483,7 @@ pub fn emit_from_bytes(
 ) -> TokenStream {
     quote! {
         impl #imp DekuContainerRead for #ident #wher {
+            #[allow(non_snake_case)]
             fn from_bytes(input: (&[u8], usize)) -> Result<((&[u8], usize), Self), DekuError> {
                 #body
             }
diff --git a/deku-derive/src/macros/deku_write.rs b/deku-derive/src/macros/deku_write.rs
index c6ffc8a..91460d5 100644
--- a/deku-derive/src/macros/deku_write.rs
+++ b/deku-derive/src/macros/deku_write.rs
@@ -161,9 +161,6 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
     let mut variant_writes = vec![];
     let mut variant_updates = vec![];
 
-    /*
-        FIXME: The loop body is too big.
-    */
     for variant in variants {
         // check if the first field has an ident, if not, it's a unnamed struct
         let variant_is_named = variant
@@ -192,14 +189,22 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
             }
         } else if id_type.is_some() {
             if let Some(variant_id) = &variant.id {
-                let variant_id: TokenStream = variant_id.clone();
-
                 quote! {
                     let mut variant_id: #id_type = #variant_id;
                     variant_id.write(output, (#id_args))?;
                 }
-            } else {
+            } else if variant.id_pat.is_some() {
                 quote! {}
+            } else if variant.fields.style.is_unit() {
+                quote! {
+                    let mut variant_id: #id_type = Self::#variant_ident as #id_type;
+                    variant_id.write(output, (#id_args))?;
+                }
+            } else {
+                return Err(syn::Error::new(
+                    variant.ident.span(),
+                    "DekuWrite: `id` must be specified on non-unit variants",
+                ));
             }
         } else {
             // either `id` or `type` needs to be specified
@@ -298,7 +303,7 @@ fn emit_enum(input: &DekuData) -> Result<TokenStream, syn::Error> {
         Ok(())
     };
 
-    // avoid outputing `use core::convert::TryInto` if update() function is empty
+    // avoid outputting `use core::convert::TryInto` if update() function is empty
     let update_use = check_update_use(&variant_updates);
 
     tokens.extend(quote! {
diff --git a/src/attributes.rs b/src/attributes.rs
index c6bcf6c..98684d7 100644
--- a/src/attributes.rs
+++ b/src/attributes.rs
@@ -743,7 +743,9 @@ assert_eq!(ret_write, data)
 Specify the identifier of the enum variant, must be paired with [type](#type)
 or [id (top-level)](#id-top-level)
 
-**Note**: If no `id` is specified, the variant is treated as the "catch-all".
+**Note**:
+    - If no `id` is specified, it is defaulted to the discriminant value.
+    - The discriminant value is retreived using the `as` keyword.
 
 Example:
 ```rust
@@ -756,13 +758,9 @@ enum DekuTest {
     VariantA(u8),
     #[deku(id = "0x02")]
     VariantB(u8, u16),
-
-    VariantCatchAll { // Catch-all variant
-        type_: u8
-    },
 }
 
-let data: Vec<u8> = vec![0x01, 0xFF, 0x02, 0xAB, 0xEF, 0xBE, 0xFF];
+let data: Vec<u8> = vec![0x01, 0xFF, 0x02, 0xAB, 0xEF, 0xBE];
 
 let (rest, value) = DekuTest::from_bytes((data.as_ref(), 0)).unwrap();
 
@@ -783,18 +781,40 @@ assert_eq!(
 
 let variant_bytes: Vec<u8> = value.try_into().unwrap();
 assert_eq!(vec![0x02, 0xAB, 0xEF, 0xBE], variant_bytes);
+```
+
+Example discriminant
+```rust
+# use deku::prelude::*;
+# use std::convert::{TryInto, TryFrom};
+# #[derive(PartialEq, Debug, DekuRead, DekuWrite)]
+#[deku(type = "u8")]
+enum DekuTest {
+    VariantA = 0x01,
+    VariantB,
+}
+
+let data: Vec<u8> = vec![0x01, 0x02];
+
+let (rest, value) = DekuTest::from_bytes((data.as_ref(), 0)).unwrap();
+
+assert_eq!(
+    DekuTest::VariantA,
+    value
+);
+
+let variant_bytes: Vec<u8> = value.try_into().unwrap();
+assert_eq!(vec![0x01], variant_bytes);
 
 let (rest, value) = DekuTest::from_bytes(rest).unwrap();
-# assert_eq!(0, rest.0.len());
-# assert_eq!(0, rest.1);
 
 assert_eq!(
-    DekuTest::VariantCatchAll { type_: 0xFF },
+    DekuTest::VariantB,
     value
 );
 
 let variant_bytes: Vec<u8> = value.try_into().unwrap();
-assert_eq!(vec![0xFF], variant_bytes);
+assert_eq!(vec![0x02], variant_bytes);
 ```
 
 # id_pat
@@ -816,19 +836,31 @@ enum DekuTest {
     VariantB {
         id: u8
     },
+    #[deku(id_pat = "_")]
+    VariantC(u8),
 }
 
-let data: Vec<u8> = vec![0x02];
+let data: Vec<u8> = vec![0x03, 0xFF];
 
 let (rest, value) = DekuTest::from_bytes((data.as_ref(), 0)).unwrap();
 
 assert_eq!(
-    DekuTest::VariantB { id: 0x02 },
+    DekuTest::VariantB { id: 0x03 },
+    value
+);
+
+let variant_bytes: Vec<u8> = value.try_into().unwrap();
+assert_eq!(vec![0x03], variant_bytes);
+
+let (rest, value) = DekuTest::from_bytes(rest).unwrap();
+
+assert_eq!(
+    DekuTest::VariantC(0xFF),
     value
 );
 
 let variant_bytes: Vec<u8> = value.try_into().unwrap();
-assert_eq!(data, variant_bytes);
+assert_eq!(vec![0xFF], variant_bytes);
 ```
 
 # type
diff --git a/src/lib.rs b/src/lib.rs
index c56e09d..6963400 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -148,8 +148,14 @@ the incoming data.
 First the "type" is read using the `type`, then is matched against the
 variants given `id`. What happens after is the same as structs!
 
-This is implemented with the [id](/attributes/index.html#id) and
-[type](attributes#type) attributes.
+This is implemented with the [id](/attributes/index.html#id),
+[id_pat](/attributes/index.html#id_pat) and
+[type](attributes#type) attributes. See these for more examples.
+
+If no `id` is specified, the variant will default to it's discriminant value.
+
+If no variant can be matched, a [DekuError::Parse](crate::error::DekuError)
+error will be returned.
 
 Example:
 


Test patch:
diff --git a/tests/test_compile/cases/enum_validation.rs b/tests/test_compile/cases/enum_validation.rs
index 71e6bb8..9cd2008 100644
--- a/tests/test_compile/cases/enum_validation.rs
+++ b/tests/test_compile/cases/enum_validation.rs
@@ -58,4 +58,21 @@ enum Test9 {
     A
 }
 
+// test `id` cannot be `_`
+#[derive(DekuRead)]
+#[deku(type = "u8")]
+enum Test10 {
+    #[deku(id = "_")]
+    A
+}
+
+// test missing `id`
+#[derive(DekuRead)]
+#[deku(type = "u8")]
+enum Test11 {
+    #[deku(id = "1")]
+    A,
+    B(u8)
+}
+
 fn main() {}
diff --git a/tests/test_compile/cases/enum_validation.stderr b/tests/test_compile/cases/enum_validation.stderr
index 3344cae..05ea52f 100644
--- a/tests/test_compile/cases/enum_validation.stderr
+++ b/tests/test_compile/cases/enum_validation.stderr
@@ -55,3 +55,15 @@ error: error: cannot use `bytes` with `id`
    |
 57 | enum Test9 {
    |      ^^^^^
+
+error: error: `id_pat` should be used for `_`
+  --> $DIR/enum_validation.rs:66:5
+   |
+66 |     A
+   |     ^
+
+error: DekuRead: `id` must be specified on non-unit variants
+  --> $DIR/enum_validation.rs:75:5
+   |
+75 |     B(u8)
+   |     ^
diff --git a/tests/test_enum.rs b/tests/test_enum.rs
index 02951c0..b87b56e 100644
--- a/tests/test_enum.rs
+++ b/tests/test_enum.rs
@@ -26,10 +26,8 @@ pub enum TestEnum {
         #[deku(count = "field_0")] Vec<u8>,
     ),
 
-    VarDefault {
-        id: u8,
-        value: u8,
-    },
+    #[deku(id_pat = "_")]
+    VarDefault { id: u8, value: u8 },
 }
 
 #[rstest(input,expected,
@@ -61,6 +59,29 @@ fn test_enum_error() {
     }
 
     let test_data: Vec<u8> = [0x02, 0x02].to_vec();
-
     let _ret_read = TestEnum::try_from(test_data.as_ref()).unwrap();
 }
+
+#[derive(PartialEq, Debug, DekuRead, DekuWrite)]
+#[deku(type = "u8")]
+pub enum TestEnumDiscriminant {
+    VarA = 0x00,
+    VarB,
+    VarC = 0x02,
+}
+
+#[rstest(input, expected,
+    case(&hex!("00"), TestEnumDiscriminant::VarA),
+    case(&hex!("01"), TestEnumDiscriminant::VarB),
+    case(&hex!("02"), TestEnumDiscriminant::VarC),
+
+    #[should_panic(expected = "Could not match enum variant id = 3 on enum `TestEnumDiscriminant`")]
+    case(&hex!("03"), TestEnumDiscriminant::VarA),
+)]
+fn test_enum_discriminant(input: &[u8], expected: TestEnumDiscriminant) {
+    let ret_read = TestEnumDiscriminant::try_from(input).unwrap();
+    assert_eq!(expected, ret_read);
+
+    let ret_write: Vec<u8> = ret_read.try_into().unwrap();
+    assert_eq!(input.to_vec(), ret_write);
+}