[macro_util] Add trailing_field_offset! macro

This macro will be used by the derive of KnownLayout for slice DSTs.

Makes progress on #29

Co-authored-by: Jack Wrenn <jswrenn@amazon.com>

src/macro_util.rs

@@ -19,6 +19,10 @@
 
 use core::{marker::PhantomData, mem::ManuallyDrop};
 
+// remove this `cfg` when `size_of_val_raw` is stabilized
+#[cfg(__INTERNAL_USE_ONLY_NIGHLTY_FEATURES_IN_TESTS)]
+use core::ptr::{self, NonNull};
+
 /// A compile-time check that should be one particular value.
 pub trait ShouldBe<const VALUE: bool> {}
 
@@ -61,6 +65,133 @@ impl<T, U> MaxAlignsOf<T, U> {
     }
 }
 
+const _64K: usize = 1 << 16;
+
+// remove this `cfg` when `size_of_val_raw` is stabilized
+#[cfg(__INTERNAL_USE_ONLY_NIGHLTY_FEATURES_IN_TESTS)]
+#[repr(C, align(65536))]
+struct Aligned64kAllocation([u8; _64K]);
+
+/// A pointer to an aligned allocation of size 2^16.
+///
+/// # Safety
+///
+/// `ALIGNED_64K_ALLOCATION` is guaranteed to point to the entirety of an
+/// allocation with size and alignment 2^16, and to have valid provenance.
+// remove this `cfg` when `size_of_val_raw` is stabilized
+#[cfg(__INTERNAL_USE_ONLY_NIGHLTY_FEATURES_IN_TESTS)]
+pub const ALIGNED_64K_ALLOCATION: NonNull<[u8]> = {
+    const REF: &Aligned64kAllocation = &Aligned64kAllocation([0; _64K]);
+    let ptr: *const Aligned64kAllocation = REF;
+    let ptr: *const [u8] = ptr::slice_from_raw_parts(ptr.cast(), _64K);
+    // SAFETY:
+    // - `ptr` is derived from a Rust reference, which is guaranteed to be
+    //   non-null.
+    // - `ptr` is derived from an `&Aligned64kAllocation`, which has size and
+    //   alignment `_64K` as promised. Its length is initialized to `_64K`,
+    //   which means that it refers to the entire allocation.
+    // - `ptr` is derived from a Rust reference, which is guaranteed to have
+    //   valid provenance.
+    //
+    // TODO(#429): Once `NonNull::new_unchecked` docs document that it preserves
+    // provenance, cite those docs.
+    // TODO: Replace this `as` with `ptr.cast_mut()` once our MSRV >= 1.65
+    #[allow(clippy::as_conversions)]
+    unsafe {
+        NonNull::new_unchecked(ptr as *mut _)
+    }
+};
+
+/// Computes the offset of the base of the field `$trailing_field_name` within
+/// the type `$ty`.
+///
+/// `trailing_field_offset!` produces code which is valid in a `const` context.
+// remove this `cfg` when `size_of_val_raw` is stabilized
+#[cfg(__INTERNAL_USE_ONLY_NIGHLTY_FEATURES_IN_TESTS)]
+#[doc(hidden)] // `#[macro_export]` bypasses this module's `#[doc(hidden)]`.
+#[macro_export]
+macro_rules! trailing_field_offset {
+    ($ty:ty, $trailing_field_name:tt) => {{
+        let min_size = {
+            let zero_elems: *const [()] = ::core::ptr::slice_from_raw_parts(
+                ::core::ptr::NonNull::<()>::dangling().as_ptr().cast_const(),
+                0,
+            );
+            // SAFETY:
+            // - If `$ty` is `Sized`, `size_of_val_raw` is always safe to call.
+            // - Otherwise:
+            //   - If `$ty` is not a slice DST, this pointer conversion will
+            //     fail due to "mismatched vtable kinds", and compilation will
+            //     fail.
+            //   - If `$ty` is a slice DST, the safety requirement is that "the
+            //     length of the slice tail must be an initialized integer, and
+            //     the size of the entire value (dynamic tail length +
+            //     statically sized prefix) must fit in isize." The length is
+            //     initialized to 0 above, and Rust guarantees that no type's
+            //     minimum size may overflow `isize`. [1]
+            //
+            // [1] TODO(#429): Citation for this?
+            unsafe {
+                #[allow(clippy::as_conversions)]
+                ::core::mem::size_of_val_raw(zero_elems as *const $ty)
+            }
+        };
+
+        assert!(min_size <= _64K);
+
+        let ptr = ALIGNED_64K_ALLOCATION.as_ptr() as *const $ty;
+
+        // SAFETY:
+        // - Thanks to the preceding `assert!`, we know that the value with zero
+        //   elements fits in `_64K` bytes, and thus in the allocation addressed
+        //   by `ALIGNED_64K_ALLOCATION`. The offset of the trailing field is
+        //   guaranteed to be no larger than this size, so this field projection
+        //   is guaranteed to remain in-bounds of its allocation.
+        // - Because the minimum size is no larger than `_64K` bytes, and
+        //   because an object's size must always be a multiple of its alignment
+        //   [1], we know that `$ty`'s alignment is no larger than `_64K`. The
+        //   allocation addressed by `ALIGNED_64K_ALLOCATION` is guaranteed to
+        //   be aligned to `_64K`, so `ptr` is guaranteed to satisfy `$ty`'s
+        //   alignment.
+        //
+        //   Note that, as of [2], this requirement is technically unnecessary,
+        //   but no harm in guaranteeing it anyway.
+        //
+        // [1] Per https://doc.rust-lang.org/reference/type-layout.html:
+        //
+        //   The size of a value is always a multiple of its alignment.
+        //
+        // [2] https://github.com/rust-lang/reference/pull/1387
+        let field = unsafe { ::core::ptr::addr_of!((*ptr).$trailing_field_name) };
+        // SAFETY:
+        // - Both `ptr` and `field` are derived from the same allocated object.
+        // - By the preceding safety comment, `field` is in bounds of that
+        //   allocated object.
+        // - The distance, in bytes, between `ptr` and `field` is required to be
+        //   a multiple of the size of `u8`, which is trivially true because
+        //   `u8`'s size is 1.
+        // - The distance, in bytes, cannot overflow `isize`. This is guaranteed
+        //   because no allocated object can have a size larger than can fit in
+        //   `isize`. [1]
+        // - The distance being in-bounds cannot rely on wrapping around the
+        //   address space. This is guaranteed because the same is guaranteed of
+        //   allocated objects. [1]
+        //
+        // [1] TODO(#429), TODO(https://github.com/rust-lang/rust/pull/116675):
+        //     Once these are guaranteed in the Reference, cite it.
+        let offset = unsafe { field.cast::<u8>().offset_from(ptr.cast::<u8>()) };
+        // Guaranteed not to be lossy: `field` comes after `ptr`, so the offset
+        // from `ptr` to `field` is guaranteed to be positive.
+        assert!(offset >= 0);
+        Some(
+            #[allow(clippy::as_conversions)]
+            {
+                offset as usize
+            },
+        )
+    }};
+}
+
 /// Does the struct type `$t` have padding?
 ///
 /// `$ts` is the list of the type of every field in `$t`. `$t` must be a
