[PATCH WIP 2/5] rust: transmute: Cleanup + Fixes

[Matthew Maurer <mmaurer@google.com>]

This change is not intended to go upstream as-is, the original
`FromBytes`/`AsBytes` is being litigated on the list. This just fixes it
up so that I can use it for this example.

Signed-off-by: Matthew Maurer <mmaurer@google.com>
---
 rust/kernel/lib.rs       |   1 +
 rust/kernel/transmute.rs | 126 ++++++++++++++++++++++++++++-------------------
 2 files changed, 77 insertions(+), 50 deletions(-)

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 045f1088938cf646519edea2102439402fb27660..0461f25cb5aee797d25153a2004d63b6b41f4ae3 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -18,6 +18,7 @@
 //
 // Stable since Rust 1.79.0.
 #![feature(inline_const)]
+#![feature(pointer_is_aligned)]
 //
 // Stable since Rust 1.81.0.
 #![feature(lint_reasons)]
diff --git a/rust/kernel/transmute.rs b/rust/kernel/transmute.rs
index ba21fe49e4f07808c0a43f16461b535fadc033f1..452b1cfb1dbecfdddec7bb59204f7290ae5040af 100644
--- a/rust/kernel/transmute.rs
+++ b/rust/kernel/transmute.rs
@@ -46,63 +46,71 @@ fn from_bytes_mut(bytes: &mut [u8]) -> Option<&mut Self>
         Self: AsBytes;
 }
 
-/// Provide an auto-implementation of FromBytes's methods for all
-/// sized types, if you need an implementation for your type use this instead.
-///
-/// # Safety
-///
-/// All bit-patterns must be valid for this type. This type must not have interior mutability.
-pub unsafe trait FromBytesSized: Sized {}
+/// Helper for implementing `from_bytes` for sized types.
+pub fn sized_from_bytes<T: FromBytes>(bytes: &[u8]) -> Option<&T> {
+    if bytes.len() == core::mem::size_of::<T>() {
+        let slice_ptr = bytes.as_ptr().cast::<T>();
+        if !slice_ptr.is_aligned() {
+            None
+        } else {
+            // SAFETY:
+            // * T is FromBytes, so anything in the bytes array is a valid bit pattern
+            // * The pointer is aligned
+            // * The pointer points to a region of the appropriate size
+            unsafe { Some(&*slice_ptr) }
+        }
+    } else {
+        None
+    }
+}
 
