Add unit variant derive (#13)

protocol-derive/src/lib.rs

@@ -1,8 +1,8 @@
 extern crate proc_macro;
 
 use crate::parse::{parse_derive_input, DeriveInputParseResult};
-use crate::render::decoder::{render_struct_decoder, render_struct_variant_decoder};
-use crate::render::encoder::{render_struct_encoder, render_struct_variant_encoder};
+use crate::render::decoder::{render_enum_decoder, render_struct_decoder};
+use crate::render::encoder::{render_enum_encoder, render_struct_encoder};
 use proc_macro::TokenStream;
 use syn::parse_macro_input;
 use syn::DeriveInput;
@@ -18,9 +18,7 @@ pub fn derive_encoder(tokens: TokenStream) -> TokenStream {
 
     TokenStream::from(match derive_parse_result {
         DeriveInputParseResult::Struct { name, fields } => render_struct_encoder(name, &fields),
-        DeriveInputParseResult::StructVariant { name, variants } => {
-            render_struct_variant_encoder(name, &variants)
-        }
+        DeriveInputParseResult::Enum { name, variants } => render_enum_encoder(name, &variants),
     })
 }
 
@@ -31,8 +29,6 @@ pub fn derive_decoder(tokens: TokenStream) -> TokenStream {
 
     TokenStream::from(match derive_parse_result {
         DeriveInputParseResult::Struct { name, fields } => render_struct_decoder(name, &fields),
-        DeriveInputParseResult::StructVariant { name, variants } => {
-            render_struct_variant_decoder(name, &variants)
-        }
+        DeriveInputParseResult::Enum { name, variants } => render_enum_decoder(name, &variants),
     })
 }

protocol-derive/src/parse.rs

@@ -2,23 +2,23 @@ use crate::error::{DeriveInputParserError, FieldError};
 use proc_macro2::Ident;
 use std::iter::FromIterator;
 use syn::punctuated::Punctuated;
-use syn::Token;
-use syn::{Data, DeriveInput, Field, Fields, FieldsNamed, Lit, Meta, NestedMeta, Type};
+use syn::{Data, DeriveInput, ExprLit, Field, Fields, FieldsNamed, Lit, Meta, NestedMeta, Type};
 use syn::{Error as SynError, Variant};
+use syn::{Expr, Token};
 
 pub(crate) enum DeriveInputParseResult<'a> {
     Struct {
         name: &'a Ident,
         fields: Vec<FieldData<'a>>,
     },
-    StructVariant {
+    Enum {
         name: &'a Ident,
         variants: Vec<VariantData<'a>>,
     },
 }
 
 pub(crate) struct VariantData<'a> {
-    pub(crate) idx: u8,
+    pub(crate) discriminant: u8,
     pub(crate) name: &'a Ident,
     pub(crate) fields: Vec<FieldData<'a>>,
 }
@@ -53,7 +53,7 @@ pub(crate) fn parse_derive_input(
         Data::Enum(data_enum) => {
             let variants = parse_variants(&data_enum.variants)?;
 
-            Ok(DeriveInputParseResult::StructVariant { name, variants })
+            Ok(DeriveInputParseResult::Enum { name, variants })
         }
         _ => Err(DeriveInputParserError::UnsupportedData),
     }
@@ -70,6 +70,7 @@ fn parse_variants(
 }
 
 fn parse_variant(idx: u8, variant: &Variant) -> Result<VariantData, DeriveInputParserError> {
+    let discriminant = parse_variant_discriminant(variant).unwrap_or(idx);
     let name = &variant.ident;
 
     let fields = match &variant.fields {
@@ -78,7 +79,24 @@ fn parse_variant(idx: u8, variant: &Variant) -> Result<VariantData, DeriveInputP
         _ => Err(DeriveInputParserError::UnnamedDataFields),
     }?;
 
-    Ok(VariantData { idx, name, fields })
+    Ok(VariantData {
+        discriminant,
+        name,
+        fields,
+    })
+}
+
+fn parse_variant_discriminant(variant: &Variant) -> Option<u8> {
+    variant
+        .discriminant
+        .as_ref()
+        .and_then(|(_, expr)| match expr {
+            Expr::Lit(ExprLit {
+                lit: Lit::Int(lit_int),
+                ..
+            }) => lit_int.base10_parse().ok(),
+            _ => None,
+        })
 }
 
 fn parse_fields(named_fields: &FieldsNamed) -> Result<Vec<FieldData>, DeriveInputParserError> {

protocol-derive/src/render/decoder.rs

@@ -24,10 +24,7 @@ pub(crate) fn render_struct_decoder(name: &Ident, fields: &Vec<FieldData>) -> To
     }
 }
 
