Teach transmute_{ref,mut}! to handle slice DSTs

TODO:
- Convince ourselves that it's acceptable to drop the `TransmuteFrom:
  SizeCompatible` super-trait bound.
- Write Kani tests for layout computations

Makes progress on #1817

Co-authored-by: Jack Wrenn <jswrenn@amazon.com>
gherrit-pr-id: Ib4bc62202e0b3b09d155333b525087f7aa8f02c2

Author: joshlf

src/impls.rs

@@ -179,39 +179,13 @@ safety_comment! {
     });
 }
 
-// SAFETY: `str` and `[u8]` have the same layout [1].
-//
-// [1] Per https://doc.rust-lang.org/1.81.0/reference/type-layout.html#str-layout:
-//
-//   String slices are a UTF-8 representation of characters that have the same
-//   layout as slices of type `[u8]`.
-unsafe impl pointer::SizeEq<str> for [u8] {
-    fn cast_from_raw(s: NonNull<str>) -> NonNull<[u8]> {
-        cast!(s)
-    }
-}
-// SAFETY: See previous safety comment.
-unsafe impl pointer::SizeEq<[u8]> for str {
-    fn cast_from_raw(bytes: NonNull<[u8]>) -> NonNull<str> {
-        cast!(bytes)
-    }
-}
+impl_size_eq!(str, [u8]);
 
 macro_rules! unsafe_impl_try_from_bytes_for_nonzero {
     ($($nonzero:ident[$prim:ty]),*) => {
         $(
             unsafe_impl!(=> TryFromBytes for $nonzero; |n| {
-                unsafe impl pointer::SizeEq<$nonzero> for Unalign<$prim> {
-                    fn cast_from_raw(n: NonNull<$nonzero>) -> NonNull<Unalign<$prim>> {
-                        cast!(n)
-                    }
-                }
-                unsafe impl pointer::SizeEq<Unalign<$prim>> for $nonzero {
-                    fn cast_from_raw(p: NonNull<Unalign<$prim>>) -> NonNull<$nonzero> {
-                        cast!(p)
-                    }
-                }
-
+                impl_size_eq!($nonzero, Unalign<$prim>);
                 let n = n.transmute::<Unalign<$prim>, invariant::Valid, _>();
                 $nonzero::new(n.read_unaligned().into_inner()).is_some()
             });
@@ -430,7 +404,7 @@ mod atomics {
         ($($($tyvar:ident)? => $atomic:ty [$prim:ty]),*) => {
             const _: () = {
                 use core::{cell::UnsafeCell, ptr::NonNull};
-                use crate::pointer::{TransmuteFrom, SizeEq, invariant::Valid};
+                use crate::pointer::{TransmuteFrom, SizeCompat, invariant::Valid};
 
                 $(
                     #[allow(unused_unsafe)] // Force the caller to call this macro inside `safety_comment!`.
@@ -443,36 +417,38 @@ mod atomics {
                     // the same size and bit validity.
                     unsafe impl<$($tyvar)?> TransmuteFrom<$prim, Valid, Valid> for $atomic {}
 
-                    // SAFETY: THe caller promised that `$atomic` and `$prim`
-                    // have the same size.
-                    unsafe impl<$($tyvar)?> SizeEq<$atomic> for $prim {
+                    // SAFETY: The caller promised that `$atomic` and `$prim`
+                    // have the same size. Thus, this cast preserves address,
+                    // referent size, and provenance.
+                    unsafe impl<$($tyvar)?> SizeCompat<$atomic> for $prim {
                         fn cast_from_raw(a: NonNull<$atomic>) -> NonNull<$prim> {
                             cast!(a)
                         }
                     }
-                    // SAFETY: THe caller promised that `$atomic` and `$prim`
-                    // have the same size.
-                    unsafe impl<$($tyvar)?> SizeEq<$prim> for $atomic {
+                    // SAFETY: See previous safety comment.
+                    unsafe impl<$($tyvar)?> SizeCompat<$prim> for $atomic {
                         fn cast_from_raw(p: NonNull<$prim>) -> NonNull<$atomic> {
                             cast!(p)
                         }
                     }
+
                     // SAFETY: The caller promised that `$atomic` and `$prim`
                     // have the same size. `UnsafeCell<T>` has the same size as
-                    // `T` [1].
+                    // `T` [1]. Thus, this cast preserves address, referent
+                    // size, and provenance.
                     //
                     // [1] Per https://doc.rust-lang.org/1.85.0/std/cell/struct.UnsafeCell.html#memory-layout:
                     //
                     //   `UnsafeCell<T>` has the same in-memory representation as
                     //   its inner type `T`. A consequence of this guarantee is that
                     //   it is possible to convert between `T` and `UnsafeCell<T>`.
-                    unsafe impl<$($tyvar)?> SizeEq<$atomic> for UnsafeCell<$prim> {
+                    unsafe impl<$($tyvar)?> SizeCompat<$atomic> for UnsafeCell<$prim> {
                         fn cast_from_raw(a: NonNull<$atomic>) -> NonNull<UnsafeCell<$prim>> {
                             cast!(a)
                         }
                     }
                     // SAFETY: See previous safety comment.
-                    unsafe impl<$($tyvar)?> SizeEq<UnsafeCell<$prim>> for $atomic {
+                    unsafe impl<$($tyvar)?> SizeCompat<UnsafeCell<$prim>> for $atomic {
                         fn cast_from_raw(p: NonNull<UnsafeCell<$prim>>) -> NonNull<$atomic> {
                             cast!(p)
                         }

src/layout.rs

@@ -92,6 +92,17 @@ pub(crate) enum MetadataCastError {
     Size,
 }
 
+pub(crate) struct CastParams {
+    pub(crate) offset_delta_elems: usize,
+    pub(crate) elem_multiple: usize,
+}
+
+impl CastParams {
+    pub(crate) const fn compute_cast(self, meta: usize) -> usize {
+        self.offset_delta_elems + (meta * self.elem_multiple)
+    }
+}
+
 impl DstLayout {
     /// The minimum possible alignment of a type.
     const MIN_ALIGN: NonZeroUsize = match NonZeroUsize::new(1) {
@@ -202,6 +213,27 @@ impl DstLayout {
         }
     }
 
+    pub(crate) const fn try_compute_cast_params(self, other: &DstLayout) -> Option<CastParams> {
+        if let (SizeInfo::SliceDst(slf), SizeInfo::SliceDst(other)) =
+            (self.size_info, other.size_info)
+        {
+            if let Some(offset_delta) = slf.offset.checked_sub(other.offset) {
+                if offset_delta % other.elem_size == 0
+                    && other.elem_size > 0
+                    && slf.elem_size >= other.elem_size
+                    && slf.elem_size % other.elem_size == 0
+                {
+                    return Some(CastParams {
+                        offset_delta_elems: offset_delta / other.elem_size,
+                        elem_multiple: slf.elem_size / other.elem_size,
+                    });
+                }
+            }
+        }
+
+        None
+    }
+
     /// Like `Layout::extend`, this creates a layout that describes a record
     /// whose layout consists of `self` followed by `next` that includes the
     /// necessary inter-field padding, but not any trailing padding.

src/macros.rs

@@ -182,22 +182,15 @@ macro_rules! transmute_ref {
             // this macro expression is `&U` where `U: 'u + Sized + FromBytes +
             // Immutable`, and that `'t` outlives `'u`.
 
