Generalize split_at_unchecked

jswrenn · jswrenn · commit c83d465b838c · 2025-03-10T19:05:15.000Z
Makes progress towards #1290.
diff --git a/src/lib.rs b/src/lib.rs
@@ -749,6 +749,11 @@ pub unsafe trait KnownLayout {
     #[doc(hidden)]
     type MaybeUninit: ?Sized + KnownLayout<PointerMetadata = Self::PointerMetadata>;
 
+    /// The type of the trailing element slice.
+    ///
+    /// If `Self` is a DST, then `TrailingElem` is `0`.
+    type TrailingElem: Sized;
+
     /// The layout of `Self`.
     ///
     /// # Safety
@@ -830,6 +835,22 @@ pub trait PointerMetadata: Copy + Eq + Debug {
     /// `size_for_metadata` promises to only return `None` if the resulting size
     /// would not fit in a `usize`.
     fn size_for_metadata(&self, layout: DstLayout) -> Option<usize>;
+
+    /// Computes the trailing padding bytes 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, `padding_needed_for` may
+    /// panic.
+    ///
+    /// # Safety
+    ///
+    /// `padding_needed_for` promises to only return `None` if the resulting
+    /// padding would not fit in a `usize`. (TODO: Can we prove this is
+    /// impossible?)
+    fn padding_needed_for(&self, layout: DstLayout) -> Option<usize>;
 }
 
 impl PointerMetadata for () {
@@ -846,6 +867,14 @@ impl PointerMetadata for () {
             SizeInfo::SliceDst(_) => None,
         }
     }
+
+    #[inline]
+    fn padding_needed_for(&self, layout: DstLayout) -> Option<usize> {
+        match layout.size_info {
+            SizeInfo::Sized { .. } => Some(0),
+            SizeInfo::SliceDst(..) => None,
+        }
+    }
 }
 
 impl PointerMetadata for usize {
@@ -867,6 +896,18 @@ impl PointerMetadata for usize {
             SizeInfo::Sized { .. } => None,
         }
     }
+
+    #[inline]
+    fn padding_needed_for(&self, layout: DstLayout) -> Option<usize> {
+        match layout.size_info {
+            SizeInfo::Sized { .. } => Some(0),
+            SizeInfo::SliceDst(TrailingSliceLayout { offset, elem_size }) => {
+                let trailing_size = self.checked_mul(elem_size)?;
+                let unpadded_size = trailing_size.checked_add(offset)?;
+                Some(util::padding_needed_for(unpadded_size, layout.align))
+            }
+        }
+    }
 }
 
 // SAFETY: Delegates safety to `DstLayout::for_slice`.
@@ -913,6 +954,8 @@ unsafe impl<T> KnownLayout for [T] {
     //   size_of::<T>()` bytes.
     type MaybeUninit = [CoreMaybeUninit<T>];
 
+    type TrailingElem = T;
+
     const LAYOUT: DstLayout = DstLayout::for_slice::<T>();
 
     // SAFETY: `.cast` preserves address and provenance. The returned pointer
diff --git a/src/pointer/inner.rs b/src/pointer/inner.rs
@@ -139,6 +139,159 @@ impl<'a, T: ?Sized> PtrInner<'a, T> {
     }
 }
 
