[derive] Correctly handle #[repr(align)] enums (google#1784)

* [derive] Don't support IntoBytes on repr(Rust) types

Closes google#1764

* [derive] Correctly handle #[repr(align)] enums

Closes google#1758

Author: joshlf

zerocopy-derive/src/enum.rs

@@ -13,8 +13,8 @@ use syn::{parse_quote, DataEnum, Error, Fields, Generics, Ident};
 use crate::{derive_try_from_bytes_inner, repr::EnumRepr, Trait};
 
 /// Generates a tag enum for the given enum. This generates an enum with the
-/// same `repr`s, variants, and corresponding discriminants, but none of the
-/// fields.
+/// same non-align `repr`s, variants, and corresponding discriminants, but none
+/// of the fields.
 pub(crate) fn generate_tag_enum(repr: &EnumRepr, data: &DataEnum) -> TokenStream {
     let variants = data.variants.iter().map(|v| {
         let ident = &v.ident;
@@ -25,6 +25,14 @@ pub(crate) fn generate_tag_enum(repr: &EnumRepr, data: &DataEnum) -> TokenStream
         }
     });
 
+    // Don't include any `repr(align)` when generating the tag enum, as that
+    // could add padding after the tag but before any variants, which is not the
+    // correct behavior.
+    let repr = match repr {
+        EnumRepr::Transparent(span) => quote::quote_spanned! { *span => #[repr(transparent)] },
+        EnumRepr::Compound(c, _) => quote! { #c },
+    };
+
     quote! {
         #repr
         #[allow(dead_code)]

zerocopy-derive/tests/enum_to_bytes.rs

@@ -121,3 +121,13 @@ enum HasData32 {
 }
 
 util_assert_impl_all!(HasData: imp::IntoBytes);
+
+// After #1752 landed but before #1758 was fixed, this failed to compile because
+// the padding check treated the tag type as being `#[repr(u8, align(2))] struct
+// Tag { A }`, which is two bytes long, rather than the correct `#[repr(u8)]
+// struct Tag { A }`, which is one byte long.
+#[derive(imp::IntoBytes)]
+#[repr(u8, align(2))]
+enum BadTagWouldHavePadding {
+    A(u8, u16),
+}

zerocopy-derive/tests/enum_try_from_bytes.rs

@@ -178,6 +178,8 @@ fn test_weird_discriminants() {
     );
 }
 
+// Technically non-portable since this is only `IntoBytes` if the discriminant
+// is an `i32` or `u32`, but we'll cross that bridge when we get to it...
 #[derive(
     Eq, PartialEq, Debug, imp::KnownLayout, imp::Immutable, imp::TryFromBytes, imp::IntoBytes,
 )]
@@ -205,6 +207,43 @@ fn test_has_fields() {
     );
 }
 
+#[derive(Eq, PartialEq, Debug, imp::KnownLayout, imp::Immutable, imp::TryFromBytes)]
+#[repr(C, align(16))]
+enum HasFieldsAligned {
+    A(u32),
+    B { foo: ::core::num::NonZeroU32 },
+}
+
+util_assert_impl_all!(HasFieldsAligned: imp::TryFromBytes);
+
+#[test]
+fn test_has_fields_aligned() {
+    const SIZE: usize = ::core::mem::size_of::<HasFieldsAligned>();
+
+    #[derive(imp::IntoBytes)]
+    #[repr(C)]
+    struct BytesOfHasFieldsAligned {
+        has_fields: HasFields,
+        padding: [u8; 8],
+    }
+
+    let wrap = |has_fields| BytesOfHasFieldsAligned { has_fields, padding: [0; 8] };
+
+    let bytes: [u8; SIZE] = ::zerocopy::transmute!(wrap(HasFields::A(10)));
+    imp::assert_eq!(
+        <HasFieldsAligned as imp::TryFromBytes>::try_read_from_bytes(&bytes[..]),
+        imp::Ok(HasFieldsAligned::A(10)),
+    );
+
+    let bytes: [u8; SIZE] = ::zerocopy::transmute!(wrap(HasFields::B {
+        foo: ::core::num::NonZeroU32::new(123456).unwrap()
+    }));
+    imp::assert_eq!(
+        <HasFieldsAligned as imp::TryFromBytes>::try_read_from_bytes(&bytes[..]),
+        imp::Ok(HasFieldsAligned::B { foo: ::core::num::NonZeroU32::new(123456).unwrap() }),
+    );
+}
+
 #[derive(
     Eq, PartialEq, Debug, imp::KnownLayout, imp::Immutable, imp::TryFromBytes, imp::IntoBytes,
 )]
@@ -233,6 +272,45 @@ fn test_has_fields_primitive() {
     );
 }
 
+#[derive(Eq, PartialEq, Debug, imp::KnownLayout, imp::Immutable, imp::TryFromBytes)]
+#[repr(u32, align(16))]
+enum HasFieldsPrimitiveAligned {
+    A(u32),
+    B { foo: ::core::num::NonZeroU32 },
+}
+
+util_assert_impl_all!(HasFieldsPrimitiveAligned: imp::TryFromBytes);
+
+#[test]
+fn test_has_fields_primitive_aligned() {
+    const SIZE: usize = ::core::mem::size_of::<HasFieldsPrimitiveAligned>();
+
+    #[derive(imp::IntoBytes)]
+    #[repr(C)]
+    struct BytesOfHasFieldsPrimitiveAligned {
+        has_fields: HasFieldsPrimitive,
+        padding: [u8; 8],
+    }
+
+    let wrap = |has_fields| BytesOfHasFieldsPrimitiveAligned { has_fields, padding: [0; 8] };
+
+    let bytes: [u8; SIZE] = ::zerocopy::transmute!(wrap(HasFieldsPrimitive::A(10)));
+    imp::assert_eq!(
+        <HasFieldsPrimitiveAligned as imp::TryFromBytes>::try_read_from_bytes(&bytes[..]),
+        imp::Ok(HasFieldsPrimitiveAligned::A(10)),
+    );
+
+    let bytes: [u8; SIZE] = ::zerocopy::transmute!(wrap(HasFieldsPrimitive::B {
+        foo: ::core::num::NonZeroU32::new(123456).unwrap()
+    }));
+    imp::assert_eq!(
+        <HasFieldsPrimitiveAligned as imp::TryFromBytes>::try_read_from_bytes(&bytes[..]),
+        imp::Ok(HasFieldsPrimitiveAligned::B {
+            foo: ::core::num::NonZeroU32::new(123456).unwrap()
+        }),
+    );
+}
+
 #[derive(imp::TryFromBytes)]
 #[repr(align(4), u32)]
 enum HasReprAlignFirst {