From e6dc8ee2a12dec3a404b028709bbfc72954d4206 Mon Sep 17 00:00:00 2001
From: Jack Wrenn <jack@wrenn.fyi>
Date: Mon, 4 Mar 2024 22:20:44 +0000
Subject: [PATCH] Implement `try_transmute!`

Closes #1013.

Makes progress towards #5.
---
 src/lib.rs        | 102 ++++++++++++++++++++++++++++++++++++++++++++++
 src/macro_util.rs |  19 +++++++++
 2 files changed, 121 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index c6b55fc0586..b362fd68a5d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1521,6 +1521,47 @@ pub unsafe trait TryFromBytes {
     #[doc(hidden)]
     fn is_bit_valid<A: invariant::at_least::Shared>(candidate: Maybe<'_, Self, A>) -> bool;
 
+    /// Is a given source value a valid instance of `Self`?
+    ///
+    /// # Safety
+    ///
+    /// Unsafe code may assume that, if `is_src_valid(src)` returns true, `*src`
+    /// is a bit-valid instance of `Self`.
+    ///
+    /// # Panics
+    ///
+    /// `is_src_valid` panics under exactly the same circumstances as
+    /// [`is_bit_valid`].
+    ///
+    /// [`is_bit_valid`]: TryFromBytes::is_bit_valid
+    #[doc(hidden)]
+    #[inline]
+    fn is_src_valid<Src>(src: &mut Src) -> bool
+    where
+        Src: IntoBytes,
+        Self: Sized,
+    {
+        let c_ptr = crate::Ptr::from_mut(src);
+
+        if mem::size_of::<Src>() > mem::size_of::<Self>() {
+            return false;
+        }
+
+        // SAFETY:
+        // - We just validated that `size_of::<Src>()` <= `size_of::<Self>()`
+        // - `c_ptr` is exclusively aliased, so we do not need to reason about
+        //   `UnsafeCell`
+        #[allow(clippy::as_conversions)]
+        let c_ptr = unsafe { c_ptr.cast_unsized(|p| p as *mut Self) };
+
+        // SAFETY: `c_ptr` is derived from `src` which is `IntoBytes`. By
+        // invariant on `IntoByte`s, `c_ptr`'s referent consists entirely of
+        // initialized bytes.
+        let c_ptr = unsafe { c_ptr.assume_validity::<crate::invariant::Initialized>() };
+
+        Self::is_bit_valid(c_ptr)
+    }
+
     /// Attempts to interpret a byte slice as a `Self`.
     ///
     /// `try_from_ref` validates that `bytes` contains a valid `Self`, and that
@@ -4732,6 +4773,67 @@ macro_rules! transmute_mut {
     }}
 }
 
+/// Conditionally transmutes a value of one type to a value of another type of
+/// the same size.
+///
+/// This macro consumes an expression of type  `Src` and produces a `Result<Dst,
+/// Src>`. It produces `Ok(dst)` if `src` is a bit-valid instance of `Dst`;
+/// otherwise it produces `Err(src)`.
+///
+/// The expression `$e` must have a concrete type, `T`, which implements
+/// [`IntoBytes`]. The `try_transmute!` expression must also have a concrete
+/// type, `U` (`U` is inferred from the calling context), and `U` must implement
+/// [`TryFromBytes`].
+///
+/// Note that the `T` produced by the expression `$e` will *not* be dropped.
+/// Semantically, its bits will be copied into a new value of type `U`, the
+/// original `T` will be forgotten, and the value of type `U` will be returned.
+///
+/// # Examples
+///
+/// ```
+/// # use zerocopy::try_transmute;
+/// assert_eq!(try_transmute!(0u8), Ok(false));
+/// assert_eq!(try_transmute!(1u8), Ok(true));
+/// assert_eq!(try_transmute!(255u8), Result::<bool, _>::Err(255));
+/// ```
+#[macro_export]
+macro_rules! try_transmute {
+    ($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. `core::mem::transmute` uses compiler magic
+        // to enforce this so long as the types are concrete.
+
+        let e = $e;
+        if false {
+            // This branch, though never taken, ensures that the type of `e` is
+            // `IntoBytes` and that the type of this macro invocation expression
+            // is `TryFromBytes`.
+
+            struct AssertIsIntoBytes<T: $crate::IntoBytes>(T);
+            let _ = AssertIsIntoBytes(e);
+
+            struct AssertIsTryFromBytes<U: $crate::TryFromBytes>(U);
+            #[allow(unused, unreachable_code)]
+            let u = AssertIsTryFromBytes(loop {});
+            Ok(u.0)
+        } else if false {
+            // Check that the sizes of the source and destination types are
+            // equal.
+
+            // SAFETY: This code is never executed.
+            Ok(unsafe {
+                // Clippy: It's okay to transmute a type to itself.
+                #[allow(clippy::useless_transmute)]
+                $crate::macro_util::core_reexport::mem::transmute(e)
+            })
+        } else {
+            $crate::macro_util::try_transmute::<_, _>(e)
+        }
+    }}
+}
+
 /// Includes a file and safely transmutes it to a value of an arbitrary type.
 ///
 /// The file will be included as a byte array, `[u8; N]`, which will be
diff --git a/src/macro_util.rs b/src/macro_util.rs
index 3744375ef3b..064d2f21964 100644
--- a/src/macro_util.rs
+++ b/src/macro_util.rs
@@ -404,6 +404,25 @@ pub unsafe fn transmute_mut<'dst, 'src: 'dst, Src: 'src, Dst: 'dst>(
     unsafe { &mut *dst }
 }
 
+/// A helper for `try_transmute!`.
+#[inline(always)]
+pub fn try_transmute<Src, Dst>(mut src: Src) -> Result<Dst, Src>
+where
+    Src: crate::IntoBytes,
+    Dst: crate::TryFromBytes,
+{
+    if !Dst::is_src_valid(&mut src) {
+        return Err(src);
+    }
+
+    let src = ManuallyDrop::new(src);
+
+    // SAFETY:
+    // - `is_src_valid` validates that `src` is a bit-valid instance of `Dst`
+    // - `src` is `ManuallyDrop`
+    Ok(unsafe { core::mem::transmute_copy(&*src) })
+}
+
 // NOTE: We can't change this to a `pub use core as core_reexport` until [1] is
 // fixed or we update to a semver-breaking version (as of this writing, 0.8.0)
 // on the `main` branch.