-            struct AssertSrcIsSized<'a, T: ::core::marker::Sized>(&'a T);
             struct AssertSrcIsIntoBytes<'a, T: ?::core::marker::Sized + $crate::IntoBytes>(&'a T);
             struct AssertSrcIsImmutable<'a, T: ?::core::marker::Sized + $crate::Immutable>(&'a T);
-            struct AssertDstIsSized<'a, T: ::core::marker::Sized>(&'a T);
             struct AssertDstIsFromBytes<'a, U: ?::core::marker::Sized + $crate::FromBytes>(&'a U);
             struct AssertDstIsImmutable<'a, T: ?::core::marker::Sized + $crate::Immutable>(&'a T);
 
-            let _ = AssertSrcIsSized(e);
             let _ = AssertSrcIsIntoBytes(e);
             let _ = AssertSrcIsImmutable(e);
 
             if true {
-                #[allow(unused, unreachable_code)]
-                let u = AssertDstIsSized(loop {});
-                u.0
-            } else if true {
                 #[allow(unused, unreachable_code)]
                 let u = AssertDstIsFromBytes(loop {});
                 u.0
@@ -206,36 +199,13 @@ macro_rules! transmute_ref {
                 let u = AssertDstIsImmutable(loop {});
                 u.0
             }
-        } else if false {
-            // This branch, though never taken, ensures that `size_of::<T>() ==
-            // size_of::<U>()` and that that `align_of::<T>() >=
-            // align_of::<U>()`.
-
-            // `t` is inferred to have type `T` because it's assigned to `e` (of
-            // type `&T`) as `&t`.
-            let mut t = loop {};
-            e = &t;
-
-            // `u` is inferred to have type `U` because it's used as `&u` as the
-            // value returned from this branch.
-            let u;
-
-            $crate::assert_size_eq!(t, u);
-            $crate::assert_align_gt_eq!(t, u);
-
-            &u
         } else {
-            // SAFETY: For source type `Src` and destination type `Dst`:
-            // - We know that `Src: IntoBytes + Immutable` and `Dst: FromBytes +
-            //   Immutable` thanks to the uses of `AssertSrcIsIntoBytes`,
-            //   `AssertSrcIsImmutable`, `AssertDstIsFromBytes`, and
-            //   `AssertDstIsImmutable` above.
-            // - We know that `size_of::<Src>() == size_of::<Dst>()` thanks to
-            //   the use of `assert_size_eq!` above.
-            // - We know that `align_of::<Src>() >= align_of::<Dst>()` thanks to
-            //   the use of `assert_align_gt_eq!` above.
-            let u = unsafe { $crate::util::macro_util::transmute_ref(e) };
-            $crate::util::macro_util::must_use(u)
+            use $crate::util::macro_util::TransmuteRefDst;
+            let t = $crate::util::macro_util::Wrap(e);
+            // SAFETY: todo
+            unsafe {
+                t.transmute_ref()
+            }
         }
     }}
 }
