#![warn(missing_docs)]
#![warn(missing_doc_code_examples)]
#![warn(rust_2018_idioms)]

//! Allocate heap memory with user-specified alignment.

/// Error type for custom errors of `AlignedBox`.
#[derive(Debug)]
pub enum AlignedBoxError {
    /// Too many elements. The requested size exceeds the maximum size for a slice.
    TooManyElements,
    /// Memory allocation failed. This is likely caused by an out of memory situation.
    OutOfMemory,
    /// Zero byte allocation are currently not supported by AlignedBox.
    ZeroAlloc,
}

impl std::error::Error for AlignedBoxError {}

impl std::fmt::Display for AlignedBoxError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AlignedBoxError::TooManyElements => write!(f, "Too many elements for a slice."),
            AlignedBoxError::OutOfMemory => write!(f, "Memory allocation failed. Out of memory?"),
            AlignedBoxError::ZeroAlloc => write!(f, "Zero byte allocations not supported."),
        }
    }
}

/// A wrapper around `std::boxed::Box` which allows allocating aligned heap memory. An instance of
/// `AlignedBox<T>` consists of a `Box<T>` and the `std::alloc::Layout` that has been used to
/// allocate the referenced memory.
pub struct AlignedBox<T: ?Sized> {
    // container is not a Box<T> but an Option<Box<T>> for purely technical reasons:
    // When drop(&mut self) is called, we need to be able to get the raw pointer from the Box.
    // Therefore we need to be able to take ownership of the Box. Option::take() allows that.
    // The Option is Some as long as the AlignedBox exist. It is only set to None during drop()
    // and into_raw_parts(). In both cases, the AlignedBox is destroyed directly afterwards.
    container: std::option::Option<std::boxed::Box<T>>,
    layout: std::alloc::Layout,
}

impl<T: ?Sized> std::ops::Deref for AlignedBox<T> {
    type Target = T;

    fn deref(&self) -> &T {
        // self.container is always Some, so we can just unwrap
        return self.container.as_deref().unwrap();
    }
}

impl<T: ?Sized> std::ops::DerefMut for AlignedBox<T> {
    fn deref_mut(&mut self) -> &mut T {
        // self.container is always Some, so we can just unwrap
        return self.container.as_deref_mut().unwrap();
    }
}

impl<T: ?Sized> Drop for AlignedBox<T> {
    fn drop(&mut self) {
        // self.container is always Some, so we can just unwrap
        let container = self.container.take().unwrap();
        unsafe {
            std::alloc::dealloc(std::boxed::Box::into_raw(container) as *mut u8, self.layout);
        }
    }
}

impl<T: Clone + ?Sized> Clone for AlignedBox<T> {
    fn clone(&self) -> Self {
        // layout is certainly valid as it has already been used to create self
        let ptr = unsafe { std::alloc::alloc(self.layout) as *mut T };
        if ptr.is_null() {
            panic!("Error when allocating memory for a clone of AlignedBox");
        }

        // SAFETY: The pointer is non-null, refers to properly sized and aligned memory and it is
        // consumed such that it cannot be used from anywhere outside the Box.
        let mut b = unsafe { AlignedBox::<T>::from_raw_parts(ptr, self.layout) };

        // *b is not a valid instance of T but uninitialized memory. We have to write to it without
        // dropping the old (invalid) value. Also the original value must not be dropped.
        // self.container is always Some, so we can just unwrap.
        // SAFETY: *b points to valid and properly aligned memory and clone() also provides us with
        // a valid value.
        unsafe { std::ptr::write(&mut *b, *self.container.clone().unwrap()) };

        b
    }
}

impl<T: ?Sized> AlignedBox<T> {
    /// Decompose the `AlignedBox` into a raw pointer and the layout used during allocation.
    /// The caller of this function becomes responsible for proper deallocation of the memory
    /// behind the pointer. This can for example be done by reconstructing the `AlignedBox` using
    /// `AlignedBox::from_raw_parts`.
    pub fn into_raw_parts(mut from: AlignedBox<T>) -> (*mut T, std::alloc::Layout) {
        // self.container is always Some, so we can just unwrap
        let container = from.container.take().unwrap();
        let ptr = std::boxed::Box::into_raw(container);
        let layout = from.layout;
        std::mem::forget(from); // AlignedBox::drop() must not be called
        (ptr, layout)
    }

