Do not produce useless match cases unit enums

Given enums like these: ```rust enum Foo { A(u16), B, C(i32, i32), } enum Bar { A, B, C, } ``` serialise derive produces something like these: ```rust // Recursive expansion of borsh::BorshSerialize macro // =================================================== impl borsh::ser::BorshSerialize for Foo { fn serialize<W: borsh::io::Write>( &self, writer: &mut W, ) -> ::core::result::Result<(), borsh::io::Error> { let variant_idx: u8 = match self { Foo::A(..) => 0u8, Foo::B => 1u8, Foo::C(..) => 2u8, }; writer.write_all(&variant_idx.to_le_bytes())?; match self { Foo::A(id0) => { borsh::BorshSerialize::serialize(id0, writer)?; } Foo::B => {} Foo::C(id0, id1) => { borsh::BorshSerialize::serialize(id0, writer)?; borsh::BorshSerialize::serialize(id1, writer)?; } } Ok(()) } } // Recursive expansion of borsh::BorshSerialize macro // =================================================== impl borsh::ser::BorshSerialize for Bar { fn serialize<W: borsh::io::Write>( &self, writer: &mut W, ) -> ::core::result::Result<(), borsh::io::Error> { let variant_idx: u8 = match self { Bar::A => 0u8, Bar::B => 1u8, Bar::C => 2u8, }; writer.write_all(&variant_idx.to_le_bytes())?; match self { Bar::A => {} Bar::B => {} Bar::C => {} } Ok(()) } } ``` Notably in `bar` case, the whole `match self` is useless because there's nothing left to serialise. With this patch, the derives now look like this: ```rust // Recursive expansion of borsh::BorshSerialize macro // =================================================== impl borsh::ser::BorshSerialize for Foo { fn serialize<W: borsh::io::Write>( &self, writer: &mut W, ) -> ::core::result::Result<(), borsh::io::Error> { let variant_idx: u8 = match self { Foo::A(..) => 0u8, Foo::B => 1u8, Foo::C(..) => 2u8, }; writer.write_all(&variant_idx.to_le_bytes())?; match self { Foo::A(id0) => { borsh::BorshSerialize::serialize(id0, writer)?; } Foo::C(id0, id1) => { borsh::BorshSerialize::serialize(id0, writer)?; borsh::BorshSerialize::serialize(id1, writer)?; } _ => {} } Ok(()) } } // Recursive expansion of borsh::BorshSerialize macro // =================================================== impl borsh::ser::BorshSerialize for Bar { fn serialize<W: borsh::io::Write>( &self, writer: &mut W, ) -> ::core::result::Result<(), borsh::io::Error> { let variant_idx: u8 = match self { Bar::A => 0u8, Bar::B => 1u8, Bar::C => 2u8, }; writer.write_all(&variant_idx.to_le_bytes())?; Ok(()) } } ``` Notably, the whole `match self` is gone for `Bar`. For `Foo`, any unit field cases are now inside a `_ => {}` catch-all. What's the point? Well, it's just nice to produce less for compiler to deal with, reducing upstream build times. Further, it makes it much nicer for anyone reading/copying the expanded macros. However patch this does add some amount of code complexity so it's up to the maintainers to decide if it's worth taking.
near · Dec 6, 2023 · ed840a1 · ed840a1
1 parent 903601b
commit ed840a1
Showing 1 changed file with 79 additions and 39 deletions.
diff --git a/borsh-derive/src/internals/serialize/enums/mod.rs b/borsh-derive/src/internals/serialize/enums/mod.rs
@@ -18,6 +18,7 @@ pub fn process(input: &ItemEnum, cratename: Path) -> syn::Result<TokenStream2> {
     let mut fields_body = TokenStream2::new();
     let use_discriminant = item::contains_use_discriminant(input)?;
     let discriminants = Discriminants::new(&input.variants);
+    let mut blank_variants = false;
 
     for (variant_idx, variant) in input.variants.iter().enumerate() {
         let variant_ident = &variant.ident;
@@ -30,15 +31,42 @@ pub fn process(input: &ItemEnum, cratename: Path) -> syn::Result<TokenStream2> {
             &mut generics_output,
         )?;
         all_variants_idx_body.extend(variant_output.variant_idx_body);
-        let (variant_header, variant_body) = (variant_output.header, variant_output.body);
-        fields_body.extend(quote!(
-            #enum_ident::#variant_ident #variant_header => {
-                #variant_body
-            }
-        ))
+        match variant_output.body {
+            VariantBody::Blank => blank_variants = true,
+            VariantBody::Fields(VariantFields { header, body }) => fields_body.extend(quote!(
+                #enum_ident::#variant_ident #header => {
+                    #body
+                }
+            )),
+        }
     }
     generics_output.extend(&mut where_clause, &cratename);
 
