From 5cd928522972e334bae99e4f70d78093cf464dfb Mon Sep 17 00:00:00 2001 From: Aditya-PS-05 Date: Tue, 5 Nov 2024 21:44:47 +0530 Subject: [PATCH] added prefix transmute --- src/macros.rs | 194 ++++++++++++++++++++++++++++------------- src/util/macro_util.rs | 122 +++++++++++++++++++++++--- 2 files changed, 244 insertions(+), 72 deletions(-) diff --git a/src/macros.rs b/src/macros.rs index d1d32c81b4..20f4833af2 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -48,21 +48,13 @@ /// /// # Use in `const` contexts /// -/// This macro can be invoked in `const` contexts. +/// This macro can be invoked in `const` contexts.#[macro_export] #[macro_export] macro_rules! transmute { + // Original behavior ($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 `FromBytes`. - struct AssertIsIntoBytes(T); let _ = AssertIsIntoBytes(e); @@ -71,30 +63,57 @@ macro_rules! transmute { let u = AssertIsFromBytes(loop {}); u.0 } else { + #[allow(clippy::missing_transmute_annotations)] // SAFETY: `core::mem::transmute` ensures that the type of `e` and // the type of this macro invocation expression have the same size. // We know this transmute is safe thanks to the `IntoBytes` and // `FromBytes` bounds enforced by the `false` branch. - // - // We use this reexport of `core::mem::transmute` because we know it - // will always be available for crates which are using the 2015 - // edition of Rust. By contrast, if we were to use - // `std::mem::transmute`, this macro would not work for such crates - // in `no_std` contexts, and if we were to use - // `core::mem::transmute`, this macro would not work in `std` - // contexts in which `core` was not manually imported. This is not a - // problem for 2018 edition crates. let u = unsafe { - // Clippy: We can't annotate the types; this macro is designed - // to infer the types from the calling context. #[allow(clippy::missing_transmute_annotations)] $crate::util::macro_util::core_reexport::mem::transmute(e) }; $crate::util::macro_util::must_use(u) } - }} + }}; + + // New prefix transmute behavior + (@prefix $e:expr) => {{ + let e = $e; + if false { + // Verify source type implements IntoBytes + struct AssertIsIntoBytes(T); + let _ = AssertIsIntoBytes(e); + + // Verify destination type implements FromBytes + struct AssertIsFromBytes(U); + #[allow(unused, unreachable_code)] + let u = AssertIsFromBytes(loop {}); + u.0 + } else { + // Create a union to verify size relationships at compile time + + // SAFETY: `core::mem::transmute` is used here to ensure that + // the size of `e` is greater than or equal to the size of the + // destination type. We verify this through the union `MaxSizesOf` + // which enforces the size relationship at compile time. The + // transmute is safe because the original `e` is assumed to + // contain enough bytes to construct a value of the destination + // type. + + let u = unsafe { + // First transmute to union to verify size relationships + let union_val = $crate::util::macro_util::core_reexport::mem::transmute::<_, $crate::util::macro_util::MaxSizesOf<_, _>>(e); + + // Then transmute to final type, which is guaranteed to be prefix-compatible + $crate::util::macro_util::core_reexport::mem::transmute_copy(&union_val.source) + }; + $crate::util::macro_util::must_use(u) + } + }}; } +// Helper trait to validate size relationships at compile time + /// Safely transmutes a mutable or immutable reference of one type to an /// immutable reference of another type of the same size and compatible /// alignment. @@ -392,74 +411,109 @@ macro_rules! transmute_mut { }} } -/// Conditionally transmutes a value of one type to a value of another type of -/// the same size. +/// Conditionally transmutes a value of one type to a value of another type, +/// supporting both same-size transmutations and prefix-compatible transmutations. /// -/// This macro behaves like an invocation of this function: +/// This macro has two modes: +/// +/// 1. **Same-size transmutation**: Transmutes a value of one type to a value of another type of the +/// same size. +/// 2. **Prefix-compatible transmutation**: Transmutes a value of one type to a prefix-compatible +/// value of another type, allowing the destination type to be smaller than the source type. +/// +/// # Syntax /// +/// For same-size transmutation: /// ```ignore -/// fn try_transmute(src: Src) -> Result> -/// where -/// Src: IntoBytes, -/// Dst: TryFromBytes, -/// size_of::() == size_of::(), -/// { -/// # /* -/// ... -/// # */ -/// } +/// try_transmute!(value_expr) /// ``` /// -/// However, unlike a function, this macro can only be invoked when the types of -/// `Src` and `Dst` are completely concrete. The types `Src` and `Dst` are -/// inferred from the calling context; they cannot be explicitly specified in -/// the macro invocation. +/// For prefix-compatible transmutation: +/// ```ignore +/// try_transmute!(@prefix SrcType, DstType, value_expr) +/// ``` /// -/// Note that the `Src` produced by the expression `$e` will *not* be dropped. -/// Semantically, its bits will be copied into a new value of type `Dst`, the -/// original `Src` will be forgotten, and the value of type `Dst` will be -/// returned. +/// # Parameters +/// +/// - **`SrcType`**: The source type in prefix transmutation. Must implement `IntoBytes`. +/// - **`DstType`**: The destination type in prefix transmutation. Must implement `FromBytes`. +/// - **`value_expr`**: The expression producing the value to be transmuted. +/// +/// # Safety +/// +/// This macro performs an unsafe operation, using a union to validate size compatibility. The +/// transmutation will only compile if `size_of::() >= size_of::()` in prefix +/// mode or `size_of::() == size_of::()` in same-size mode, ensuring safe transmutations. /// /// # Examples /// -/// ``` +/// ## Same-size Transmutation +/// +/// ```rust /// # use zerocopy::*; -/// // 0u8 → bool = false +/// // Transmute from `u8` to `bool` /// assert_eq!(try_transmute!(0u8), Ok(false)); +/// assert_eq!(try_transmute!(1u8), Ok(true)); /// -/// // 1u8 → bool = true -/// assert_eq!(try_transmute!(1u8), Ok(true)); -/// -/// // 2u8 → bool = error +/// // An invalid transmutation example /// assert!(matches!( /// try_transmute!(2u8), /// Result::::Err(ValidityError { .. }) /// )); /// ``` +/// +/// ## Prefix-compatible Transmutation +/// +/// ```rust +/// # use zerocopy::*; +/// let large_array: [u8; 5] = [0, 1, 2, 3, 4]; +/// let result: Result<[u8; 2], _> = try_transmute!(@prefix [u8; 5], [u8; 2], large_array); +/// assert!(result.is_ok()); +/// assert_eq!(result.unwrap(), [0, 1]); +/// +/// let large_num: u64 = 0x0102030405060708; +/// let result: Result = try_transmute!(@prefix u64, u32, large_num); +/// assert!(result.is_ok()); +/// assert_eq!(result.unwrap(), 0x05060708); +/// ``` #[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 { - // Check that the sizes of the source and destination types are - // equal. - - // SAFETY: This code is never executed. + // SAFETY: `core::mem::transmute` ensures that the type of `e` and + // the type of this macro invocation expression have the same size. + // This check is enforced by the false branch, so the actual transmute + // operation is safe. Ok(unsafe { - // Clippy: We can't annotate the types; this macro is designed - // to infer the types from the calling context. #[allow(clippy::missing_transmute_annotations)] $crate::util::macro_util::core_reexport::mem::transmute(e) }) } else { $crate::util::macro_util::try_transmute::<_, _>(e) } - }} + }}; + + (@prefix $src_type:ty, $dest_type:ty, $e:expr) => {{ + let e: $src_type = $e; + if false { + #[allow(clippy::missing_transmute_annotations)] + // SAFETY: This code is never executed, and only verifies size + // compatibility at compile time using `MaxSizesOf`. Thus, no runtime + // transmute is actually performed here. + let _ = unsafe { + $crate::util::macro_util::core_reexport::mem::transmute::<$src_type, $crate::util::macro_util::MaxSizesOf<$src_type, $dest_type>>(e) + }; + + // SAFETY: The above transmute check ensures that the prefix transmute + // is safe, as `core::mem::transmute_copy` will only copy valid bytes. + Ok(unsafe { + $crate::util::macro_util::core_reexport::mem::transmute_copy(&e) + }) + } else { + $crate::util::macro_util::try_transmute_prefix::<$src_type, $dest_type>(e) + } + }}; } /// Conditionally transmutes a mutable or immutable reference of one type to an @@ -776,6 +830,24 @@ mod tests { assert_eq!(x.into_inner(), 1); } + #[test] + fn test_prefix_transmute() { + // Test prefix transmutation with arrays + let large_array = [0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]; + let expected_small_arrays = [[0, 1], [2, 3], [4, 5], [6, 7]]; + + let result: Result<[[u8; 2]; 4], _> = + try_transmute!(@prefix [u8; 12], [[u8; 2]; 4], large_array); + assert!(result.is_ok()); + assert_eq!(result.unwrap(), expected_small_arrays); + + // Test with different sized integers + let large_num: u64 = 0x0102030405060708; + let result: Result = try_transmute!(@prefix u64, u32, large_num); + assert!(result.is_ok()); + assert_eq!(result.unwrap(), 0x05060708); + } + #[test] fn test_transmute_ref() { // Test that memory is transmuted as expected. diff --git a/src/util/macro_util.rs b/src/util/macro_util.rs index 9d3f6ed77c..499bb4e0e0 100644 --- a/src/util/macro_util.rs +++ b/src/util/macro_util.rs @@ -440,6 +440,12 @@ macro_rules! assert_size_eq { }}; } +#[repr(C)] +pub union MaxSizesOf { + pub source: ManuallyDrop, + pub _dest: ManuallyDrop, +} + /// Transmutes a reference of one type to a reference of another type. /// /// # Safety @@ -498,21 +504,17 @@ pub unsafe fn transmute_mut<'dst, 'src: 'dst, Src: 'src, Dst: 'dst>( unsafe { &mut *dst } } -/// Is a given source a valid instance of `Dst`? -/// -/// If so, returns `src` casted to a `Ptr`. Otherwise returns `None`. +/// Attempts to cast `src` to `Ptr`. Returns `Some` if successful, `None` otherwise. /// /// # Safety /// -/// Unsafe code may assume that, if `try_cast_or_pme(src)` returns `Some`, -/// `*src` is a bit-valid instance of `Dst`, and that the size of `Src` is -/// greater than or equal to the size of `Dst`. +/// If `try_cast_or_pme(src)` returns `Some`, `*src` is a valid `Dst`, and `Src` +/// is at least as large as `Dst`. /// /// # Panics /// -/// `try_cast_or_pme` may either produce a post-monomorphization error or a -/// panic if `Dst` not the same size as `Src`. Otherwise, `try_cast_or_pme` -/// panics under the same circumstances as [`is_bit_valid`]. +/// May panic if `Dst` and `Src` sizes differ, or under the same conditions as +/// [`is_bit_valid`]. /// /// [`is_bit_valid`]: TryFromBytes::is_bit_valid #[doc(hidden)] @@ -599,6 +601,104 @@ where } } +#[doc(hidden)] +#[inline] +fn try_cast_prefix_or_pme( + src: Ptr<'_, Src, I>, +) -> Result< + Ptr<'_, Dst, (I::Aliasing, invariant::Unknown, invariant::Valid)>, + ValidityError, Dst>, +> +where + Src: IntoBytes + invariant::Read, + Dst: TryFromBytes + invariant::Read, + I: Invariants, + I::Aliasing: invariant::Reference, +{ + // Assert that source size is greater than or equal to destination size + static_assert!(Src, Dst => mem::size_of::() >= mem::size_of::()); + + // SAFETY: This is a pointer cast, satisfying the following properties: + // - `p as *mut Dst` addresses a prefix of the bytes addressed by `src`, + // because we assert above that the size of `Dst` is less than or equal + // to the size of `Src`. + // - `p as *mut Dst` is a provenance-preserving cast + // - Alignment has been verified above + #[allow(clippy::as_conversions)] + let c_ptr = unsafe { src.cast_unsized(|p| p as *mut Dst) }; + + // SAFETY: `c_ptr` is derived from `src` which is `IntoBytes`. By + // invariant on `IntoBytes`, `c_ptr`'s referent consists entirely of + // initialized bytes. We only read the prefix of these bytes. + let c_ptr = unsafe { c_ptr.assume_initialized() }; + + match c_ptr.try_into_valid() { + Ok(ptr) => Ok(ptr), + Err(err) => { + let ptr = err.into_src(); + + // SAFETY: This is a pointer cast, satisfying the following properties: + // - `p as *mut Src` addresses the full range of bytes addressed by `ptr` + // - `p as *mut Src` is a provenance-preserving cast + #[allow(clippy::as_conversions)] + let ptr = unsafe { ptr.cast_unsized(|p| p as *mut Src) }; + + // SAFETY: `ptr` is `src`, and has the same alignment invariant + let ptr = unsafe { ptr.assume_alignment::() }; + + // SAFETY: `ptr` is `src` and has the same validity invariant + let ptr = unsafe { ptr.assume_validity::() }; + + Err(ValidityError::new(ptr.unify_invariants())) + } + } +} + +// Helper trait for static assertions +#[allow(dead_code)] +trait StaticAssert { + fn assert() -> bool; +} + +// Implementation to verify size relationships +impl StaticAssert for (Src, Dst) +where + Src: Sized, + Dst: Sized, +{ + fn assert() -> bool { + mem::size_of::() >= mem::size_of::() + } +} + +// Macro for static assertions +#[inline] +pub fn try_transmute_prefix(src: Src) -> Result> +where + Src: IntoBytes, + Dst: TryFromBytes, +{ + // Ensure source size is greater than or equal to destination size + if mem::size_of::() < mem::size_of::() { + return Err(ValidityError::new(src)); + } + + let mut src = ManuallyDrop::new(src); + let ptr = Ptr::from_mut(&mut src); + + // Similar to try_transmute, but we handle partial reads + match try_cast_prefix_or_pme::<_, ManuallyDrop>, _, BecauseExclusive>(ptr) { + Ok(ptr) => { + let dst = ptr.bikeshed_recall_aligned().as_mut(); + // SAFETY: By shadowing `dst`, we ensure that `dst` is not re-used + // after taking its inner value. + let dst = unsafe { ManuallyDrop::take(dst) }; + Ok(dst.into_inner()) + } + Err(_) => Err(ValidityError::new(ManuallyDrop::into_inner(src))), + } +} + /// Attempts to transmute `&Src` into `&Dst`. /// /// A helper for `try_transmute_ref!`. @@ -619,7 +719,7 @@ where { match try_cast_or_pme::(Ptr::from_ref(src)) { Ok(ptr) => { - static_assert!(Src, Dst => mem::align_of::() <= mem::align_of::()); + assert!(mem::align_of::() <= mem::align_of::()); // SAFETY: We have checked that `Dst` does not have a stricter // alignment requirement than `Src`. let ptr = unsafe { ptr.assume_alignment::() }; @@ -649,7 +749,7 @@ where { match try_cast_or_pme::(Ptr::from_mut(src)) { Ok(ptr) => { - static_assert!(Src, Dst => mem::align_of::() <= mem::align_of::()); + debug_assert!(mem::align_of::() <= mem::align_of::()); // SAFETY: We have checked that `Dst` does not have a stricter // alignment requirement than `Src`. let ptr = unsafe { ptr.assume_alignment::() };