    /// Construct an `AlignedBox` from a raw pointer and the layout that has been used to allocate
    /// the memory behind that pointer. After calling this function, the pointer is owned by the
    /// `AlignedBox`. In particular, the memory will be freed when the `AlignedBox` is dropped.
    /// This is only safe if the given layout is the same as the one that was used during memory
    /// allocation.
    ///
    /// # Safety
    /// The function is unsafe because improper use can lead to issues, such as double-free. Also,
    /// behavior is undefined if the given layout does not correspond to the one used for
    /// allocation.
    pub unsafe fn from_raw_parts(ptr: *mut T, layout: std::alloc::Layout) -> AlignedBox<T> {
        let container = Some(std::boxed::Box::from_raw(ptr));
        AlignedBox::<T> { container, layout }
    }
}

impl<T> AlignedBox<T> {
    /// Store `value` of type `T` on the heap, making sure that it is aligned to a multiple of
    /// `alignment`. It is also checked if `alignment` is a valid alignment for type `T` or
    /// increased to a valid alignment otherwise.
    ///
    /// # Example
    /// Place value 17 of type `i32` on the heap, aligned to 64 bytes:
    /// ```
    /// use aligned_box::AlignedBox;
    ///
    /// let b = AlignedBox::<i32>::new(64, 17);
    /// ```
    pub fn new(
        mut alignment: usize,
        value: T,
    ) -> std::result::Result<AlignedBox<T>, std::boxed::Box<dyn std::error::Error>> {
        if alignment < std::mem::align_of::<T>() {
            alignment = std::mem::align_of::<T>();
        }

        let memsize: usize = std::mem::size_of::<T>();
        if memsize == 0 {
            return Err(AlignedBoxError::ZeroAlloc.into());
        }

        let layout = std::alloc::Layout::from_size_align(memsize, alignment)?;

        // SAFETY: Requirements on layout are enforced by using from_size_align().
        let ptr = unsafe { std::alloc::alloc(layout) as *mut T };
        if ptr.is_null() {
            return Err(AlignedBoxError::OutOfMemory.into());
        }

        // SAFETY: The pointer is non-null, refers to properly sized and aligned memory and it is
        // consumed such that it cannot be used from anywhere outside the Box.
        let mut b = unsafe { AlignedBox::<T>::from_raw_parts(ptr, layout) };

        // *b is not a valid instance of T but uninitialized memory. We have to write to it without
        // dropping the old (invalid) value. Also the original value must not be dropped.
        // SAFETY: Both value and *b point to valid and properly aligned memory.
        unsafe { std::ptr::write(&mut *b, value) };

        Ok(b)
    }
}

impl<T: Default> AlignedBox<[T]> {
    /// Allocate memory for `nelems` values of type `T` on the heap, making sure that it is aligned
    /// to a multiple of `alignment`. All values are initialized by the default value of type `T`.
    /// It is also checked if `alignment` is a valid alignment for type `T` or increased to a
    /// valid alignment otherwise.
    ///
    /// # Example
    /// Allocate memory for 1024 values of type `f32` on the heap, aligned to 128 bytes. Values
    /// are initialized by their default value:
    /// ```
    /// use aligned_box::AlignedBox;
    ///
    /// let b = AlignedBox::<[f32]>::slice_from_default(128, 1024);
    /// ```
    pub fn slice_from_default(
        mut alignment: usize,
        nelems: usize,
    ) -> std::result::Result<AlignedBox<[T]>, std::boxed::Box<dyn std::error::Error>> {
        if alignment < std::mem::align_of::<T>() {
            alignment = std::mem::align_of::<T>();
        }

        // Make sure the requested amount of Ts will fit into a slice.
        let maxelems = (isize::MAX as usize) / std::mem::size_of::<T>();
        if nelems > maxelems {
            return Err(AlignedBoxError::TooManyElements.into());
        }

        let memsize: usize = std::mem::size_of::<T>() * nelems;
        if memsize == 0 {
            return Err(AlignedBoxError::ZeroAlloc.into());
        }

        let layout = std::alloc::Layout::from_size_align(memsize, alignment)?;

        // SAFETY: Requirements on layout are enforced by using from_size_align().
        let ptr = unsafe { std::alloc::alloc(layout) as *mut T };
        if ptr.is_null() {
            return Err(AlignedBoxError::OutOfMemory.into());
        }

        // SAFETY: Requirements on ptr and nelems have been verified: ptr is non-null, nelems does
        // not exceed the maximum size. The referenced memory is not accessed as long as slice
        // exists.
        let slice = unsafe { std::slice::from_raw_parts_mut(ptr, nelems) };

        // SAFETY: We only create a single Box from the given slice. The slice itself is consumed
        // so that the referenced memory can be modified from now on.
        let mut b = unsafe { AlignedBox::<[T]>::from_raw_parts(slice, layout) };

        for i in (*b).iter_mut() {
            let d = T::default(); // create new default value

            // *i is not a valid instance of T but uninitialized memory. We have to write to it
            // without dropping the old (invalid) value. Also d must not be dropped.
            // SAFETY: Both value and b point to valid and properly aligned memory.
            unsafe { std::ptr::write(&mut *i, d) };
        }

        Ok(b)
    }
}

