stumpalo 0.5.0

use core::alloc::LayoutError;
use core::{alloc::Layout, mem::size_of, ptr};
use core_alloc::alloc::{alloc, dealloc, handle_alloc_error};

use crate::errors::Error;

pub(crate) const HEADER_SIZE: usize = size_of::<ChunkHeader>();
pub const INITIAL_CHUNK_CAPACITY: usize = 256 - HEADER_SIZE;
pub(crate) const CHUNK_ALIGN: usize = 16;

pub(crate) struct ChunkHeader {
    // When this chunk is in use, this points to the previously
    // used chunk.
    // When this chunk is free, this points to the next free chunk.
    pub(crate) next: *mut ChunkHeader,
    pub(crate) capacity: usize,
}

pub(crate) const FIRST_HEADER: ChunkHeader = ChunkHeader {
    next: ptr::null_mut::<ChunkHeader>(),
    capacity: 0,
};

#[inline(never)]
pub(crate) fn count_bytes_allocated_in_chunk_seq_with_meta(mut head: *mut ChunkHeader) -> usize {
    let mut res = 0;
    while !head.is_null() {
        unsafe {
            let node = head.as_ref_unchecked();
            res += node.capacity + HEADER_SIZE;
            head = node.next;
        }
    }
    res
}

#[inline(never)]
pub(crate) fn count_bytes_allocated_in_chunk_seq(mut head: *mut ChunkHeader) -> usize {
    let mut res = 0;
    while !head.is_null() {
        unsafe {
            let node = head.as_ref_unchecked();
            res += node.capacity;
            head = node.next;
        }
    }
    res
}

/// ChunkHeaders form a linked list.
/// [`join_chunk_chains`] appends 'b' to 'a'.
pub(crate) fn join_chunk_chains(a: *mut ChunkHeader, b: *mut ChunkHeader) -> *mut ChunkHeader {
    if a.is_null() {
        return b;
    }
    if b.is_null() {
        return a;
    }
    let mut curr = a;
    loop {
        let next = unsafe { (*curr).next };
        if next.is_null() {
            break;
        }
        curr = next;
    }
    unsafe { (*curr).next = b };
    a
}

enum NewChunkRes {
    LayoutError(LayoutError),
    AllocErr(Layout),
    Success(*mut ChunkHeader),
}

impl ChunkHeader {
    // Returns a pointer to the chunk_header.
    // Panics if the requested size is too large to form a valid Layout,
    // or if the underlying allocator returns null (OOM).
    #[inline]
    fn new_internal(size: usize, next: *mut ChunkHeader) -> NewChunkRes {
        let total = size + HEADER_SIZE;
        let layout = match Layout::from_size_align(total, CHUNK_ALIGN) {
            Ok(a) => a,
            Err(e) => return NewChunkRes::LayoutError(e),
        };
        let chunk: *mut ChunkHeader = unsafe { alloc(layout).cast::<ChunkHeader>() };
        if chunk.is_null() {
            return NewChunkRes::AllocErr(layout);
        }
        let header = ChunkHeader {
            next,
            capacity: size,
        };
        unsafe {
            ptr::write(chunk, header);
        }
        NewChunkRes::Success(chunk)
    }

    #[inline]
    pub(crate) fn new(size: usize, next: *mut ChunkHeader) -> *mut ChunkHeader {
        match Self::new_internal(size, next) {
            NewChunkRes::LayoutError(err) => {
                panic!("{}", err);
            }
            NewChunkRes::AllocErr(layout) => handle_alloc_error(layout),
            NewChunkRes::Success(res) => res,
        }
    }

    #[inline]
    pub(crate) fn try_new(size: usize, next: *mut ChunkHeader) -> Result<*mut ChunkHeader, Error> {
        match Self::new_internal(size, next) {
            NewChunkRes::LayoutError(err) => Err(Error::LayoutError(err)),
            NewChunkRes::AllocErr(_) => Err(Error::OOM),
            NewChunkRes::Success(res) => Ok(res),
        }
    }

    pub(crate) unsafe fn free_chunk_chain(mut current: *mut ChunkHeader) {
        while !current.is_null() {
            let prev = unsafe { (*current).next };
            let capacity = unsafe { (*current).capacity };
            if capacity > 0 {
                let layout = unsafe {
                    Layout::from_size_align_unchecked(capacity + HEADER_SIZE, CHUNK_ALIGN)
                };
                unsafe { dealloc(current.cast::<u8>(), layout) };
            }
            current = prev;
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::chunk::HEADER_SIZE;

    #[test]
    fn size_of_header() {
        assert_eq!(HEADER_SIZE, 16);
    }
}