-macro_rules! impl_frombytessized {
+/// Helper for implementing `from_bytes_mut` for sized types.
+pub fn sized_from_bytes_mut<T: FromBytes + AsBytes>(bytes: &mut [u8]) -> Option<&mut T> {
+    if bytes.len() == core::mem::size_of::<T>() {
+        let slice_ptr = bytes.as_mut_ptr().cast::<T>();
+        if !slice_ptr.is_aligned() {
+            None
+        } else {
+            // SAFETY:
+            // * T is FromBytes, so anything in the bytes array is a valid bit pattern
+            // * T is AsBytes, so mutating T will not expose padding to the byte array
+            // * The pointer is aligned
+            // * The pointer points to a region of the appropriate size
+            unsafe { Some(&mut *slice_ptr) }
+        }
+    } else {
+        None
+    }
+}
+
+macro_rules! impl_from_bytes{
     ($($({$($generics:tt)*})? $t:ty, )*) => {
         // SAFETY: Safety comments written in the macro invocation.
-        $(unsafe impl$($($generics)*)? FromBytesSized for $t {})*
+        $(unsafe impl$($($generics)*)? FromBytes for $t {
+            fn from_bytes(bytes: &[u8]) -> Option<&$t> {
+                sized_from_bytes(bytes)
+            }
+
+            fn from_bytes_mut(bytes: &mut [u8]) -> Option<&mut $t>
+            where
+            Self: AsBytes,
+            {
+                sized_from_bytes_mut(bytes)
+            }
+        })*
     };
 }
 
-impl_frombytessized! {
+impl_from_bytes! {
     // SAFETY: All bit patterns are acceptable values of the types below.
+    // Checking the pointer size makes this operation safe and it's necessary
+    // to dereference to get the value and return it as a reference to `Self`.
     u8, u16, u32, u64, usize,
     i8, i16, i32, i64, isize,
 
     // SAFETY: If all bit patterns are acceptable for individual values in an array, then all bit
     // patterns are also acceptable for arrays of that type.
-    {<T: FromBytesSized, const N: usize>} [T; N],
-}
-
-// SAFETY: The `FromBytesSized` implementation guarantees that all bit
-// patterns are acceptable values of the types and in array case if
-// all bit patterns are acceptable for individual values in an array,
-// then all bit patterns are also acceptable for arrays of that type.
-unsafe impl<T> FromBytes for T
-where
-    T: FromBytesSized,
-{
-    fn from_bytes(bytes: &[u8]) -> Option<&Self> {
-        let slice_ptr = bytes.as_ptr().cast::<T>();
-        let size = ::core::mem::size_of::<T>();
-        if bytes.len() == size && slice_ptr.is_aligned() {
-            // SAFETY: Since the code checks the size and alignment, the slice is valid.
-            unsafe { Some(&*slice_ptr) }
-        } else {
-            None
-        }
-    }
-
-    fn from_bytes_mut(bytes: &mut [u8]) -> Option<&mut Self>
-    where
-        Self: AsBytes,
-    {
-        let slice_ptr = bytes.as_mut_ptr().cast::<T>();
-        let size = ::core::mem::size_of::<T>();
-        if bytes.len() == size && slice_ptr.is_aligned() {
-            // SAFETY: Since the code checks the size and alignment, the slice is valid.
-            unsafe { Some(&mut *slice_ptr) }
-        } else {
-            None
-        }
-    }
+    {<T: FromBytes, const N: usize>} [T; N],
 }
 
 // SAFETY: If all bit patterns are acceptable for individual values in an array, then all bit
@@ -110,7 +118,7 @@ fn from_bytes_mut(bytes: &mut [u8]) -> Option<&mut Self>
 unsafe impl<T: FromBytes> FromBytes for [T] {
     fn from_bytes(bytes: &[u8]) -> Option<&Self> {
         let size = ::core::mem::size_of::<T>();
-        build_assert!(size == 0, "Can't create a slice with zero elements");
+        build_assert!(size != 0, "Can't create a slice with zero-sized elements");
         let slice_ptr = bytes.as_ptr().cast::<T>();
         if bytes.len() % size == 0 && slice_ptr.is_aligned() {
             // SAFETY: Since the number of elements is different from
@@ -126,7 +134,7 @@ fn from_bytes_mut(bytes: &mut [u8]) -> Option<&mut Self>
         Self: AsBytes,
     {
         let size = ::core::mem::size_of::<T>();
-        build_assert!(size == 0, "Can't create a slice with zero elements");
+        build_assert!(size != 0, "Can't create a slice with zero-sized elements");
         let slice_ptr = bytes.as_mut_ptr().cast::<T>();
         if bytes.len() % size == 0 && slice_ptr.is_aligned() {
             // SAFETY: Since the number of elements is different from
@@ -158,16 +166,34 @@ fn from_bytes_mut(bytes: &mut [u8]) -> Option<&mut Self>
 ///
 /// Values of this type may not contain any uninitialized bytes. This type must not have interior
 /// mutability.
-pub unsafe trait AsBytes {}
+pub unsafe trait AsBytes {
+    /// View data structure as a buffer
+    fn as_bytes(&self) -> &[u8] {
+        let len = core::mem::size_of_val(self);
+        // SAFETY: By unsafe trait impl precondition, there's no interior mutability and no
+        // uninitialized bytes.
+        unsafe { core::slice::from_raw_parts(core::ptr::from_ref(self).cast::<u8>(), len) }
+    }
+    /// View data structure as a mutable buffer
+    fn as_mut_bytes(&mut self) -> &mut [u8]
+    where
+        Self: FromBytes,
+    {
+        let len = core::mem::size_of_val(self);
+        // SAFETY: By unsafe trait impl precondition, there's no interior mutability, and no
+        // unititialized bytes. By FromBytes trait impl precondition, all bit patterns are valid.
+        unsafe { core::slice::from_raw_parts_mut(core::ptr::from_mut(self).cast::<u8>(), len) }
+    }
+}
 
-macro_rules! impl_asbytes {
+macro_rules! impl_as_bytes {
     ($($({$($generics:tt)*})? $t:ty, )*) => {
         // SAFETY: Safety comments written in the macro invocation.
         $(unsafe impl$($($generics)*)? AsBytes for $t {})*
     };
 }
 
-impl_asbytes! {
+impl_as_bytes! {
     // SAFETY: Instances of the following types have no uninitialized portions.
     u8, u16, u32, u64, usize,
     i8, i16, i32, i64, isize,

-- 
2.51.0.rc1.167.g924127e9c0-goog