@@ -276,6 +407,32 @@ mod tests {
         assert_t_align_gteq_u_align!(u8, AU64, false);
     }
 
+    // remove this `cfg` when `size_of_val_raw` is stabilized
+    #[cfg(__INTERNAL_USE_ONLY_NIGHLTY_FEATURES_IN_TESTS)]
+    #[test]
+    fn test_trailing_field_offset() {
+        macro_rules! test {
+            (#[$cfg:meta] ($($ts:ty),* ; $trailing_field_ty:ty) => $expect:expr) => {{
+                #[$cfg]
+                struct Test($($ts,)* $trailing_field_ty);
+                assert_eq!(test!(@offset $($ts),* ; $trailing_field_ty), $expect);
+            }};
+            (#[$cfg:meta] $(#[$cfgs:meta])* ($($ts:ty),* ; $trailing_field_ty:ty) => $expect:expr) => {
+                test!(#[$cfg] ($($ts),* ; $trailing_field_ty) => $expect);
+                test!($(#[$cfgs])* ($($ts),* ; $trailing_field_ty) => $expect);
+            };
+            (@offset ; $_trailing:ty) => { trailing_field_offset!(Test, 0) };
+            (@offset $_t:ty ; $_trailing:ty) => { trailing_field_offset!(Test, 1) };
+        }
+
+        test!(#[repr(C)] #[repr(transparent)] #[repr(packed)] #[repr(C)] (; u8) => Some(0));
+        test!(#[repr(C)] #[repr(packed)] #[repr(C)] (u8; u8) => Some(1));
+        test!(#[repr(C)] (u8; [AU64]) => Some(8));
+        test!(#[repr(C)] (u8; Nested<u8, AU64>) => Some(8));
+
+        // TODO: More tests.
+    }
+
     #[test]
     fn test_struct_has_padding() {
         // Test that, for each provided repr, `struct_has_padding!` reports the

src/util.rs

@@ -601,6 +601,25 @@ pub(crate) mod testutil {
             Display::fmt(&self.0, f)
         }
     }
+
+    #[derive(
+        KnownLayout,
+        FromZeroes,
+        FromBytes,
+        Eq,
+        PartialEq,
+        Ord,
+        PartialOrd,
+        Default,
+        Debug,
+        Copy,
+        Clone,
+    )]
+    #[repr(C)]
+    pub(crate) struct Nested<T, U>{
+        _t: T,
+        _u: U,
+    }
 }
 
 #[cfg(test)]