[derive] Support TryFromBytes for structs

TODO: - Should DataExt::fields have a different shape? Extracting field names for enums doesn't really make sense - Lots and lots of tests - If we manage to land derive(KnownLayout) first, replace manual impls of KnownLayout with the derive - Is there any way to make sure the code we emit doesn't emit warnings? I tested with an earlier version of the code that had lots of code smells like pointer coercions that emitted warnings, and even after emitting `#[deny(...)]` attributes, I couldn't get code which made use of the derive to complain at all - the compiler just happily accepted it. Makes progress on #5
google · Sep 19, 2023 · 3dd3531 · 3dd3531
1 parent 3008eb5
commit 3dd3531
Show file tree

Hide file tree

Showing 7 changed files with 199 additions and 47 deletions.
diff --git a/src/byteorder.rs b/src/byteorder.rs
@@ -176,16 +176,20 @@ example of how it can be used for parsing UDP packets.
 [`AsBytes`]: crate::AsBytes
 [`Unaligned`]: crate::Unaligned"),
             #[derive(Copy, Clone, Eq, PartialEq, Hash)]
-            #[cfg_attr(any(feature = "derive", test), derive(FromZeroes, FromBytes, AsBytes, Unaligned))]
+            #[cfg_attr(any(feature = "derive", test), derive(TryFromBytes, FromZeroes, FromBytes, AsBytes, Unaligned))]
             #[repr(transparent)]
             pub struct $name<O>([u8; $bytes], PhantomData<O>);
         }
 