impl<T: Copy> AlignedBox<[T]> {
    /// Allocate memory for `nelems` values of type `T` on the heap, making sure that it is aligned
    /// to a multiple of `alignment`. All values are initialized by copies of `value`. It is also
    /// checked if `alignment` is a valid alignment for type `T` or increased to a
    /// valid alignment otherwise.
    ///
    /// # Example
    /// Allocate memory for 1024 values of type `f32` on the heap, aligned to 128 bytes. All values
    /// are initialized with PI:
    /// ```
    /// use aligned_box::AlignedBox;
    ///
    /// let b = AlignedBox::<[f32]>::slice_from_value(128, 1024, std::f32::consts::PI);
    pub fn slice_from_value(
        mut alignment: usize,
        nelems: usize,
        value: T,
    ) -> std::result::Result<AlignedBox<[T]>, std::boxed::Box<dyn std::error::Error>> {
        if alignment < std::mem::align_of::<T>() {
            alignment = std::mem::align_of::<T>();
        }

        // Make sure the requested amount of Ts will fit into a slice.
        let maxelems = (isize::MAX as usize) / std::mem::size_of::<T>();
        if nelems > maxelems {
            return Err(AlignedBoxError::TooManyElements.into());
        }

        let memsize: usize = std::mem::size_of::<T>() * nelems;
        if memsize == 0 {
            return Err(AlignedBoxError::ZeroAlloc.into());
        }

        let layout = std::alloc::Layout::from_size_align(memsize, alignment)?;

        // SAFETY: Requirements on layout are enforced by using from_size_align().
        let ptr = unsafe { std::alloc::alloc(layout) as *mut T };
        if ptr.is_null() {
            return Err(AlignedBoxError::OutOfMemory.into());
        }

        // SAFETY: Requirements on ptr and nelems have been verified: ptr is non-null, nelems does
        // not exceed the maximum size. The referenced memory is not accessed as long as slice
        // exists.
        let slice = unsafe { std::slice::from_raw_parts_mut(ptr, nelems) };

        // SAFETY: We only create a single Box from the given slice. The slice itself is consumed
        // so that the referenced memory can be modified from now on.
        let mut b = unsafe { AlignedBox::<[T]>::from_raw_parts(slice, layout) };

        for i in (*b).iter_mut() {
            // T is Copy and therefore also !Drop. We can simply copy from value to *i without
            // worrying about dropping.
            *i = value;
        }

        Ok(b)
    }
}

#[cfg(test)]
mod tests {
    use super::AlignedBox;
    use lazy_static::lazy_static;

    lazy_static! {
        static ref SEQ_TEST_MUTEX: std::sync::RwLock<()> = std::sync::RwLock::new(());
    }

    #[test]
    fn alignment() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let alignments = [
            1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536,
            131072, 262144, 524288, 1048576,
        ];

