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

Implement a `transmute_unchecked/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 <[email protected]>
---
 rust/kernel/device_id.rs   |  3 +--
 rust/kernel/lib.rs         |  1 +
 rust/kernel/mem.rs         | 37 +++++++++++++++++++++++++++++++++++++
 rust/kernel/sync/atomic.rs |  4 ++--
 4 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/rust/kernel/device_id.rs b/rust/kernel/device_id.rs
index c81fca5b4986..cb8661dec2ab 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_unchecked(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..a0901cbe1b2d
--- /dev/null
+++ b/rust/kernel/mem.rs
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Basic utilities for dealing with memory, values, and types.
+
+use crate::prelude::*;
+
+/// 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.
+///
+/// The signature is equivalent after 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`].
+#[inline(always)]
+pub const unsafe fn transmute_unchecked<Src, Dst>(val: Src) -> Dst {
+    const_assert!(size_of::<Src>() == size_of::<Dst>());
+
+    // SAFETY: This is identical to `transmute` except that we bypassed the size check; which we
+    // used `const_assert!` to check above.
+    unsafe { core::mem::transmute_copy(&core::mem::ManuallyDrop::new(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.
+#[inline(always)]
+pub const fn transmute<Src: IntoBytes, Dst: FromBytes>(val: Src) -> Dst {
+    // SAFETY: transmute is safe with `IntoBytes` and `FromBytes` bounds.
+    unsafe { transmute_unchecked(val) }
+}
diff --git a/rust/kernel/sync/atomic.rs b/rust/kernel/sync/atomic.rs
index 9cd009d57e35..d5f4ef285193 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_unchecked(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_unchecked(r) }
 }
 
 impl<T: AtomicType> Atomic<T> {

-- 
2.54.0