+        impl_known_layout!(O => $name<O>);
+
         safety_comment! {
             /// SAFETY:
             /// `$name<O>` is `repr(transparent)`, and so it has the same layout
             /// as its only non-zero field, which is a `u8` array. `u8` arrays
-            /// are `FromZeroes`, `FromBytes`, `AsBytes`, and `Unaligned`.
+            /// are `TryFromBytes`, `FromZeroes`, `FromBytes`, `AsBytes`, and
+            /// `Unaligned`.
+            impl_or_verify!(O => TryFromBytes for $name<O>);
             impl_or_verify!(O => FromZeroes for $name<O>);
             impl_or_verify!(O => FromBytes for $name<O>);
             impl_or_verify!(O => AsBytes for $name<O>);

diff --git a/src/derive_util.rs b/src/derive_util.rs
@@ -66,6 +66,7 @@ macro_rules! union_has_padding {
 #[cfg(test)]
 mod tests {
     use crate::util::testutil::*;
+    use crate::*;
 
     #[test]
     fn test_struct_has_padding() {
@@ -124,4 +125,15 @@ mod tests {
         // anyway.
         test!(#[repr(C)] #[repr(packed)] {a: u8, b: u64} => true);
     }
+
+    #[test]
+    fn foo() {
+        #[derive(TryFromBytes)]
+        struct Foo {
+            f: u8,
+            b: bool,
+        }
+
+        impl_known_layout!(Foo);
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -1087,6 +1087,11 @@ pub unsafe trait FromBytes: FromZeroes {
 /// [github-repo]: https://github.com/google/zerocopy
 /// [`try_from_ref`]: TryFromBytes::try_from_ref
 pub unsafe trait TryFromBytes: KnownLayout {
+    #[doc(hidden)]
+    fn only_derive_is_allowed_to_implement_this_trait()
+    where
+        Self: Sized;
+
     /// Does a given memory range contain a valid instance of `Self`?
     ///
     /// # Safety
@@ -3754,10 +3759,16 @@ mod tests {
     //
     // This is used to test the custom derives of our traits. The `[u8]` type
     // gets a hand-rolled impl, so it doesn't exercise our custom derives.
-    #[derive(Debug, Eq, PartialEq, FromZeroes, FromBytes, AsBytes, Unaligned)]
+    #[derive(Debug, Eq, PartialEq, TryFromBytes, FromZeroes, FromBytes, AsBytes, Unaligned)]
     #[repr(transparent)]
     struct Unsized([u8]);
 
+    // SAFETY: `Unsized` is a `#[repr(transparent)]` wrapper around `[u8]`.
+    unsafe_impl_known_layout!(
+        #[repr([u8])]
+        Unsized
+    );
+
     impl Unsized {
         fn from_mut_slice(slc: &mut [u8]) -> &mut Unsized {
             // SAFETY: This *probably* sound - since the layouts of `[u8]` and
@@ -5122,14 +5133,16 @@ mod tests {
             assert_eq!(too_many_bytes[0], 123);
         }
 
-        #[derive(Debug, Eq, PartialEq, FromZeroes, FromBytes, AsBytes)]
+        #[derive(Debug, Eq, PartialEq, TryFromBytes, FromZeroes, FromBytes, AsBytes)]
         #[repr(C)]
         struct Foo {
             a: u32,
             b: Wrapping<u32>,
             c: Option<NonZeroU32>,
         }
 
+        impl_known_layout!(Foo);
+
         let expected_bytes: Vec<u8> = if cfg!(target_endian = "little") {
             vec![1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0]
         } else {
@@ -5151,12 +5164,14 @@ mod tests {
 
     #[test]
     fn test_array() {
-        #[derive(FromZeroes, FromBytes, AsBytes)]
+        #[derive(TryFromBytes, FromZeroes, FromBytes, AsBytes)]
         #[repr(C)]
         struct Foo {
             a: [u16; 33],
         }
 
+        impl_known_layout!(Foo);
+
         let foo = Foo { a: [0xFFFF; 33] };
         let expected = [0xFFu8; 66];
         assert_eq!(foo.as_bytes(), &expected[..]);
@@ -5215,23 +5230,25 @@ mod tests {
 
     #[test]
     fn test_transparent_packed_generic_struct() {
-        #[derive(AsBytes, FromZeroes, FromBytes, Unaligned)]
+        #[derive(AsBytes, TryFromBytes, FromZeroes, FromBytes, Unaligned)]
         #[repr(transparent)]
         struct Foo<T> {
             _t: T,
             _phantom: PhantomData<()>,
         }
 
+        impl_known_layout!(T => Foo<T>);
         assert_impl_all!(Foo<u32>: FromZeroes, FromBytes, AsBytes);
         assert_impl_all!(Foo<u8>: Unaligned);
 
-        #[derive(AsBytes, FromZeroes, FromBytes, Unaligned)]
+        #[derive(AsBytes, TryFromBytes, FromZeroes, FromBytes, Unaligned)]
         #[repr(packed)]
         struct Bar<T, U> {
             _t: T,
             _u: U,
         }
 
+        impl_known_layout!(T, U => Bar<T, U>);
         assert_impl_all!(Bar<u8, AU64>: FromZeroes, FromBytes, AsBytes, Unaligned);
     }
 
@@ -5495,7 +5512,7 @@ mod tests {
         assert_impls!(Wrapping<bool>: TryFromBytes, FromZeroes, AsBytes, Unaligned, !FromBytes);
         assert_impls!(Wrapping<NotZerocopy>: !TryFromBytes, !FromZeroes, !FromBytes, !AsBytes, !Unaligned);
 
-        assert_impls!(Unalign<u8>: FromZeroes, FromBytes, AsBytes, Unaligned, !TryFromBytes);
+        assert_impls!(Unalign<u8>: TryFromBytes, FromZeroes, FromBytes, AsBytes, Unaligned, TryFromBytes);
         assert_impls!(Unalign<NotZerocopy>: Unaligned, !TryFromBytes, !FromZeroes, !FromBytes, !AsBytes);
 
         assert_impls!([u8]: TryFromBytes, FromZeroes, FromBytes, AsBytes, Unaligned);

diff --git a/src/macros.rs b/src/macros.rs
@@ -130,6 +130,9 @@ macro_rules! unsafe_impl {
             let $candidate = unsafe { &*(candidate.as_ptr() as *const $repr) };
             $is_bit_valid
         }
+
+        #[allow(clippy::missing_inline_in_public_items)]
+        fn only_derive_is_allowed_to_implement_this_trait() {}
     };
     (@method TryFromBytes ; |$candidate:ident: NonNull<$repr:ty>| $is_bit_valid:expr) => {
         #[inline]
@@ -141,8 +144,15 @@ macro_rules! unsafe_impl {
             let $candidate = unsafe { NonNull::new_unchecked(candidate.as_ptr() as *mut $repr) };
             $is_bit_valid
         }
+
+        #[allow(clippy::missing_inline_in_public_items)]
+        fn only_derive_is_allowed_to_implement_this_trait() {}
+    };
+    (@method TryFromBytes) => {
+        #[inline(always)] unsafe fn is_bit_valid(_: NonNull<Self>) -> bool { true }
+        #[allow(clippy::missing_inline_in_public_items)]
+        fn only_derive_is_allowed_to_implement_this_trait() {}
     };
-    (@method TryFromBytes) => { #[inline(always)] unsafe fn is_bit_valid(_: NonNull<Self>) -> bool { true } };
     (@method $trait:ident) => {
         #[allow(clippy::missing_inline_in_public_items)]
         fn only_derive_is_allowed_to_implement_this_trait() {}
@@ -253,17 +263,17 @@ macro_rules! impl_known_layout {
     ($(const $constvar:ident : $constty:ty, $tyvar:ident $(: ?$optbound:ident)? => $ty:ty),* $(,)?) => {
         $(impl_known_layout!(@inner const $constvar: $constty, $tyvar $(: ?$optbound)? => $ty);)*
     };
-    ($($tyvar:ident $(: ?$optbound:ident)? => $ty:ty),* $(,)?) => {
-        $(impl_known_layout!(@inner , $tyvar $(: ?$optbound)? => $ty);)*
+    ($($($tyvar:ident $(: ?$optbound:ident)?),* => $ty:ty),* $(,)?) => {
+        $(impl_known_layout!(@inner , $($tyvar $(: ?$optbound)?),* => $ty);)*
     };
     ($($ty:ty),*) => { $(impl_known_layout!(@inner , => $ty);)* };
-    (@inner $(const $constvar:ident : $constty:ty)? , $($tyvar:ident $(: ?$optbound:ident)?)? => $ty:ty) => {
+    (@inner $(const $constvar:ident : $constty:ty)? , $($tyvar:ident $(: ?$optbound:ident)?),* => $ty:ty) => {
         const _: () = {
             use core::ptr::NonNull;
 
-            impl<$(const $constvar : $constty,)? $($tyvar $(: ?$optbound)?)?> sealed::KnownLayoutSealed for $ty {}
+            impl<$(const $constvar : $constty,)? $($tyvar $(: ?$optbound)?),*> sealed::KnownLayoutSealed for $ty {}
             // SAFETY: Delegates safety to `DstLayout::for_type`.
-            unsafe impl<$(const $constvar : $constty,)? $($tyvar $(: ?$optbound)?)?> KnownLayout for $ty {
+            unsafe impl<$(const $constvar : $constty,)? $($tyvar $(: ?$optbound)?),*> KnownLayout for $ty {
                 const LAYOUT: DstLayout = DstLayout::for_type::<$ty>();
 
                 // SAFETY: `.cast` preserves address and provenance.

diff --git a/src/wrappers.rs b/src/wrappers.rs
@@ -54,10 +54,15 @@ use super::*;
 // [3] https://github.com/google/zerocopy/issues/209
 #[allow(missing_debug_implementations)]
 #[derive(Default, Copy)]
-#[cfg_attr(any(feature = "derive", test), derive(FromZeroes, FromBytes, AsBytes, Unaligned))]
+#[cfg_attr(
+    any(feature = "derive", test),
+    derive(TryFromBytes, FromZeroes, FromBytes, AsBytes, Unaligned)
+)]
 #[repr(C, packed)]
 pub struct Unalign<T>(T);
 
+impl_known_layout!(T => Unalign<T>);
+
 safety_comment! {
     /// SAFETY:
     /// - `Unalign<T>` is `repr(packed)`, so it is unaligned regardless of the

diff --git a/zerocopy-derive/src/ext.rs b/zerocopy-derive/src/ext.rs
@@ -2,48 +2,74 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-use syn::{Data, DataEnum, DataStruct, DataUnion, Type};
+use proc_macro2::TokenStream;
+use quote::ToTokens;
+use syn::{Data, DataEnum, DataStruct, DataUnion, Field, Index, Type};
 
 pub trait DataExt {
-    /// Extract the types of all fields. For enums, extract the types of fields
-    /// from each variant.
-    fn field_types(&self) -> Vec<&Type>;
+    /// Extract the names and types of all fields. For enums, extract the names
+    /// and types of fields from each variant. For tuple structs, the names are
+    /// the indices used to index into the struct (ie, `0`, `1`, etc).
+    ///
+    /// TODO: Extracting field names for enums doesn't really make sense. Types
+    /// makes sense because we don't care about where they live - we just care
+    /// about transitive ownership. But for field names, we'd only use them when
+    /// generating is_bit_valid, which cares about where they live.
+    fn fields(&self) -> Vec<(TokenStream, &Type)>;
 }
 
 impl DataExt for Data {
-    fn field_types(&self) -> Vec<&Type> {
+    fn fields(&self) -> Vec<(TokenStream, &Type)> {
         match self {
-            Data::Struct(strc) => strc.field_types(),
-            Data::Enum(enm) => enm.field_types(),
-            Data::Union(un) => un.field_types(),
+            Data::Struct(strc) => strc.fields(),
+            Data::Enum(enm) => enm.fields(),
+            Data::Union(un) => un.fields(),
         }
     }
 }
 
 impl DataExt for DataStruct {
-    fn field_types(&self) -> Vec<&Type> {
-        self.fields.iter().map(|f| &f.ty).collect()
+    fn fields(&self) -> Vec<(TokenStream, &Type)> {
+        map_fields(&self.fields)
     }
 }
 
 impl DataExt for DataEnum {
-    fn field_types(&self) -> Vec<&Type> {
-        self.variants.iter().flat_map(|var| &var.fields).map(|f| &f.ty).collect()
+    fn fields(&self) -> Vec<(TokenStream, &Type)> {
+        map_fields(self.variants.iter().flat_map(|var| &var.fields))
     }
 }
 
 impl DataExt for DataUnion {
-    fn field_types(&self) -> Vec<&Type> {
-        self.fields.named.iter().map(|f| &f.ty).collect()
+    fn fields(&self) -> Vec<(TokenStream, &Type)> {
+        map_fields(&self.fields.named)
     }
 }
 
+fn map_fields<'a>(
+    fields: impl 'a + IntoIterator<Item = &'a Field>,
+) -> Vec<(TokenStream, &'a Type)> {
+    fields
+        .into_iter()
+        .enumerate()
+        .map(|(idx, f)| {
+            (
+                f.ident
+                    .as_ref()
+                    .map(ToTokens::to_token_stream)
+                    .unwrap_or_else(|| Index::from(idx).to_token_stream()),
+                &f.ty,
+            )
+        })
+        .collect()
+}
+
 pub trait EnumExt {
     fn is_c_like(&self) -> bool;
 }
 
 impl EnumExt for DataEnum {
     fn is_c_like(&self) -> bool {
-        self.field_types().is_empty()
+        self.fields().is_empty()
     }
 }