@@ -334,26 +304,18 @@ macro_rules! transmute_mut {
             // writing, mutable references are banned), the error message
             // appears to originate in the user's code rather than in the
             // internals of this macro.
-            struct AssertSrcIsSized<'a, T: ::core::marker::Sized>(&'a T);
             struct AssertSrcIsFromBytes<'a, T: ?::core::marker::Sized + $crate::FromBytes>(&'a T);
             struct AssertSrcIsIntoBytes<'a, T: ?::core::marker::Sized + $crate::IntoBytes>(&'a T);
-            struct AssertDstIsSized<'a, T: ::core::marker::Sized>(&'a T);
             struct AssertDstIsFromBytes<'a, T: ?::core::marker::Sized + $crate::FromBytes>(&'a T);
             struct AssertDstIsIntoBytes<'a, T: ?::core::marker::Sized + $crate::IntoBytes>(&'a T);
 
             if true {
-                let _ = AssertSrcIsSized(&*e);
-            } else if true {
                 let _ = AssertSrcIsFromBytes(&*e);
             } else {
                 let _ = AssertSrcIsIntoBytes(&*e);
             }
 
             if true {
-                #[allow(unused, unreachable_code)]
-                let u = AssertDstIsSized(loop {});
-                &mut *u.0
-            } else if true {
                 #[allow(unused, unreachable_code)]
                 let u = AssertDstIsFromBytes(loop {});
                 &mut *u.0
@@ -362,32 +324,13 @@ macro_rules! transmute_mut {
                 let u = AssertDstIsIntoBytes(loop {});
                 &mut *u.0
             }
-        } else if false {
-            // This branch, though never taken, ensures that `size_of::<T>() ==
-            // size_of::<U>()` and that that `align_of::<T>() >=
-            // align_of::<U>()`.
-
-            // `t` is inferred to have type `T` because it's assigned to `e` (of
-            // type `&mut T`) as `&mut t`.
-            let mut t = loop {};
-            e = &mut t;
-
-            // `u` is inferred to have type `U` because it's used as `&mut u` as
-            // the value returned from this branch.
-            let u;
-
-            $crate::assert_size_eq!(t, u);
-            $crate::assert_align_gt_eq!(t, u);
-
-            &mut u
         } else {
-            // SAFETY: For source type `Src` and destination type `Dst`:
-            // - We know that `size_of::<Src>() == size_of::<Dst>()` thanks to
-            //   the use of `assert_size_eq!` above.
-            // - We know that `align_of::<Src>() >= align_of::<Dst>()` thanks to
-            //   the use of `assert_align_gt_eq!` above.
-            let u = unsafe { $crate::util::macro_util::transmute_mut(e) };
-            $crate::util::macro_util::must_use(u)
+            use $crate::util::macro_util::TransmuteMutDst;
+            let t = $crate::util::macro_util::Wrap(e);
+            // SAFETY: todo
+            unsafe {
+                t.transmute_mut()
+            }
         }
     }}
 }
