KnownLayout tracks pointer metadata type

src/lib.rs

@@ -1065,14 +1065,73 @@ pub unsafe trait KnownLayout {
     where
         Self: Sized;
 
+    /// `()` for sized types and `usize` for slice DSTs.
+    #[doc(hidden)]
+    type PointerMetadata: PointerMetadata;
+
     #[doc(hidden)]
     const LAYOUT: DstLayout;
 
     /// SAFETY: The returned pointer has the same address and provenance as
     /// `bytes`. If `Self` is a DST, the returned pointer's referent has `elems`
-    /// elements in its trailing slice. If `Self` is sized, `elems` is ignored.
+    /// elements in its trailing slice.
     #[doc(hidden)]
-    fn raw_from_ptr_len(bytes: NonNull<u8>, elems: usize) -> NonNull<Self>;
+    fn raw_from_ptr_len(bytes: NonNull<u8>, meta: Self::PointerMetadata) -> NonNull<Self>;
+}
+
+/// The metadata associated with a [`KnownLayout`] type.
+#[doc(hidden)]
+pub trait PointerMetadata {
+    /// Constructs a `Self` from an element count.
+    ///
+    /// If `Self = ()`, this returns `()`. If `Self = usize`, this returns
+    /// `elems`. No other types are currently supported.
+    fn from_elem_count(elems: usize) -> Self;
+
+    /// Computes the size of the object with the given layout and pointer
+    /// metadata.
+    ///
+    /// # Panics
+    ///
+    /// If `Self = ()`, `layout` must describe a sized type. If `Self = usize`,
+    /// `layout` must describe a slice DST. Otherwise, `size_for_metadata` may
+    /// panic.
+    fn size_for_metadata(&self, layout: DstLayout) -> Option<usize>;
+}
+
+impl PointerMetadata for () {
+    #[inline]
+    #[allow(clippy::unused_unit)]
+    fn from_elem_count(_elems: usize) -> () {}
+
+    #[inline]
+    fn size_for_metadata(&self, layout: DstLayout) -> Option<usize> {
+        match layout.size_info {
+            SizeInfo::Sized { size } => Some(size),
+            // NOTE: This branch is unreachable, but we return `None` rather
+            // than `unreachable!()` to avoid generating panic paths.
+            SizeInfo::SliceDst(_) => None,
+        }
+    }
+}
+
+impl PointerMetadata for usize {
+    #[inline]
+    fn from_elem_count(elems: usize) -> usize {
+        elems
+    }
+
+    #[inline]
+    fn size_for_metadata(&self, layout: DstLayout) -> Option<usize> {
+        match layout.size_info {
+            SizeInfo::SliceDst(TrailingSliceLayout { offset, elem_size }) => {
+                let slice_len = elem_size.checked_mul(*self)?;
+                let without_padding = offset.checked_add(slice_len)?;
+                without_padding.checked_add(util::padding_needed_for(without_padding, layout.align))
+            }
+            SizeInfo::Sized { .. } => unreachable!(),
+        }
+    }
 }
 
 // SAFETY: Delegates safety to `DstLayout::for_slice`.
