nibblecode 0.1.0

use std::alloc::{AllocError, Allocator, Global, Layout};
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::ptr::NonNull;

use crate::Serialize;

/// A [`core::alloc::Allocator`] that aligns all allocations to at least `T::ALIGN` bytes.
pub struct AlignedAlloc<T: Serialize + ?Sized, A: Allocator = Global>(A, PhantomData<T>);

impl<T: Serialize + ?Sized, A: Allocator> AlignedAlloc<T, A> {
	pub fn new_in(alloc: A) -> Self {
		Self(alloc, PhantomData)
	}
}

impl<T: Serialize + ?Sized> Default for AlignedAlloc<T, Global> {
	fn default() -> AlignedAlloc<T, Global> {
		Self(Global, PhantomData)
	}
}

unsafe impl<T: Serialize + ?Sized, A: Allocator> Allocator for AlignedAlloc<T, A> {
	fn allocate(&self, mut layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
		// If the requested align is smaller than ours, update it
		if layout.align() < T::ALIGN.as_usize() {
			// Make sure the size isn't more than the maax size for this align
			if isize::MAX as usize + 1 - T::ALIGN.as_usize() < layout.size() {
				return Err(AllocError);
			}

			// Safety: we know `ALIGN` is a power of 2
			layout =
				unsafe { Layout::from_size_align_unchecked(layout.size(), T::ALIGN.as_usize()) };
		}

		self.0.allocate(layout)
	}

	unsafe fn deallocate(&self, ptr: NonNull<u8>, mut layout: Layout) {
		unsafe {
			// If the requested align is smaller than ours, update it
			if layout.align() < T::ALIGN.as_usize() {
				// Don't bother checking for valid size for align like in `allocate`, because we
				// should be recieving a layout that was able to allocate. Also, we can't return an
				// error from `deallocate`.

				// Safety: we know `ALIGN` is a power of 2
				layout = Layout::from_size_align_unchecked(layout.size(), T::ALIGN.as_usize());
			}

			self.0.deallocate(ptr, layout);
		}
	}
}

impl<T: Serialize + ?Sized, A: Allocator + Clone> Clone for AlignedAlloc<T, A> {
	fn clone(&self) -> AlignedAlloc<T, A> {
		AlignedAlloc(self.0.clone(), PhantomData)
	}
}

pub fn new_uninit_boxed_slice<T: Serialize + ?Sized>(
	heap_size: usize,
) -> Box<[MaybeUninit<u8>], AlignedAlloc<T>> {
	Box::new_uninit_slice_in(
		size_of::<T::Archived>() + heap_size,
		AlignedAlloc(Global, PhantomData),
	)
}