@@ -1038,6 +981,18 @@ mod tests {
         assert_eq!(x.into_inner(), 1);
     }
 
+    // A `Sized` type which doesn't implement `KnownLayout` (it is "not
+    // `KnownLayout`", or `Nkl`).
+    //
+    // This permits us to test that `transmute_ref!` and `transmute_mut!` work
+    // for types which are `Sized + !KnownLayout`. When we added support for
+    // slice DSTs in #1924, this new support relied on `KnownLayout`, but we
+    // need to make sure to remain backwards-compatible with code which ueses
+    // these macros with types which are `!KnownLayout`.
+    #[derive(FromBytes, IntoBytes, Immutable, PartialEq, Eq, Debug)]
+    #[repr(transparent)]
+    struct Nkl<T>(T);
+
     #[test]
     fn test_transmute_ref() {
         // Test that memory is transmuted as expected.
@@ -1055,6 +1010,49 @@ mod tests {
         const X: &'static [[u8; 2]; 4] = transmute_ref!(&ARRAY_OF_U8S);
         assert_eq!(*X, ARRAY_OF_ARRAYS);
 
+        // Call through a generic function to make sure our autoref
+        // specialization trick works even when types are generic.
+        const fn transmute_ref<T, U>(t: &T) -> &U
+        where
+            T: IntoBytes + Immutable,
+            U: FromBytes + Immutable,
+        {
+            transmute_ref!(t)
+        }
+
+        // Test that `transmute_ref!` supports non-`KnownLayout` `Sized` types.
+        const ARRAY_OF_NKL_U8S: Nkl<[u8; 8]> = Nkl([0u8, 1, 2, 3, 4, 5, 6, 7]);
+        const ARRAY_OF_NKL_ARRAYS: Nkl<[[u8; 2]; 4]> = Nkl([[0, 1], [2, 3], [4, 5], [6, 7]]);
+        const X_NKL: &Nkl<[[u8; 2]; 4]> = transmute_ref(&ARRAY_OF_NKL_U8S);
+        assert_eq!(*X_NKL, ARRAY_OF_NKL_ARRAYS);
+
+        // Test that `transmute_ref!` works on slice DSTs in and that memory is
+        // transmuted as expected.
+        #[derive(KnownLayout, Immutable, FromBytes, IntoBytes, PartialEq, Debug)]
+        #[repr(C)]
+        struct SliceDst<T> {
+            a: U16,
+            b: [T],
+        }
+
+        use crate::byteorder::native_endian::U16;
+        let slice_dst_of_u8s =
+            SliceDst::<[u8; 2]>::ref_from_bytes(&[0, 1, 2, 3, 4, 5][..]).unwrap();
+        let slice_dst_of_u16s = SliceDst::<U16>::ref_from_bytes(&[0, 1, 2, 3, 4, 5][..]).unwrap();
+        let x: &SliceDst<U16> = transmute_ref!(slice_dst_of_u8s);
+        assert_eq!(x, slice_dst_of_u16s);
+
+        let slice_dst_of_u8s = SliceDst::<u8>::ref_from_bytes(&[0, 1, 2, 3, 4, 5][..]).unwrap();
+        let x: &[u8] = transmute_ref!(slice_dst_of_u8s);
+        assert_eq!(x, [0, 1, 2, 3, 4, 5]);
+
+        let x: &[u8] = transmute_ref!(slice_dst_of_u16s);
+        assert_eq!(x, [0, 1, 2, 3, 4, 5]);
+
+        let x: &[U16] = transmute_ref!(slice_dst_of_u16s);
+        let slice_of_u16s: &[U16] = <[U16]>::ref_from_bytes(&[0, 1, 2, 3, 4, 5][..]).unwrap();
+        assert_eq!(x, slice_of_u16s);
+
         // Test that it's legal to transmute a reference while shrinking the
         // lifetime (note that `X` has the lifetime `'static`).
         let x: &[u8; 8] = transmute_ref!(X);
@@ -1198,6 +1196,15 @@ mod tests {
             let x: &mut [u8; 8] = transmute_mut!(&mut array_of_arrays);
             assert_eq!(*x, array_of_u8s);
         }
+
+        // Test that `transmute_mut!` supports non-`KnownLayout` types.
+        let mut array_of_u8s = Nkl([0u8, 1, 2, 3, 4, 5, 6, 7]);
+        let mut array_of_arrays = Nkl([[0, 1], [2, 3], [4, 5], [6, 7]]);
+        let x: &mut Nkl<[[u8; 2]; 4]> = transmute_mut!(&mut array_of_u8s);
+        assert_eq!(*x, array_of_arrays);
+        let x: &mut Nkl<[u8; 8]> = transmute_mut!(&mut array_of_arrays);
+        assert_eq!(*x, array_of_u8s);
+
         // Test that `transmute_mut!` supports decreasing alignment.
         let mut u = AU64(0);
         let array = [0, 0, 0, 0, 0, 0, 0, 0];
@@ -1209,6 +1216,30 @@ mod tests {
         #[allow(clippy::useless_transmute)]
         let y: &u8 = transmute_mut!(&mut x);
         assert_eq!(*y, 0);
+
+        // Test that `transmute_mut!` works on slice DSTs in and that memory is
+        // transmuted as expected.
+        #[derive(KnownLayout, Immutable, FromBytes, IntoBytes, PartialEq, Debug)]
+        #[repr(C)]
+        struct SliceDst<T> {
+            a: u8,
+            b: [T],
+        }
+
+        use crate::byteorder::native_endian::U16;
+        let mut bytes = [0, 1, 2, 3, 4, 5, 6];
+        let slice_dst_of_u8s = SliceDst::<[u8; 2]>::mut_from_bytes(&mut bytes[..]).unwrap();
+        let mut bytes = [0, 1, 2, 3, 4, 5, 6];
+        let slice_dst_of_u16s = SliceDst::<U16>::mut_from_bytes(&mut bytes[..]).unwrap();
+        let x: &mut SliceDst<U16> = transmute_mut!(slice_dst_of_u8s);
+        assert_eq!(x, slice_dst_of_u16s);
+
+        // Test that `transmute_mut!` works on slices that memory is transmuted
+        // as expected.
+        let array_of_u16s: &mut [u16] = &mut [0u16, 1, 2];
+        let array_of_i16s: &mut [i16] = &mut [0i16, 1, 2];
+        let x: &mut [i16] = transmute_mut!(array_of_u16s);
+        assert_eq!(x, array_of_i16s);
     }
 
     #[test]