+#[allow(clippy::needless_lifetimes)]
+impl<'a, T> PtrInner<'a, T>
+where
+    T: ?Sized + KnownLayout,
+{
+    /// Produces the metadata of this `ptr`.
+    ///
+    /// # Safety
+    ///
+    /// Unsafe code my rely on `len` satisfying the above contract.
+    pub(crate) fn meta(self) -> T::PointerMetadata {
+        T::pointer_to_metadata(self.as_non_null().as_ptr())
+    }
+
+    /// Produces a `PtrInner` with the same address and provenance as `self` but
+    /// the given `meta`.
+    ///
+    /// # Safety
+    ///
+    /// The caller promises that if `self`'s referent is not zero sized, then
+    /// a pointer constructed from its address with the given `meta` metadata
+    /// will address a subset of the allocation pointed to by `self`.
+    #[inline]
+    pub(crate) unsafe fn with_meta(self, meta: T::PointerMetadata) -> Self
+    where
+        T: KnownLayout,
+    {
+        let raw = T::raw_from_ptr_len(self.as_non_null().cast(), meta);
+
+        // SAFETY:
+        //
+        // Lemma 0: `raw` either addresses zero bytes, or addresses a subset of
+        //          the allocation pointed to by `self` and has the same
+        //          provenance as `self`. Proof: `raw` is constructed using
+        //          provenance-preserving operations, and the caller has
+        //          promised that, if `self`'s referent is not zero-sized, the
+        //          resulting pointer addresses a subset of the allocation
+        //          pointed to by `self`.
+        //
+        // 0. Per Lemma 0 and by invariant on `self`, if `ptr`'s referent is not
+        //    zero sized, then `ptr` is derived from some valid Rust allocation,
+        //    `A`.
+        // 1. Per Lemma 0 and by invariant on `self`, if `ptr`'s referent is not
+        //    zero sized, then `ptr` has valid provenance for `A`.
+        // 2. Per Lemma 0 and by invariant on `self`, if `ptr`'s referent is not
+        //    zero sized, then `ptr` addresses a byte range which is entirely
+        //    contained in `A`.
+        // 3. Per Lemma 0 and by invariant on `self`, `ptr` addresses a byte
+        //    range whose length fits in an `isize`.
+        // 4. Per Lemma 0 and by invariant on `self`, `ptr` addresses a byte
+        //    range which does not wrap around the address space.
+        // 5. Per Lemma 0 and by invariant on `self`, if `ptr`'s referent is not
+        //    zero sized, then `A` is guaranteed to live for at least `'a`.
+        unsafe { PtrInner::new(raw) }
+    }
+}
+
+#[allow(clippy::needless_lifetimes)]
+impl<'a, T> PtrInner<'a, T>
+where
+    T: ?Sized + KnownLayout<PointerMetadata = usize>,
+{
+    /// The number of trailing slice elements in the object referenced by
+    /// `self`.
+    ///
+    /// # Safety
+    ///
+    /// Unsafe code my rely on `len` satisfying the above contract.
+    pub(crate) fn len(&self) -> usize {
+        self.meta()
+    }
+
+    /// Splits `T` into two.
+    ///
+    /// # Safety
+    ///
+    /// The caller promises that:
+    ///  - `l_len <= self.len()`.
+    ///
+    /// Given `let (left, right) = ptr.split_at(l_len)`, it is guaranteed
+    /// that `left` and `right` are contiguous and non-overlapping if
+    /// `self.padding_needed_for_len(l_len) == Some(0)`. This is true for
+    /// all `[T]`.
+    ///
+    /// If `self.padding_needed_for_len(l_len) != Some(0)`, then the left
+    /// pointer will overlap the right pointer to satisfy `T`'s padding
+    /// requirements.
+    pub(crate) unsafe fn split_at_unchecked(
+        self,
+        l_len: usize,
+    ) -> (Self, PtrInner<'a, [T::TrailingElem]>) {
+        // SAFETY: The caller promises that `l_len <= self.len()`.
+        // Trivially, `0 <= l_len`.
+        let left = unsafe { self.with_meta(l_len) };
+
+        let right = self.trailing_slice();
+        // SAFETY: The caller promises that `l_len <= self.len() =
+        // slf.len()`. Trivially, `slf.len() <= slf.len()`.
+        let right = unsafe { right.slice_unchecked(l_len..self.len()) };
+
+        // SAFETY: If `self.padding_needed_for_len(l_len) == Some(0)`, then
+        // `left` and `right` are non-overlapping. Proof: `left` is constructed
+        // `slf` with `l_len` as its (exclusive) upper bound. If
+        // `self.padding_needed_for_len(l_len) == Some(0)`, then `left` requires
+        // no trailing padding following its final element. Since `right` is
+        // constructed from `slf`'s trailing slice with `l_len` as its
+        // (inclusive) lower bound, no byte is referred to by both pointers.
+        //
+        // Conversely, `self.padding_needed_for_len(l_len) == Some(N)`, where `N
+        // > 0`, `left` requires `N` bytes of trailing padding following its
+        // final element. Since `right` is constructed from the trailing slice
+        // of `slf` with `l_len` as its (inclusive) lower bound, the first `N`
+        // bytes of `right` are aliased by `left`.
+        (left, right)
+    }
+
+    /// Produces the number of padding bytes that must follow the trailing
+    /// element of `self` if its trailing slice has length `len`.
+    ///
+    /// # Safety
+    ///
+    /// Unsafe code my rely on `len` satisfying the above contract.
+    #[allow(unused)]
+    pub(crate) fn padding_needed_for_len(&self, len: usize) -> Option<usize> {
+        len.padding_needed_for(T::LAYOUT)
+    }
+
+    /// Produces the trailing slice of `self`.
+    ///
+    /// # Safety
+    ///
+    /// Unsafe code my rely on `len` satisfying the above contract.
+    pub(crate) fn trailing_slice(self) -> PtrInner<'a, [T::TrailingElem]> {
+        let offset = match T::LAYOUT.size_info {
+            crate::SizeInfo::Sized { size } => size,
+            crate::SizeInfo::SliceDst(crate::TrailingSliceLayout { offset, .. }) => offset,
+        };
+
+        let bytes = self.as_non_null().cast::<u8>().as_ptr();
+
+        // SAFETY: TODO.
+        let bytes = unsafe { bytes.add(offset) };
+
+        // SAFETY: TODO
+        let bytes = unsafe { NonNull::new_unchecked(bytes) };
+
+        let raw = KnownLayout::raw_from_ptr_len(bytes, self.len());
+
+        // SAFETY: TODO
+        unsafe { PtrInner::new(raw) }
+    }
+}
+
 #[allow(clippy::needless_lifetimes)]
 impl<'a, T> PtrInner<'a, [T]> {
     /// Creates a pointer which addresses the given `range` of self.
@@ -200,31 +353,6 @@ impl<'a, T> PtrInner<'a, [T]> {
         unsafe { PtrInner::new(ptr) }
     }
 
-    /// Splits the slice in two.
-    ///
-    /// # Safety
-    ///
-    /// The caller promises that `l_len <= self.len()`.
-    ///
-    /// Given `let (left, right) = ptr.split_at(l_len)`, it is guaranteed
-    /// that `left` and `right` are contiguous and non-overlapping.
-    pub(crate) unsafe fn split_at(self, l_len: usize) -> (Self, Self) {
-        // SAFETY: The caller promises that `l_len <= self.len()`.
-        // Trivially, `0 <= l_len`.
-        let left = unsafe { self.slice_unchecked(0..l_len) };
-
-        // SAFETY: The caller promises that `l_len <= self.len() =
-        // slf.len()`. Trivially, `slf.len() <= slf.len()`.
-        let right = unsafe { self.slice_unchecked(l_len..self.len()) };
-
-        // SAFETY: `left` and `right` are non-overlapping. Proof: `left` is
-        // constructed from `slf` with `l_len` as its (exclusive) upper
-        // bound, while `right` is constructed from `slf` with `l_len` as
-        // its (inclusive) lower bound. Thus, no index is a member of both
-        // ranges.
-        (left, right)
-    }
-
     /// Iteratively projects the elements `PtrInner<T>` from `PtrInner<[T]>`.
     pub(crate) fn iter(&self) -> impl Iterator<Item = PtrInner<'a, T>> {
         // TODO(#429): Once `NonNull::cast` documents that it preserves
@@ -297,32 +425,6 @@ impl<'a, T> PtrInner<'a, [T]> {
             unsafe { PtrInner::new(elem) }
         })
     }
