const-generic-alignment 0.3.0

#![no_std]
#![cfg_attr(feature = "freeze", feature(freeze))]
#![cfg_attr(feature = "unstable_docs", feature(doc_auto_cfg))]

#[cfg(feature = "freeze")]
use core::marker::Freeze;
use core::{
    borrow::{Borrow, BorrowMut},
    fmt,
    hash::Hash,
    mem::MaybeUninit,
    panic::{RefUnwindSafe, UnwindSafe},
};

/// A wrapper type with the specified alignment.
///
/// This type implements most relevant traits if the wrapped type does,
/// including [`Default`], [`Ord`], [`Copy`], [`Send`], [`Sync`], [`Unpin`],
/// etc.
///
/// With the `bytemuck` optional feature enabled, this type implements the
#[cfg_attr(feature = "bytemuck", doc = "[`bytemuck::AnyBitPattern`]")]
#[cfg_attr(
    not(feature = "bytemuck"),
    doc = "[`bytemuck::AnyBitPattern`](https://docs.rs/bytemuck/latest/bytemuck/trait.AnyBitPattern.html)"
)]
/// trait if the wrapped type does. (It cannot in general implement `Pod`,
/// because there may be padding.)
///
/// With the `freeze` optional feature enabled, this type implements the
/// unstable [`core::marker::Freeze`] trait if the wrapped type does.
///
/// # Layout
///
/// The wrapped field `inner` is at offset `0`, and the size and alignment of
/// `Self` are minimal to meet the wrapped field's size and alignment, and to
/// meet the requested alignment, under Rust's normal `size % align == 0` rule.
///
/// This is *not* `repr(transparent)`.
#[derive(Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(C)]
pub struct Aligned<T: ?Sized, const ALIGNMENT: usize>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    pub align: AlignedZst<ALIGNMENT>,
    pub inner: T,
}

impl<T: ?Sized + Hash, const ALIGNMENT: usize> Hash for Aligned<T, ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    #[inline]
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
        T::hash(&self.inner, state)
    }
}

impl<T: ?Sized, const ALIGNMENT: usize> AsRef<T> for Aligned<T, ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    #[inline]
    fn as_ref(&self) -> &T {
        &self.inner
    }
}

impl<T: ?Sized, const ALIGNMENT: usize> AsMut<T> for Aligned<T, ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    #[inline]
    fn as_mut(&mut self) -> &mut T {
        &mut self.inner
    }
}

impl<T: ?Sized, const ALIGNMENT: usize> Borrow<T> for Aligned<T, ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    #[inline]
    fn borrow(&self) -> &T {
        &self.inner
    }
}

impl<T: ?Sized, const ALIGNMENT: usize> BorrowMut<T> for Aligned<T, ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    #[inline]
    fn borrow_mut(&mut self) -> &mut T {
        &mut self.inner
    }
}

impl<T, const ALIGNMENT: usize> Aligned<T, ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    /// Creates a new `Aligned` wrapper containing the given value.
    pub const fn new(value: T) -> Self {
        Self { align: AlignedZst::new(), inner: value }
    }

    /// Consumes this `Aligned` wrapper, returning the wrapped value.
    pub fn into_inner(self) -> T {
        let Self { inner, .. } = self;
        inner
    }
}

// Safety: The only fields are an inhabited ZST and `T`, which are both
// `Zeroable`. Padding is not relevant for `Zeroable`.
#[cfg(feature = "bytemuck")]
unsafe impl<T, const ALIGNMENT: usize> bytemuck::Zeroable
    for Aligned<T, ALIGNMENT>
where
    T: bytemuck::Zeroable,
    Alignment<ALIGNMENT>: SupportedAlignment,
{
}

// Safety: The only fields are an inhabited ZST and `T`, which are both
// `AnyBitPattern`. Padding is not relevant for `AnyBitPattern`.
#[cfg(feature = "bytemuck")]
unsafe impl<T, const ALIGNMENT: usize> bytemuck::AnyBitPattern
    for Aligned<T, ALIGNMENT>
where
    T: bytemuck::AnyBitPattern,
    Alignment<ALIGNMENT>: SupportedAlignment,
{
}

/// A Zero-sized type with the specified alignment.
///
/// To facilitate deriving traits on types containing this type, this type
/// implements most relevant traits, including [`Default`], [`Ord`], [`Copy`],
/// [`Send`], [`Sync`], [`Unpin`], etc.
///
/// With the `bytemuck` optional feature enabled, this type implements the
#[cfg_attr(feature = "bytemuck", doc = "[`bytemuck::Pod`]")]
#[cfg_attr(
    not(feature = "bytemuck"),
    doc = "[`bytemuck::Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)"
)]
/// trait.
///
/// With the `freeze` optional feature enabled, this type implements the
/// unstable [`core::marker::Freeze`] trait.
#[derive(Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct AlignedZst<const ALIGNMENT: usize>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    inner: <Alignment<ALIGNMENT> as SupportedAlignment>::AlignedZst,
}

impl<const ALIGNMENT: usize> AlignedZst<ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    /// Creates a new `AlignedZst`.
    ///
    /// This does the same thing as `<Self as Default>::default()`, but is
    /// usable in `const`.
    pub const fn new() -> Self {
        // Safety: This is an inhabited ZST.
        unsafe { MaybeUninit::uninit().assume_init() }
    }
}

