[macros] Support shrinking reference transmutes

In `transmute_ref!` and `transmute_mut!`, support an `#![allow(shrink)]`
attribute which is invoked as follows:

  transmute_ref!(#![allow(shrink)] src);

When this attribute is provided, the macros will permit shrinking
transmutes, in which the destination value may be smaller than the
source value.

Makes progress on #1817

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

Author: joshlf

src/layout.rs

@@ -623,11 +623,15 @@ mod cast_from_raw {
     /// [cast_from_raw]: crate::pointer::SizeFrom::cast_from_raw
     //
     // FIXME(#1817): Support Sized->Unsized and Unsized->Sized casts
-    pub(crate) fn cast_from_raw<Src, Dst>(src: PtrInner<'_, Src>) -> PtrInner<'_, Dst>
+    pub(crate) fn cast_from_raw<Src, Dst, const ALLOW_SHRINK: bool>(
+        src: PtrInner<'_, Src>,
+    ) -> PtrInner<'_, Dst>
     where
         Src: KnownLayout<PointerMetadata = usize> + ?Sized,
         Dst: KnownLayout<PointerMetadata = usize> + ?Sized,
     {
+        // TODO: Update this comment.
+        //
         // At compile time (specifically, post-monomorphization time), we need
         // to compute two things:
         // - Whether, given *any* `*Src`, it is possible to construct a `*Dst`
@@ -694,12 +698,19 @@ mod cast_from_raw {
         /// `Src`'s alignment must not be smaller than `Dst`'s alignment.
         #[derive(Copy, Clone)]
         struct CastParams {
-            offset_delta_elems: usize,
-            elem_multiple: usize,
+            // `offset_delta / dst.elem_size = offset_delta_elems_num / denom`
+            offset_delta_elems_num: usize,
+            // `src.elem_size / dst.elem_size = elem_multiple_num / denom`
+            elem_multiple_num: usize,
+            denom: NonZeroUsize,
         }
 
         impl CastParams {
-            const fn try_compute(src: &DstLayout, dst: &DstLayout) -> Option<CastParams> {
+            const fn try_compute(
+                src: &DstLayout,
+                dst: &DstLayout,
+                allow_shrink: bool,
+            ) -> Option<CastParams> {
                 if src.align.get() < dst.align.get() {
                     return None;
                 }
@@ -724,33 +735,44 @@ mod cast_from_raw {
                     return None;
                 };
 
-                // PANICS: `dst_elem_size: NonZeroUsize`, so this won't div by zero.
-                #[allow(clippy::arithmetic_side_effects)]
-                let delta_mod_other_elem = offset_delta % dst_elem_size.get();
+                const fn gcd(a: usize, b: usize) -> usize {
+                    if a == 0 {
+                        b
+                    } else {
+                        #[allow(clippy::arithmetic_side_effects)]
+                        gcd(b % a, a)
+                    }
+                }
 
-                // PANICS: `dst_elem_size: NonZeroUsize`, so this won't div by zero.
-                #[allow(clippy::arithmetic_side_effects)]
-                let elem_remainder = src.elem_size % dst_elem_size.get();
+                let gcd = gcd(gcd(offset_delta, src.elem_size), dst_elem_size.get());
 
-                if delta_mod_other_elem != 0 || src.elem_size < dst.elem_size || elem_remainder != 0
-                {
-                    return None;
-                }
+                // PANICS: `gcd` is non-zero.
+                #[allow(clippy::arithmetic_side_effects)]
+                let offset_delta_elems_num = offset_delta / gcd;
 
-                // PANICS: `dst_elem_size: NonZeroUsize`, so this won't div by zero.
+                // PANICS: `gcd` is non-zero.
                 #[allow(clippy::arithmetic_side_effects)]
-                let offset_delta_elems = offset_delta / dst_elem_size.get();
+                let elem_multiple_num = src.elem_size / gcd;
 
-                // PANICS: `dst_elem_size: NonZeroUsize`, so this won't div by zero.
+                // PANICS: `dst_elem_size` is non-zero, and `gcd` is no greater
+                // than it by construction. Thus, this should be at least 1.
                 #[allow(clippy::arithmetic_side_effects)]
-                let elem_multiple = src.elem_size / dst_elem_size.get();
+                let denom = match NonZeroUsize::new(dst_elem_size.get() / gcd) {
+                    Some(d) => d,
+                    None => const_panic!("CastParams::try_compute: denom should be non-zero"),
+                };
+
+                if denom.get() != 1 && !allow_shrink {
+                    return None;
+                }
 
                 // SAFETY: We checked above that `src.align >= dst.align`.
                 Some(CastParams {
                     // SAFETY: We checked above that this is an exact ratio.
-                    offset_delta_elems,
+                    offset_delta_elems_num,
                     // SAFETY: We checked above that this is an exact ratio.
-                    elem_multiple,
+                    elem_multiple_num,
+                    denom,
                 })
             }
 
@@ -759,12 +781,17 @@ mod cast_from_raw {
             /// `src_meta` describes a `Src` whose size is no larger than
             /// `isize::MAX`.
             ///
-            /// The returned metadata describes a `Dst` of the same size as the
-            /// original `Src`.
+            /// If `self.denom == 1`, then the returned metadata describes a
+            /// `Dst` of the same size as the original `Src`. Otherwise, the
+            /// returned metadata describes a `Dst` whose size is no greater
+            /// than the size of the original `Src`.
             unsafe fn cast_metadata(self, src_meta: usize) -> usize {
                 #[allow(unused)]
                 use crate::util::polyfills::*;
 
+                // TODO: Update this safety comment. Make sure that even if
+                // `denom > 1`, these arithmetic operations will not overflow.
+                //
                 // SAFETY: `self` is a witness that the following equation
                 // holds:
                 //
@@ -774,24 +801,25 @@ mod cast_from_raw {
                 // metadata, this math will not overflow, and the returned value
                 // will describe a `Dst` of the same size.
                 #[allow(unstable_name_collisions)]
-                unsafe {
-                    self.offset_delta_elems
-                        .unchecked_add(src_meta.unchecked_mul(self.elem_multiple))
-                }
+                let num = unsafe {
+                    self.offset_delta_elems_num
+                        .unchecked_add(src_meta.unchecked_mul(self.elem_multiple_num))
+                };
+                num / self.denom
             }
         }
 
-        trait Params<Src: ?Sized> {
+        trait Params<Src: ?Sized, const ALLOW_SHRINK: bool> {
             const CAST_PARAMS: CastParams;
         }
 
-        impl<Src, Dst> Params<Src> for Dst
+        impl<Src, Dst, const ALLOW_SHRINK: bool> Params<Src, ALLOW_SHRINK> for Dst
         where
             Src: KnownLayout + ?Sized,
             Dst: KnownLayout<PointerMetadata = usize> + ?Sized,
         {
             const CAST_PARAMS: CastParams =
-                match CastParams::try_compute(&Src::LAYOUT, &Dst::LAYOUT) {
+                match CastParams::try_compute(&Src::LAYOUT, &Dst::LAYOUT, ALLOW_SHRINK) {
                     Some(params) => params,
                     None => const_panic!(
                         "cannot `transmute_ref!` or `transmute_mut!` between incompatible types"
@@ -800,7 +828,7 @@ mod cast_from_raw {
         }
 
         let src_meta = <Src as KnownLayout>::pointer_to_metadata(src.as_non_null().as_ptr());
-        let params = <Dst as Params<Src>>::CAST_PARAMS;
+        let params = <Dst as Params<Src, ALLOW_SHRINK>>::CAST_PARAMS;
 
         // SAFETY: `src: PtrInner`, and so by invariant on `PtrInner`, `src`'s
         // referent is no larger than `isize::MAX`.
@@ -809,11 +837,11 @@ mod cast_from_raw {
         let dst = <Dst as KnownLayout>::raw_from_ptr_len(src.as_non_null().cast(), dst_meta);
 
         // SAFETY: By post-condition on `params.cast_metadata`, `dst` addresses
-        // the same number of bytes as `src`. Since `src: PtrInner`, `src` has
-        // provenance for its entire referent, which lives inside of a single
-        // allocation. Since `dst` has the same address as `src` and was
-        // constructed using provenance-preserving operations, it addresses a
-        // subset of those bytes, and has provenance for those bytes.
+        // no more bytes than `src`. Since `src: PtrInner`, `src` has provenance
+        // for its entire referent, which lives inside of a single allocation.
+        // Since `dst` has the same address as `src` and was constructed using
+        // provenance-preserving operations, it addresses a subset of those
+        // bytes, and has provenance for those bytes.
         unsafe { PtrInner::new(dst) }
     }
 }

src/macros.rs

@@ -208,6 +208,24 @@ macro_rules! transmute {
 /// assert_eq!(size_of_val(src), size_of_val(dst));
 /// ```
 ///
+/// ## `#![allow(shrink)]`
+///
+/// If `#![allow(shrink)]` is provided, `transmute_ref!` additionally supports
+/// transmutations that shrink the size of the referent; e.g.:
+///
+/// ```
+/// # use zerocopy::transmute_ref;
+/// # use core::mem::size_of_val; // Not in the prelude on our MSRV
+/// let src: &[[u8; 3]] = &[[0, 1, 2], [3, 4, 5], [6, 7, 8]][..];
+/// let dst: &[[u8; 2]] = transmute_ref!(#![allow(shrink)] src);
+///
+/// assert_eq!(src.len(), 3);
+/// assert_eq!(dst.len(), 4);
+/// assert_eq!(size_of_val(src), 9);
+/// assert_eq!(size_of_val(dst), 8);
+/// assert_eq!(dst, [[0, 1], [2, 3], [4, 5], [6, 7]]);
+/// ```
+///
 /// # Errors
 ///
 /// Violations of the alignment and size compatibility checks are detected
@@ -294,7 +312,18 @@ macro_rules! transmute {
 /// `Dst: Sized`.
 #[macro_export]
 macro_rules! transmute_ref {
-    ($e:expr) => {{
+    (#![allow(shrink)] $e:expr) => {
+        $crate::__transmute_ref_inner!(true, $e)
+    };
+    ($e:expr) => {
+        $crate::__transmute_ref_inner!(false, $e)
+    };
+}
+
+#[macro_export]
+#[doc(hidden)]
+macro_rules! __transmute_ref_inner {
+    ($allow_shrink:literal, $e:expr) => {{
         // NOTE: This must be a macro (rather than a function with trait bounds)
         // because there's no way, in a generic context, to enforce that two
         // types have the same size or alignment.
@@ -333,10 +362,10 @@ macro_rules! transmute_ref {
             // - `Src: IntoBytes + Immutable`
             // - `Dst: FromBytes + Immutable`
             unsafe {
-                t.transmute_ref()
+                t.transmute_ref::<$allow_shrink>()
             }
         }
-    }}
+    }};
 }
 
 /// Safely transmutes a mutable reference of one type to a mutable reference of
@@ -382,6 +411,29 @@ macro_rules! transmute_ref {
 /// assert_eq!(size_of_val(src), dst_size);
 /// ```
 ///
+/// ## `#![allow(shrink)]`
+///
+/// If `#![allow(shrink)]` is provided, `transmute_mut!` additionally supports
+/// transmutations that shrink the size of the referent; e.g.:
+///
+/// ```
+/// # use zerocopy::transmute_mut;
+/// # use core::mem::size_of_val; // Not in the prelude on our MSRV
+/// let src: &mut [[u8; 3]] = &mut [[0, 1, 2], [3, 4, 5], [6, 7, 8]][..];
+/// let dst: &mut [[u8; 2]] = transmute_mut!(#![allow(shrink)] src);
+///
+///
+/// let dst_len = dst.len();
+/// let dst_size = size_of_val(dst);
+/// assert_eq!(dst, [[0, 1], [2, 3], [4, 5], [6, 7]]);
+///
+/// assert_eq!(src.len(), 3);
+/// assert_eq!(dst_len, 4);
+///
+/// assert_eq!(size_of_val(src), 9);
+/// assert_eq!(dst_size, 8);
+/// ```
+///
 /// # Errors
 ///
 /// Violations of the alignment and size compatibility checks are detected
@@ -470,7 +522,18 @@ macro_rules! transmute_ref {
 /// ```
 #[macro_export]
 macro_rules! transmute_mut {
-    ($e:expr) => {{
+    (#![allow(shrink)] $e:expr) => {
+        $crate::__transmute_mut_inner!(true, $e)
+    };
+    ($e:expr) => {
+        $crate::__transmute_mut_inner!(false, $e)
+    };
+}
+
+#[doc(hidden)]
+#[macro_export]
+macro_rules! __transmute_mut_inner {
+    ($allow_shrink:literal, $e:expr) => {{
         // NOTE: This must be a macro (rather than a function with trait bounds)
         // because, for backwards-compatibility on v0.8.x, we use the autoref
         // specialization trick to dispatch to different `transmute_mut`
@@ -484,7 +547,7 @@ macro_rules! transmute_mut {
         #[allow(unused)]
         use $crate::util::macro_util::TransmuteMutDst as _;
         let t = $crate::util::macro_util::Wrap::new(e);
-        t.transmute_mut()
+        t.transmute_mut::<$allow_shrink>()
     }}
 }
 
@@ -1232,6 +1295,11 @@ mod tests {
         let slice_of_u16s: &[U16] = <[U16]>::ref_from_bytes(&[0, 1, 2, 3, 4, 5][..]).unwrap();
         assert_eq!(x, slice_of_u16s);
 
+        // Test that transmuting from a larger sized type to a smaller sized
+        // type works.
+        let x: &u8 = transmute_ref!(#![allow(shrink)] &0u16);
+        assert_eq!(*x, 0);
+
         // Test that transmuting from a type with larger trailing slice offset
         // and larger trailing slice element works.
         let bytes = &[0, 1, 2, 3, 4, 5, 6, 7][..];
@@ -1240,6 +1308,15 @@ mod tests {
         let x: &SliceDst<U16, u8> = transmute_ref!(slice_dst_big);
         assert_eq!(x, slice_dst_small);
 
+        let bytes = &[0, 1, 2, 3, 4, 5, 6, 7][..];
+        let slice_dst_big = SliceDst::<[u8; 4], [u8; 4]>::ref_from_bytes(bytes).unwrap();
+        let slice_dst_small = SliceDst::<[u8; 3], [u8; 3]>::ref_from_bytes(&bytes[..6]).unwrap();
+        let x: &SliceDst<[u8; 3], [u8; 3]> = transmute_ref!(
+            #![allow(shrink)]
+            slice_dst_big
+        );
+        assert_eq!(x, slice_dst_small);
+
         // 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);
@@ -1420,6 +1497,14 @@ mod tests {
         let x: &mut [i16] = transmute_mut!(array_of_u16s);
         assert_eq!(x, array_of_i16s);
 
+        // Test that transmuting from a larger sized type to a smaller sized
+        // type works.
+        let mut large: [u8; 2] = [1, 1];
+        let x: &mut u8 = transmute_mut!(#![allow(shrink)] &mut large);
+        assert_eq!(*x, 1);
+        *x = 0;
+        assert_eq!(large, [0, 1]);
+
         // Test that transmuting from a type with larger trailing slice offset
         // and larger trailing slice element works.
         let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7];
@@ -1428,6 +1513,16 @@ mod tests {
         let slice_dst_small = SliceDst::<U16, u8>::mut_from_bytes(&mut bytes[..]).unwrap();
         let x: &mut SliceDst<U16, u8> = transmute_mut!(slice_dst_big);
         assert_eq!(x, slice_dst_small);
+
+        let mut bytes = [0, 1, 2, 3, 4, 5, 6, 7];
+        let slice_dst_big = SliceDst::<[u8; 4], [u8; 4]>::mut_from_bytes(&mut bytes[..]).unwrap();
+        let mut bytes = [0, 1, 2, 3, 4, 5];
+        let slice_dst_small = SliceDst::<[u8; 3], [u8; 3]>::mut_from_bytes(&mut bytes[..]).unwrap();
+        let x: &mut SliceDst<[u8; 3], [u8; 3]> = transmute_mut!(
+            #![allow(shrink)]
+            slice_dst_big
+        );
+        assert_eq!(x, slice_dst_small);
     }
 
     #[test]