[PATCH v4 02/16] rust: mem: add `transmute` with deferred size check

[Gary Guo]

Implement a `transmute/safe_transmute` that checks size at
monomorphization time instead of type-checking time. This allows more cases
where we know that the size matches but this is not generically checkable.

The signature is equivalent to the unstable `transmute_neo` function in the
standard library. A safe variant is provided to use with types implementing
`FromBytes` and `IntoBytes`.

Existing users of `transmute_copy` to bypass size checks are converted.

Signed-off-by: Gary Guo <gary@xxxxxxxxxxx>
---
Changes since v3:
- Renamed the methods to `transmute` and `safe_transmute`, so the name
  `transmute_unchecked` is for the completely unchecked variant instead,
  matching that of Rust intrinsics.
---
 rust/kernel/device_id.rs   |  3 +-
 rust/kernel/lib.rs         |  1 +
 rust/kernel/mem.rs         | 99 ++++++++++++++++++++++++++++++++++++++++++++++
 rust/kernel/sync/atomic.rs |  4 +-
 4 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/rust/kernel/device_id.rs b/rust/kernel/device_id.rs
index c81fca5b4986..f0b9cb84e58e 100644
--- a/rust/kernel/device_id.rs
+++ b/rust/kernel/device_id.rs
@@ -146,8 +146,7 @@ impl<T: RawDeviceId, const N: usize> IdArray<T, (), N> {
     /// If the device implements [`RawDeviceIdIndex`], consider using [`IdArray::new`] instead.
     pub const fn new_without_index(ids: [T; N]) -> Self {
         // SAFETY: `T` is layout-wise compatible with `T::RawType`, so is the array of them.
-        let raw_ids: [MaybeUninit<T::RawType>; N] = unsafe { core::mem::transmute_copy(&ids) };
-        core::mem::forget(ids);
+        let raw_ids: [MaybeUninit<T::RawType>; N] = unsafe { crate::mem::transmute(ids) };
 
         Self {
             ids: raw_ids,
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 4d5c96ddc49c..7225abc64084 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -98,6 +98,7 @@
 pub mod kunit;
 pub mod list;
 pub mod maple_tree;
+pub mod mem;
 pub mod miscdevice;
 pub mod mm;
 pub mod module;
diff --git a/rust/kernel/mem.rs b/rust/kernel/mem.rs
new file mode 100644
index 000000000000..958e43bbcc3a
--- /dev/null
+++ b/rust/kernel/mem.rs
@@ -0,0 +1,99 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Basic utilities for dealing with memory, values, and types.
+
+use crate::prelude::*;
+
+/// Transmute between two types.
+///
+/// Use this instead of [`core::mem::transmute`] when it is known that sizes are identical but this
+/// cannot be proven by the compiler.
+///
+/// This is equivalent to Rust's `transmute_unchecked` intrinsics.
+///
+/// # Safety
+///
+/// All safety requirements of [`core::mem::transmute`] apply, plus that the size `Src` and `Dst`
+/// must match.
+///
+/// # Example
+///
+/// This can be used when types are known to have the same size, but only at runtime.
+/// ```no_run
+/// # use core::any::TypeId;
+/// fn to_u32<T: 'static>(v: T) -> Option<u32> {
+///     if TypeId::of::<T>() != TypeId::of::<u32>() {
+///         return None;
+///     }
+///
+///     // `core::mem::transmute` won't work here.
+///     // SAFETY: We've checked that `T` is u32!
+///     Some(unsafe { kernel::mem::transmute_unchecked(v) })
+/// }
+///
+/// to_u32(1u32);
+#[inline(always)]
+pub const unsafe fn transmute_unchecked<Src, Dst>(val: Src) -> Dst {
+    // SAFETY: This is identical to `transmute` except that we bypassed the size check; which is
+    // true per safety requirement.
+    unsafe { core::mem::transmute_copy(&core::mem::ManuallyDrop::new(val)) }
+}
+
+/// Version of `transmute` that performs size check at monomorphization-time.
+///
+/// Use this instead of [`core::mem::transmute`] when it is known that sizes are identical but this
+/// cannot be proven by the compiler during type checking and can be proven during monomorphization.
+///
+/// The signature is equivalent to Rust standard library's unstable `transmute_neo` and that of
+/// [RFC 3844](https://github.com/rust-lang/rfcs/pull/3844).
+///
+/// # Safety
+///
+/// Same as [`core::mem::transmute`].
+///
+/// # Examples
+///
+/// This is typically used in generic code where it's known that type will have the same size, but
+/// the compiler cannot prove it generically.
+/// ```no_run
+/// trait IsU32 {}
+/// impl IsU32 for u32 {}
+///
+/// fn to_u32<T: IsU32>(v: T) -> u32 {
+///     // `core::mem::transmute` won't work here.
+///     // SAFETY: We know that `v` is u32!
+///     unsafe { kernel::mem::transmute(v) }
+/// }
+///
+/// to_u32(1u32);
+/// ```
+#[inline(always)]
+pub const unsafe fn transmute<Src, Dst>(val: Src) -> Dst {
+    const_assert!(size_of::<Src>() == size_of::<Dst>());
+
+    // SAFETY: Size is checked above. Other safety requirements follow those of the function.
+    unsafe { transmute_unchecked(val) }
+}
+
+/// Safely transmutes a value of one type to a value of another type of the same size.
+///
+/// The sizes are checked during monomorphization.
+///
+/// This can be considered as generic version of [`zerocopy::transmute!`] macro that defers the size
+/// check and thus can be used in more cases.
+///
+/// # Examples
+///
+/// ```no_run
+/// fn to_u32<T: FromBytes + IntoBytes>(v: T) -> u32 {
+///     // `zerocopy::transmute!` won't work here.
+///     kernel::mem::safe_transmute(v)
+/// }
+///
+/// to_u32(1i32);
+/// ```
+#[inline(always)]
+pub const fn safe_transmute<Src: IntoBytes, Dst: FromBytes>(val: Src) -> Dst {
+    // SAFETY: transmute is safe with `IntoBytes` and `FromBytes` bounds.
+    unsafe { transmute(val) }
+}
diff --git a/rust/kernel/sync/atomic.rs b/rust/kernel/sync/atomic.rs
index 9cd009d57e35..6d27898add42 100644
--- a/rust/kernel/sync/atomic.rs
+++ b/rust/kernel/sync/atomic.rs
@@ -140,7 +140,7 @@ pub unsafe trait AtomicAdd<Rhs = Self>: AtomicType {
 const fn into_repr<T: AtomicType>(v: T) -> T::Repr {
     // SAFETY: Per the safety requirement of `AtomicType`, `T` is round-trip transmutable to
     // `T::Repr`, therefore the transmute operation is sound.
-    unsafe { core::mem::transmute_copy(&v) }
+    unsafe { crate::mem::transmute(v) }
 }
 
 /// # Safety
@@ -149,7 +149,7 @@ const fn into_repr<T: AtomicType>(v: T) -> T::Repr {
 #[inline(always)]
 const unsafe fn from_repr<T: AtomicType>(r: T::Repr) -> T {
     // SAFETY: Per the safety requirement of the function, the transmute operation is sound.
-    unsafe { core::mem::transmute_copy(&r) }
+    unsafe { crate::mem::transmute(r) }
 }
 
 impl<T: AtomicType> Atomic<T> {

-- 
2.54.0