impl<const ALIGNMENT: usize> fmt::Debug for AlignedZst<ALIGNMENT>
where
    Alignment<ALIGNMENT>: SupportedAlignment,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "AlignedZst<{}> {{ .. }}", ALIGNMENT)
    }
}

// Safety: This is an inhabited ZST.
#[cfg(feature = "bytemuck")]
unsafe impl<const ALIGNMENT: usize> bytemuck::Zeroable for AlignedZst<ALIGNMENT> where
    Alignment<ALIGNMENT>: SupportedAlignment
{
}

// Safety: This is an inhabited ZST.
#[cfg(feature = "bytemuck")]
unsafe impl<const ALIGNMENT: usize> bytemuck::Pod for AlignedZst<ALIGNMENT> where
    Alignment<ALIGNMENT>: SupportedAlignment
{
}

mod sealed {
    pub trait Sealed {}
}

impl<const N: usize> sealed::Sealed for Alignment<N> {}

/// Marker type to specify the alignment of [`AlignedZst`] in the type system.
///
/// For valid alignments `N`, `Alignment<N>` implements [`SupportedAlignment`].
#[non_exhaustive]
pub struct Alignment<const N: usize>;

/// Statically guarantees that an alignment is supported
///
/// This trait is *sealed*: the list of implementors below is total. Users do
/// not have the ability to mark additional `Alignment<N>` values as supported.
/// Only alignments supported by rustc/LLVM are constructable, i.e. powers of
/// two from 1 to `2^29` (or `2^15` on 16-bit targets).
pub unsafe trait SupportedAlignment: sealed::Sealed {
    cfg_if::cfg_if! {
        if #[cfg(feature = "freeze")] {
            /// Safety: This must be an inhabited ZST with the specified alignment.
            #[doc(hidden)]
            type AlignedZst: Sized + Default + Copy + Ord + Hash + Send + Sync + Unpin + UnwindSafe + RefUnwindSafe + Freeze + 'static;
        } else {
            /// Safety: This must be an inhabited ZST with the specified alignment.
            #[doc(hidden)]
            type AlignedZst: Sized + Default + Copy + Ord + Hash + Send + Sync + Unpin + UnwindSafe + RefUnwindSafe + 'static;
        }
    }
}

macro_rules! make_zsts {
    ( $($typename:ident = $alignment:literal;)* ) => { $(
        /// Helper type used only as a field in [`AlignedZst`].
        ///
        #[doc = concat!("Is a ZST with alignment ", stringify!($alignment), ".")]
        #[doc(hidden)]
        #[derive(Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
        #[repr(align($alignment))]
        pub struct $typename;

        unsafe impl SupportedAlignment for Alignment<$alignment> {
            type AlignedZst = $typename;
        }
     )* };
}

make_zsts! {
    Align1 = 1;
    Align2 = 2;
    Align4 = 4;
    Align8 = 8;
    Align16 = 16;
    Align32 = 32;
    Align64 = 64;
    Align128 = 128;
    Align256 = 256;
    Align512 = 512;
    Align1024 = 1024;
    Align2048 = 2048;
    Align4096 = 4096;
    Align8192 = 8192;
    Align16384 = 16384;
    Align32768 = 32768;
}
#[cfg(not(target_pointer_width = "16"))]
make_zsts! {
    Align65536 = 65536;
    Align131072 = 131072;
    Align262144 = 262144;
    Align524288 = 524288;
    Align1048576 = 1048576;
    Align2097152 = 2097152;
    Align4194304 = 4194304;
    Align8388608 = 8388608;
    Align16777216 = 16777216;
    Align33554432 = 33554432;
    Align67108864 = 67108864;
    Align134217728 = 134217728;
    Align268435456 = 268435456;
    Align536870912 = 536870912;
}

#[cfg(test)]
mod tests {
    extern crate std;

    use crate::{Aligned, AlignedZst, Alignment, SupportedAlignment};

    #[test]
    fn it_works() {
        #[derive(Default, Clone, Copy)]
        struct AlignedThing<const N: usize, T>
        where
            Alignment<N>: SupportedAlignment,
        {
            _aligned: AlignedZst<N>,
            value: T,
        }

        assert_eq!(core::mem::align_of::<AlignedThing<64, u8>>(), 64);
        assert_eq!(
            core::mem::align_of::<AlignedThing<4, u64>>(),
            core::cmp::max(4, core::mem::align_of::<u64>())
        );
        assert_eq!(
            core::mem::align_of::<AlignedThing<536870912, u8>>(),
            536870912
        );

        let x: [AlignedThing<64, u8>; 2] =
            [AlignedThing { value: 42, _aligned: Default::default() }; 2];
        assert_eq!(
            (&x[1].value as *const _ as usize)
                - (&x[0].value as *const _ as usize),
            64
        );
    }

    #[test]
    fn wrapper_works() {
        assert_eq!(core::mem::align_of::<Aligned<u8, 64>>(), 64);
        assert_eq!(
            core::mem::align_of::<Aligned<u64, 4>>(),
            core::cmp::max(4, core::mem::align_of::<u64>())
        );
        assert_eq!(core::mem::align_of::<Aligned<u8, 536870912>>(), 536870912);

        let x: [Aligned<u8, 64>; 2] = [Aligned::new(42); 2];
        assert_eq!(
            (&x[1].inner as *const _ as usize)
                - (&x[0].inner as *const _ as usize),
            64
        );
    }
}