        for a in alignments.iter() {
            let b = AlignedBox::<[u8]>::slice_from_value(*a, 42, 0).unwrap();
            assert_eq!((b.as_ptr() as usize) % *a, 0);
        }
    }

    #[test]
    fn size() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let sizes = [
            1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536,
            131072, 262144, 524288, 1048576,
        ];

        for s in sizes.iter() {
            let b = AlignedBox::<[u8]>::slice_from_value(1, *s, 0).unwrap();
            assert_eq!(b.len(), *s);
        }
    }

    #[test]
    fn free() {
        // This test should not run concurrently with other tests since multiple threads influence
        // each other and addresses are not reproducible.
        let _m = SEQ_TEST_MUTEX.write().unwrap();

        const ATTEMPTS: usize = 1000;
        let alignment = 1024 * 1024; // 1MB
        let size = 1024; // 1KB

        let mut addrs = std::collections::HashSet::new();

        // Test if dropping the box actually frees the memory. If any two allocations give us the
        // same address, this is the case.
        for _ in 0..ATTEMPTS {
            let b = AlignedBox::<[u8]>::slice_from_value(alignment, size, 0).unwrap();
            addrs.insert(b.as_ptr() as usize);
        }
        assert_ne!(addrs.len(), ATTEMPTS);
    }

    #[test]
    fn aliasing() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let size = 1024; // 1KB

        let b1 = AlignedBox::<[u8]>::slice_from_value(1, size, 0).unwrap();
        let b2 = AlignedBox::<[u8]>::slice_from_value(1, size, 0).unwrap();
        let addr1 = b1.as_ptr() as usize;
        let addr2 = b2.as_ptr() as usize;
        assert_ne!(addr1, addr2);
        if addr1 > addr2 {
            assert!((addr1 - addr2) >= size);
        } else {
            assert!((addr2 - addr1) >= size);
        }
    }

    #[test]
    fn read_write() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let mut b = AlignedBox::<[f32]>::slice_from_value(128, 1024, 3.1415).unwrap();
        for i in b.iter() {
            assert_eq!(*i, 3.1415);
        }
        let mut ctr: f32 = 0.0;
        for i in b.iter_mut() {
            *i = ctr;
            ctr += 1.0;
        }
        ctr = 0.0;
        for i in b.iter() {
            assert_eq!(*i, ctr);
            ctr += 1.0;
        }
    }

    #[test]
    fn defaults() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        #[derive(PartialEq, Eq, Debug)]
        struct SomeVaryingDefault {
            i: i32,
        }

        static COUNTER: std::sync::atomic::AtomicI32 = std::sync::atomic::AtomicI32::new(0);

        impl Default for SomeVaryingDefault {
            fn default() -> SomeVaryingDefault {
                SomeVaryingDefault {
                    i: COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
                }
            }
        }

        let b = AlignedBox::<[SomeVaryingDefault]>::slice_from_default(128, 1024).unwrap();
        assert_eq!(SomeVaryingDefault::default().i, 1024);
        let mut ctr = 0;
        for i in b.iter() {
            assert_eq!(i.i, ctr);
            ctr += 1;
        }
    }

    #[test]
    fn move_sem() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let v = vec![1, 2, 3];
        let b = AlignedBox::new(128, v).unwrap();
        assert_eq!(*b, vec![1, 2, 3]);
    }

    #[test]
    fn copy_sem() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let v = 17;
        let _ = AlignedBox::new(128, v).unwrap();
        let _ = AlignedBox::slice_from_value(128, 1024, v).unwrap();
        assert_eq!(v, 17);
    }

    #[test]
    fn min_align() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        #[repr(C, align(131072))]
        struct LargeAlign {
            x: u8,
        }

        assert_eq!(std::mem::align_of::<LargeAlign>(), 131072);

        let a = LargeAlign { x: 28 };
        let b = AlignedBox::<LargeAlign>::new(1, a).unwrap();
        let (ptr, layout) = AlignedBox::into_raw_parts(b);
        assert_eq!((ptr as usize) % std::mem::align_of::<LargeAlign>(), 0);
        let _ = unsafe { AlignedBox::from_raw_parts(ptr, layout) };
    }

    #[test]
    fn clone() {
        let _m = SEQ_TEST_MUTEX.read().unwrap();

        let mut b = AlignedBox::new(128, vec![47, 11]).unwrap();
        let mut another_b = b.clone();

        assert_eq!(*b, *another_b);

        b[0] = 48;
        another_b[1] = 12;

        assert_eq!(b[0], 48);
        assert_eq!(b[1], 11);
        assert_eq!(another_b[0], 47);
        assert_eq!(another_b[1], 12);
    }
}