@@ -1083,6 +1142,9 @@ unsafe impl<T> KnownLayout for [T] {
         Self: Sized,
     {
     }
+
+    type PointerMetadata = usize;
+
     const LAYOUT: DstLayout = DstLayout::for_slice::<T>();
 
     // SAFETY: `.cast` preserves address and provenance. The returned pointer

src/macros.rs

@@ -542,6 +542,8 @@ macro_rules! impl_known_layout {
                 #[allow(clippy::missing_inline_in_public_items)]
                 fn only_derive_is_allowed_to_implement_this_trait() where Self: Sized {}
 
+                type PointerMetadata = ();
+
                 #[allow(unused_qualifications)]
                 const LAYOUT: crate::DstLayout = crate::DstLayout::for_type::<$ty>();
 
@@ -550,7 +552,7 @@ macro_rules! impl_known_layout {
                 // TODO(#429): Add documentation to `.cast` that promises that
                 // it preserves provenance.
                 #[inline(always)]
-                fn raw_from_ptr_len(bytes: NonNull<u8>, _elems: usize) -> NonNull<Self> {
+                fn raw_from_ptr_len(bytes: NonNull<u8>, _meta: ()) -> NonNull<Self> {
                     bytes.cast::<Self>()
                 }
             }
@@ -579,6 +581,8 @@ macro_rules! unsafe_impl_known_layout {
                 #[allow(clippy::missing_inline_in_public_items)]
                 fn only_derive_is_allowed_to_implement_this_trait() {}
 
+                type PointerMetadata = <$repr as KnownLayout>::PointerMetadata;
+
                 const LAYOUT: DstLayout = <$repr as KnownLayout>::LAYOUT;
 
                 // SAFETY: All operations preserve address and provenance.
@@ -588,9 +592,9 @@ macro_rules! unsafe_impl_known_layout {
                 // that it preserves provenance.
                 #[inline(always)]
                 #[allow(unused_qualifications)] // for `core::ptr::NonNull`
-                fn raw_from_ptr_len(bytes: NonNull<u8>, elems: usize) -> NonNull<Self> {
+                fn raw_from_ptr_len(bytes: NonNull<u8>, meta: <$repr as KnownLayout>::PointerMetadata) -> NonNull<Self> {
                     #[allow(clippy::as_conversions)]
-                    let ptr = <$repr>::raw_from_ptr_len(bytes, elems).as_ptr() as *mut Self;
+                    let ptr = <$repr>::raw_from_ptr_len(bytes, meta).as_ptr() as *mut Self;
                     // SAFETY: `ptr` was converted from `bytes`, which is non-null.
                     unsafe { NonNull::new_unchecked(ptr) }
                 }

src/pointer/ptr.rs

@@ -922,6 +922,7 @@ mod _transitions {
 /// Casts of the referent type.
 mod _casts {
     use super::*;
+    use crate::PointerMetadata;
 
     impl<'a, T, I> Ptr<'a, T, I>
     where
@@ -1178,6 +1179,7 @@ mod _casts {
             // produces a pointer whose address is greater than or equal to that of
             // `ptr`. Since `ptr` is a `NonNull`, `base` is also non-null.
             let base = unsafe { NonNull::new_unchecked(base) };
+            let elems = <U as KnownLayout>::PointerMetadata::from_elem_count(elems);
             let ptr = U::raw_from_ptr_len(base, elems);
 
             // SAFETY:

zerocopy-derive/src/lib.rs

@@ -161,6 +161,8 @@ fn derive_known_layout_inner(ast: &DeriveInput) -> proc_macro2::TokenStream {
         (
             SelfBounds::None,
             quote!(
+                type PointerMetadata = <#trailing_field_ty as ::zerocopy::KnownLayout>::PointerMetadata;
+
                 // SAFETY: `LAYOUT` accurately describes the layout of `Self`.
                 // The layout of `Self` is reflected using a sequence of
                 // invocations of `DstLayout::{new_zst,extend,pad_to_align}`.
@@ -195,10 +197,10 @@ fn derive_known_layout_inner(ast: &DeriveInput) -> proc_macro2::TokenStream {
                 #[inline(always)]
                 fn raw_from_ptr_len(
                     bytes: ::zerocopy::macro_util::core_reexport::ptr::NonNull<u8>,
-                    elems: usize,
+                    meta: <#trailing_field_ty as ::zerocopy::KnownLayout>::PointerMetadata,
                 ) -> ::zerocopy::macro_util::core_reexport::ptr::NonNull<Self> {
                     use ::zerocopy::{KnownLayout};
-                    let trailing = <#trailing_field_ty as KnownLayout>::raw_from_ptr_len(bytes, elems);
+                    let trailing = <#trailing_field_ty as KnownLayout>::raw_from_ptr_len(bytes, meta);
                     let slf = trailing.as_ptr() as *mut Self;
                     // SAFETY: Constructed from `trailing`, which is non-null.
                     unsafe { ::zerocopy::macro_util::core_reexport::ptr::NonNull::new_unchecked(slf) }
@@ -212,6 +214,8 @@ fn derive_known_layout_inner(ast: &DeriveInput) -> proc_macro2::TokenStream {
         (
             SelfBounds::SIZED,
             quote!(
+                type PointerMetadata = ();
+
                 // SAFETY: `LAYOUT` is guaranteed to accurately describe the
                 // layout of `Self`, because that is the documented safety
                 // contract of `DstLayout::for_type`.
@@ -224,7 +228,7 @@ fn derive_known_layout_inner(ast: &DeriveInput) -> proc_macro2::TokenStream {
                 #[inline(always)]
                 fn raw_from_ptr_len(
                     bytes: ::zerocopy::macro_util::core_reexport::ptr::NonNull<u8>,
-                    _elems: usize,
+                    _meta: (),
                 ) -> ::zerocopy::macro_util::core_reexport::ptr::NonNull<Self> {
                     bytes.cast::<Self>()
                 }