-
-    /// The number of slice elements in the object referenced by `self`.
-    ///
-    /// # Safety
-    ///
-    /// Unsafe code my rely on `len` satisfying the above contract.
-    pub(crate) fn len(&self) -> usize {
-        self.trailing_slice_len()
-    }
-}
-
-#[allow(clippy::needless_lifetimes)]
-impl<'a, T> PtrInner<'a, T>
-where
-    T: ?Sized + KnownLayout<PointerMetadata = usize>,
-{
-    /// The number of trailing slice elements in the object referenced by
-    /// `self`.
-    ///
-    /// # Safety
-    ///
-    /// Unsafe code my rely on `trailing_slice_len` satisfying the above
-    /// contract.
-    pub(super) fn trailing_slice_len(&self) -> usize {
-        T::pointer_to_metadata(self.as_non_null().as_ptr())
-    }
 }
 
 impl<'a, T, const N: usize> PtrInner<'a, [T; N]> {
@@ -442,7 +544,11 @@ impl<'a> PtrInner<'a, [u8]> {
 
         // SAFETY: `validate_cast_and_convert_metadata` promises to return
         // `split_at <= self.len()`.
-        let (l_slice, r_slice) = unsafe { self.split_at(split_at) };
+        //
+        // Lemma 0: `l_slice` and `r_slice` are non-overlapping. Proof: By
+        // contract on `PtrInner::split_at_unchecked`, the produced `PtrInner`s
+        // are always non-overlapping if `self` is a `[T]`; here it is a `[u8]`.
+        let (l_slice, r_slice) = unsafe { self.split_at_unchecked(split_at) };
 
         let (target, remainder) = match cast_type {
             CastType::Prefix => (l_slice, r_slice),
@@ -500,7 +606,7 @@ mod tests {
         for i in 0..=N {
             assert_eq!(ptr.len(), N);
             // SAFETY: `i` is in bounds by construction.
-            let (l, r) = unsafe { ptr.split_at(i) };
+            let (l, r) = unsafe { ptr.split_at_unchecked(i) };
             // SAFETY: Points to a valid value by construction.
             #[allow(clippy::undocumented_unsafe_blocks)] // Clippy false positive
             let l_sum: usize = l
diff --git a/src/util/macros.rs b/src/util/macros.rs
@@ -454,6 +454,8 @@ macro_rules! impl_known_layout {
                 //   alignment, and ABI as `T`
                 type MaybeUninit = core::mem::MaybeUninit<Self>;
 
+                type TrailingElem = ();
+
                 const LAYOUT: crate::DstLayout = crate::DstLayout::for_type::<$ty>();
 
                 // SAFETY: `.cast` preserves address and provenance.
@@ -497,6 +499,7 @@ macro_rules! unsafe_impl_known_layout {
 
                 type PointerMetadata = <$repr as KnownLayout>::PointerMetadata;
                 type MaybeUninit = <$repr as KnownLayout>::MaybeUninit;
+                type TrailingElem = <$repr as KnownLayout>::TrailingElem;
 
                 const LAYOUT: DstLayout = <$repr as KnownLayout>::LAYOUT;
 
diff --git a/zerocopy-derive/src/lib.rs b/zerocopy-derive/src/lib.rs
@@ -225,6 +225,8 @@ fn derive_known_layout_inner(ast: &DeriveInput, _top_level: Trait) -> Result<Tok
 
                 type MaybeUninit = __ZerocopyKnownLayoutMaybeUninit #ty_generics;
 
+                type TrailingElem = <#trailing_field_ty as ::zerocopy::KnownLayout>::TrailingElem;
+
                 // 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}`.
@@ -364,6 +366,8 @@ fn derive_known_layout_inner(ast: &DeriveInput, _top_level: Trait) -> Result<Tok
 
                     type MaybeUninit = Self;
 
+                    type TrailingElem = <#trailing_field_ty as ::zerocopy::KnownLayout>::TrailingElem;
+
                     const LAYOUT: ::zerocopy::DstLayout = <#ident #ty_generics as ::zerocopy::KnownLayout>::LAYOUT;
 
                     #methods
@@ -382,6 +386,7 @@ fn derive_known_layout_inner(ast: &DeriveInput, _top_level: Trait) -> Result<Tok
                 type PointerMetadata = ();
                 type MaybeUninit =
                     ::zerocopy::util::macro_util::core_reexport::mem::MaybeUninit<Self>;
+                type TrailingElem = ();
 
                 // SAFETY: `LAYOUT` is guaranteed to accurately describe the
                 // layout of `Self`, because that is the documented safety