+    let fields_match = if fields_body.is_empty() {
+        // If we no variants with fields, there's nothing to match against. Just
+        // re-use the empty token stream.
+        fields_body
+    } else {
+        let unit_fields_catchall = if blank_variants {
+            // We had some variants with unit fields, create a catch-all for
+            // these to be used at the bottom.
+            quote!(
+                _ => {}
+            )
+        } else {
+            TokenStream2::new()
+        };
+        // Create a match that serialises all the fields for each non-unit
+        // variant and add a catch-all at the bottom if we do have unit
+        // variants.
+        quote!(
+            match self {
+                #fields_body
+                #unit_fields_catchall
+            }
+        )
+    };
+
     Ok(quote! {
         impl #impl_generics #cratename::ser::BorshSerialize for #enum_ident #ty_generics #where_clause {
             fn serialize<W: #cratename::io::Write>(&self, writer: &mut W) -> ::core::result::Result<(), #cratename::io::Error> {
@@ -47,29 +75,29 @@ pub fn process(input: &ItemEnum, cratename: Path) -> syn::Result<TokenStream2> {
                 };
                 writer.write_all(&variant_idx.to_le_bytes())?;
 
-                match self {
-                    #fields_body
-                }
+                #fields_match
                 Ok(())
             }
         }
     })
 }
 
-struct VariantOutput {
+#[derive(Default)]
+struct VariantFields {
     header: TokenStream2,
     body: TokenStream2,
-    variant_idx_body: TokenStream2,
 }
 
-impl VariantOutput {
-    fn new() -> Self {
-        Self {
-            body: TokenStream2::new(),
-            header: TokenStream2::new(),
-            variant_idx_body: TokenStream2::new(),
-        }
-    }
+enum VariantBody {
+    // No body variant, unit enum variant.
+    Blank,
+    // Variant with body (fields)
+    Fields(VariantFields),
+}
+
+struct VariantOutput {
+    body: VariantBody,
+    variant_idx_body: TokenStream2,
 }
 
 fn process_variant(
@@ -80,36 +108,48 @@ fn process_variant(
     generics: &mut serialize::GenericsOutput,
 ) -> syn::Result<VariantOutput> {
     let variant_ident = &variant.ident;
-    let mut variant_output = VariantOutput::new();
-    match &variant.fields {
+    let variant_output = match &variant.fields {
         Fields::Named(fields) => {
+            let mut variant_fields = VariantFields::default();
             for field in &fields.named {
                 let field_id = serialize::FieldId::Enum(field.ident.clone().unwrap());
-                process_field(field, field_id, cratename, generics, &mut variant_output)?;
+                process_field(field, field_id, cratename, generics, &mut variant_fields)?;
+            }
+            let header = variant_fields.header;
+            VariantOutput {
+                body: VariantBody::Fields(VariantFields {
+                    // `..` pattern matching works even if all fields were specified
+                    header: quote! { { #header.. }},
+                    body: variant_fields.body,
+                }),
+                variant_idx_body: quote!(
+                    #enum_ident::#variant_ident {..} => #discriminant_value,
+                ),
             }
-            let header = variant_output.header;
-            // `..` pattern matching works even if all fields were specified
-            variant_output.header = quote! { { #header.. }};
-            variant_output.variant_idx_body = quote!(
-                #enum_ident::#variant_ident {..} => #discriminant_value,
-            );
         }
         Fields::Unnamed(fields) => {
+            let mut variant_fields = VariantFields::default();
             for (field_idx, field) in fields.unnamed.iter().enumerate() {
                 let field_id = serialize::FieldId::new_enum_unnamed(field_idx)?;
-                process_field(field, field_id, cratename, generics, &mut variant_output)?;
+                process_field(field, field_id, cratename, generics, &mut variant_fields)?;
+            }
+            let header = variant_fields.header;
+            VariantOutput {
+                body: VariantBody::Fields(VariantFields {
+                    header: quote! { ( #header )},
+                    body: variant_fields.body,
+                }),
+                variant_idx_body: quote!(
+                    #enum_ident::#variant_ident(..) => #discriminant_value,
+                ),
             }
-            let header = variant_output.header;
-            variant_output.header = quote! { ( #header )};
-            variant_output.variant_idx_body = quote!(
-                #enum_ident::#variant_ident(..) => #discriminant_value,
-            );
         }
-        Fields::Unit => {
-            variant_output.variant_idx_body = quote!(
+        Fields::Unit => VariantOutput {
+            body: VariantBody::Blank,
+            variant_idx_body: quote!(
                 #enum_ident::#variant_ident => #discriminant_value,
-            );
-        }
+            ),
+        },
     };
     Ok(variant_output)
 }
@@ -119,7 +159,7 @@ fn process_field(
     field_id: serialize::FieldId,
     cratename: &Path,
     generics: &mut serialize::GenericsOutput,
-    output: &mut VariantOutput,
+    output: &mut VariantFields,
 ) -> syn::Result<()> {
     let parsed = field::Attributes::parse(&field.attrs)?;