-pub(crate) fn render_struct_variant_decoder(
-    name: &Ident,
-    variants: &Vec<VariantData>,
-) -> TokenStream2 {
+pub(crate) fn render_enum_decoder(name: &Ident, variants: &Vec<VariantData>) -> TokenStream2 {
     let render_variants = render_variants(variants);
 
     quote! {
@@ -52,26 +49,37 @@ fn render_variants(variants: &Vec<VariantData>) -> TokenStream2 {
 }
 
 fn render_variant(variant: &VariantData) -> TokenStream2 {
-    let idx = variant.idx;
+    if variant.fields.is_empty() {
+        render_unit_variant(variant)
+    } else {
+        render_struct_variant(variant)
+    }
+}
+
+fn render_unit_variant(variant: &VariantData) -> TokenStream2 {
+    let discriminant = variant.discriminant;
+    let name = variant.name;
+
+    quote! {
+        #discriminant => Ok(Self::#name),
+    }
+}
+
+fn render_struct_variant(variant: &VariantData) -> TokenStream2 {
+    let discriminant = variant.discriminant;
     let name = variant.name;
     let fields = &variant.fields;
 
-    if fields.is_empty() {
-        quote! {
-            #idx => Ok(Self::#name),
-        }
-    } else {
-        let field_names_joined_comma = render_field_names_joined_comma(fields);
-        let render_fields = render_fields(fields);
+    let field_names_joined_comma = render_field_names_joined_comma(fields);
+    let render_fields = render_fields(fields);
 
-        quote! {
-            #idx => {
-                #render_fields
+    quote! {
+        #discriminant => {
+            #render_fields
 
-                Ok(Self::#name {
-                    #field_names_joined_comma
-                })
-            }
+            Ok(Self::#name {
+                #field_names_joined_comma
+            })
         }
     }
 }

protocol-derive/src/render/encoder.rs

@@ -18,10 +18,7 @@ pub(crate) fn render_struct_encoder(name: &Ident, fields: &Vec<FieldData>) -> To
     }
 }
 
-pub(crate) fn render_struct_variant_encoder(
-    name: &Ident,
-    variants: &Vec<VariantData>,
-) -> TokenStream2 {
+pub(crate) fn render_enum_encoder(name: &Ident, variants: &Vec<VariantData>) -> TokenStream2 {
     let render_variants = render_variants(variants);
 
     quote! {
@@ -43,28 +40,39 @@ fn render_variants(variants: &Vec<VariantData>) -> TokenStream2 {
 }
 
 fn render_variant(variant: &VariantData) -> TokenStream2 {
-    let idx = variant.idx;
+    if variant.fields.is_empty() {
+        render_unit_variant(variant)
+    } else {
+        render_struct_variant(variant)
+    }
+}
+
+fn render_unit_variant(variant: &VariantData) -> TokenStream2 {
+    let discriminant = variant.discriminant;
+    let name = variant.name;
+
+    quote! {
+        Self::#name => {
+            writer.write_u8(#discriminant)?;
+        }
+    }
+}
+
+fn render_struct_variant(variant: &VariantData) -> TokenStream2 {
+    let discriminant = variant.discriminant;
     let name = variant.name;
     let fields = &variant.fields;
 
-    if fields.is_empty() {
-        quote! {
-            Self::#name => {
-                writer.write_u8(#idx)?;
-            }
-        }
-    } else {
-        let field_names_joined_comma = render_field_names_joined_comma(fields);
-        let render_fields = render_fields(fields, false);
+    let field_names_joined_comma = render_field_names_joined_comma(fields);
+    let render_fields = render_fields(fields, false);
 
-        quote! {
-            Self::#name {
-                #field_names_joined_comma
-            } => {
-                writer.write_u8(#idx)?;
+    quote! {
+        Self::#name {
+            #field_names_joined_comma
+        } => {
+            writer.write_u8(#discriminant)?;
 
-                #render_fields
-            }
+            #render_fields
         }
     }
 }