[PATCH v4 03/16] rust: mem: add `AsRepr` and `AsReprMut`

[Gary Guo <gary@garyguo.net>]

Some API like atomics and I/O operate on primitives only; therefore other
types would need to converted to these primitive first. Add two traits
`AsRepr` and `AsReprMut` to indicate that the type can be turned into a
primitive for these operations.

`T: AsRepr` means that `&T` can be viewed as `&T::Repr` and thus it needs
to support transmutability in one direction. `T: AsReprMut` means that
`&mut T` can be viewed as `&mut T::Repr` and thus it needs to support
bi-directional transmutability.

To avoid duplicating implementations, all repr types are normalized to
unsigned integers.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
Changes since v3:
- Added `as_repr` and `as_repr_mut` so the name makes sense (Alex).
  I did not add `Copy` bound because it is not necessary (adding it won't
  simplify the implementation without also requring `Self: Copy`, which I
  don't want to add).

- Dropped round-trip transmutability as a concept, as it can be proved
  correct because `from_repr_unchecked(into_repr(v))` is basically just
  `move v`.

  Even for the `AtomicType`'s use case where copies are made, it can still
  be proven by observing that `from_repr_unchecked(copy into_repr(v))` is
  not distinguishable from `from_repr_unchecked(into_repr(copy v))`.
---
 rust/kernel/mem.rs | 132 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 132 insertions(+)

diff --git a/rust/kernel/mem.rs b/rust/kernel/mem.rs
index 958e43bbcc3a..0bfd6929ae45 100644
--- a/rust/kernel/mem.rs
+++ b/rust/kernel/mem.rs
@@ -97,3 +97,135 @@ pub const fn safe_transmute<Src: IntoBytes, Dst: FromBytes>(val: Src) -> Dst {
     // SAFETY: transmute is safe with `IntoBytes` and `FromBytes` bounds.
     unsafe { transmute(val) }
 }
+
+/// Type that is layout-compatible with a primitive representation.
+///
+/// # Safety
+///
+/// - `Self` must have the same size and alignment as [`Self::Repr`].
+/// - `Self` can be [transmutable] to [`Self::Repr`].
+/// - Neither `Self` nor [`Self::Repr`] contains interior mutability.
+///
+/// The above basically says that `&Self` can be transmuted to `&Self::Repr`.
+///
+/// [transmutable]: core::mem::transmute
+pub unsafe trait AsRepr: Sized {
+    /// Primitive representation of this type.
+    type Repr;
+
+    /// Convert from `&Self` to [`&Self::Repr`](AsRepr::Repr).
+    #[inline(always)]
+    fn as_repr(this: &Self) -> &Self::Repr {
+        // SAFETY: Per safety requirement of the trait.
+        unsafe { core::mem::transmute(this) }
+    }
+
+    /// Convert from `Self` to [`Self::Repr`].
+    #[inline(always)]
+    fn into_repr(this: Self) -> Self::Repr {
+        // SAFETY: Per safety requirement of the trait.
+        unsafe { transmute(this) }
+    }
+
+    /// Convert from [`Self::Repr`] to `Self`.
+    ///
+    /// # Safety
+    ///
+    /// `repr` must be a valid bit pattern of `Self` and satisfy type-specific invariants of it.
+    ///
+    /// Alternatively, if `repr` is previously obtained using [`Self::into_repr`], and each
+    /// `from_repr_unchecked` should corresponds to a unique `into_repr` call, then it is safe to
+    /// call as well (this means that we're undoing a `into_repr` call getting the exact bytes
+    /// back).
+    ///
+    /// This method makes no guarantee if a `into_repr` corresponds to multiple
+    /// `from_repr_unchecked` (i.e. copies are made), to allow for cases where `Repr` is a pointer
+    /// and the user of the API wants ownership transfer. Users that want the ability to call
+    /// `from_repr_unchecked` after copying can require `Copy` bound explicitly.
+    #[inline(always)]
+    unsafe fn from_repr_unchecked(repr: Self::Repr) -> Self {
+        // SAFETY: Per safety requirement, `repr` is valid repr of `Self`, or it is previously from
+        // `into_repr`, in which case we're undoing the transmute so it is also safe.
+        unsafe { transmute(repr) }
+    }
+}
+
+/// Type that is bi-directionally transmutable with a primitive representation.
+///
+/// # Safety
+///
+/// - [`Self`] must be [transmutable] from [`Self::Repr`].
+///
+/// [transmutable]: core::mem::transmute
+/// [`Self::Repr`]: AsRepr::Repr
+pub unsafe trait AsReprMut: AsRepr {
+    /// Convert from `&mut Self` to [`&mut Self::Repr`](AsRepr::Repr).
+    #[inline(always)]
+    fn as_repr_mut(this: &mut Self) -> &mut Self::Repr {
+        // SAFETY: Per safety requirement of the trait.
+        unsafe { core::mem::transmute(this) }
+    }
+
+    /// Convert from [`Self::Repr`](AsRepr::Repr) to `Self`.
+    #[inline(always)]
+    fn from_repr(repr: Self::Repr) -> Self {
+        // SAFETY: Per safety requirement of the trait.
+        unsafe { transmute(repr) }
+    }
+}
+
+// SAFETY: `bool` has the same size and alignment as `u8`, and Rust guarantees that `bool` has
+// only two valid bit patterns: 0 (false) and 1 (true). Thus `bool` can be transmuted to `u8`.
+// Neither types contain interior mutability.
+unsafe impl AsRepr for bool {
+    type Repr = u8;
+}
+
+// SAFETY: `*mut T` has the same size and alignment with `*const c_void`, and thus `*mut T` is
+// transmutable to `*const c_void`. Neither types contain interior mutability.
+unsafe impl<T> AsRepr for *mut T {
+    type Repr = *const c_void;
+}
+
+// SAFETY: `*mut T` is transmutable from `*const c_void`.
+unsafe impl<T> AsReprMut for *mut T {}
+
+// SAFETY: `*const T` has the same size and alignment with `*const c_void`, and is transmutable to
+// `*const c_void`. Neither types contain interior mutability.
+unsafe impl<T> AsRepr for *const T {
+    type Repr = *const c_void;
+}
+
+// SAFETY: `*const T` is transmutable from `*const c_void`.
+unsafe impl<T> AsReprMut for *const T {}
+
+macro_rules! int_impl {
+    ($($unsigned:ident $signed:ident ,)*) => {$(
+        // SAFETY: $unsigned has the same size and alignment with itself, and is transmutable to
+        // itself. It does not contain interior mutability.
+        unsafe impl AsRepr for $unsigned {
+            type Repr = $unsigned;
+        }
+
+        // SAFETY: $unsigned is transmutable from itself.
+        unsafe impl AsReprMut for $unsigned {}
+
+        // SAFETY: $signed has the same size and alignment with $unsigned, and is transmutable to it
+        // Neither types contain interior mutability.
+        unsafe impl AsRepr for $signed {
+            type Repr = $unsigned;
+        }
+
+        // SAFETY: $signed is transmutable from $unsigned.
+        unsafe impl AsReprMut for $signed {}
+    )*};
+}
+
+int_impl! {
+    u8 i8,
+    u16 i16,
+    u32 i32,
+    u64 i64,
+    // `usize` is not normalized to particular integer for portability.
+    usize isize,
+}

-- 
2.54.0