use core::mem::MaybeUninit;

/// Byte buffer with length `SIZE` and aligned to `2^ALIGN` bytes.
pub struct AlignedBuffer<const SIZE: usize, const ALIGN: u8>
where
    Alignment<ALIGN>: ValidAlignment,
{
    /// ZST to force the alignment of this struct.
    ///
    /// "The alignment of the type is at least the maximum alignment of its fields."
    /// https://doc.rust-lang.org/reference/type-layout.html#the-default-representation
    _alignment_marker: <Alignment<ALIGN> as sealed::ValidAlignment>::Marker,

    /// Buffer of bytes to store a value in.
    ///
    /// MaybeUninit is used to allow storing types smaller than `SIZE` and
    /// allows storing types with padding.
    pub data: [MaybeUninit<u8>; SIZE],
}

impl<const SIZE: usize, const ALIGN: u8> AlignedBuffer<SIZE, ALIGN> 
where
    Alignment<ALIGN>: ValidAlignment,
{
    /// Make a new aligned buffer.
    ///
    /// # Invariants
    /// The following invariants can be used by unsafe code.
    ///
    /// - `data` will be properly aligned to `2^ALIGN` bytes.
    /// - `data` can store `0` to `SIZE` bytes.
    ///
    /// Unsafe cannot assume `data` can store a multiple of `2^ALIGN`
    /// bytes more than `SIZE`.
    pub fn new() -> Self {
        // SAFETY: An uninitialized `[MaybeUninit<_>; N]` is valid.
        // Taken from impl of `MaybeUninit::uninit_array()`.
        let data = unsafe { MaybeUninit::<[MaybeUninit<u8>; SIZE]>::uninit().assume_init() };

        Self {
            _alignment_marker: Default::default(),
            data,
        }
    }
}

#[test]
fn forces_alignment() {
    use core::ptr::addr_of;

    // The 2^6 = 64 byte alignment forces the size to be a multiple of 64.
    // Note, the type can only safely store `SIZE` bytes. The extra bytes
    // are padding bytes and aren't safe to use.
    type A = AlignedBuffer<10, 6>;
    assert_eq!(core::mem::align_of::<A>(), 64);
    assert_eq!(core::mem::size_of::<A>(), 64);

    // Check that `data` is actually put at the beginning.
    let x = A::new();
    assert_eq!(addr_of!(x.data), addr_of!(x) as _);

    // The 2^3 = 8 byte alignment forces a round up to 104 bytes from
    // 101 bytes.
    type B = AlignedBuffer<101, 3>;
    assert_eq!(core::mem::align_of::<B>(), 8);
    assert_eq!(core::mem::size_of::<B>(), 104);
    
    // Check that `data` is actually put at the beginning.
    let x = B::new();
    assert_eq!(addr_of!(x.data), addr_of!(x) as _);

    // A 2^0 = 1 byte alignment and zero size results in a ZST.
    // 0 is always a valid multiple of any alignment.
    type C = AlignedBuffer<0, 0>;
    assert_eq!(core::mem::align_of::<C>(), 1);
    assert_eq!(core::mem::size_of::<C>(), 0);

    // Check that `data` is actually put at the beginning.
    let x = C::new();
    assert_eq!(addr_of!(x.data), addr_of!(x) as _);

    // A 2^0 = 1 byte alignment 1 byte long buffer has size 1.
    type D = AlignedBuffer<1, 0>;
    assert_eq!(core::mem::align_of::<D>(), 1);
    assert_eq!(core::mem::size_of::<D>(), 1);

    // Check that `data` is actually put at the beginning.
    let x = D::new();
    assert_eq!(addr_of!(x.data), addr_of!(x) as _);
}

/// Type to test if given `ALIGN` is a valid alignment power.
///
/// Rust only allows alignments from `2^ALIGN` where `ALIGN` is 0 to 29.
pub enum Alignment<const ALIGN: u8> {}

/// Marker for if an alignment power is valid.
///
/// This trait is sealed and cannot be implemented by an external crate.
pub trait ValidAlignment: sealed::ValidAlignment {}

/// Sealed implementation of `ValidAlignment`.
pub(crate) mod sealed {
    use super::Alignment;

    /// A valid alignment has a marker type with the given alignment
    /// applied by a `#[repr(align(ALIGNMENT))]`.
    ///
    /// A valid alignment also has a `ALIGNMENT` constant giving the
    /// alignment amount in bytes.
    pub trait ValidAlignment {
        type Marker: Default;
        const ALIGNMENT: usize;
    }

    macro_rules! impl_alignment_marker {
        ($(($power:literal, $num:literal, $name:ident)),*) => {
            $(
                /// ZST to with given alignment.
                #[repr(align($num))]
                pub struct $name;

                // Implement the public version of the trait.
                impl super::ValidAlignment for Alignment<$power> {}

                // Implement the sealed version.
                impl ValidAlignment for Alignment<$power> {
                    type Marker = $name;

                    const ALIGNMENT: usize = $num;
                }

                // Allow the marker type to be constructable without
                // needing to name it.
                impl Default for $name {
                    fn default() -> Self {
                        Self
                    }
                }
            )*
        }
    }

    // Implement `ValidAlignment` for all valid alignments.
    //
    // "The alignment value must be a power of two from 1 up to 2^29"
    // https://doc.rust-lang.org/reference/type-layout.html#the-alignment-modifiers
    impl_alignment_marker![
        (0, 1, Marker0),
        (1, 2, Marker1),
        (2, 4, Marker2),
        (3, 8, Marker3),
        (4, 16, Marker4),
        (5, 32, Marker5),
        (6, 64, Marker6),
        (7, 128, Marker7),
        (8, 256, Marker8),
        (9, 512, Marker9),
        (10, 1024, Marker10),
        (11, 2048, Marker11),
        (12, 4096, Marker12),
        (13, 8192, Marker13),
        (14, 16384, Marker14),
        (15, 32768, Marker15),
        (16, 65536, Marker16),
        (17, 131072, Marker17),
        (18, 262144, Marker18),
        (19, 524288, Marker19),
        (20, 1048576, Marker20),
        (21, 2097152, Marker21),
        (22, 4194304, Marker22),
        (23, 8388608, Marker23),
        (24, 16777216, Marker24),
        (25, 33554432, Marker25),
        (26, 67108864, Marker26),
        (27, 134217728, Marker27),
        (28, 268435456, Marker28),
        (29, 536870912, Marker29)
    ];
}