bump-scope 2.3.0

use core::{alloc::Layout, num::NonZeroUsize, ptr::NonNull};

use crate::{
    BaseAllocator, alloc::AllocError, bump_down, layout::CustomLayout, polyfill::non_null, raw_bump::RawBump,
    settings::BumpAllocatorSettings, up_align_usize_unchecked,
};

#[inline(always)]
pub fn allocate<A, S>(bump: &RawBump<A, S>, layout: Layout) -> Result<NonNull<[u8]>, AllocError>
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    Ok(NonNull::slice_from_raw_parts(
        bump.alloc::<AllocError>(layout)?,
        layout.size(),
    ))
}

#[inline(always)]
pub unsafe fn deallocate<A, S>(bump: &RawBump<A, S>, ptr: NonNull<u8>, layout: Layout)
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    if !S::DEALLOCATES {
        return;
    }

    unsafe {
        // free allocated space if this is the last allocation
        if is_last(bump, ptr, layout) {
            deallocate_assume_last(bump, ptr, layout);
        }
    }
}

#[inline(always)]
unsafe fn deallocate_assume_last<A, S>(bump: &RawBump<A, S>, ptr: NonNull<u8>, layout: Layout)
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    if !S::DEALLOCATES {
        return;
    }

    unsafe {
        let chunk = bump.chunk.get().as_non_dummy_unchecked();
        if S::UP {
            chunk.set_pos_addr_and_align(ptr.addr().get());
        } else {
            let mut addr = ptr.addr().get();
            addr += layout.size();
            chunk.set_pos_addr_and_align(addr);
        }
    }
}

/// Checks if the memory block is the last one in the bump allocator.
///
/// The given memory block must have been allocated by some bump allocator,
/// not necessarily the given one (but if it isn't then it will always return false).
///
/// This condition must be true to be able to do grow, shrink or deallocate in place.
///
/// If this function returns true then the bump is guaranteed to have a non-dummy
/// chunk since a memory block is only returned from a bump allocator with a non-dummy
/// chunk.
#[inline(always)]
unsafe fn is_last<A, S>(bump: &RawBump<A, S>, ptr: NonNull<u8>, layout: Layout) -> bool
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    if S::UP {
        (unsafe { ptr.add(layout.size()) }) == bump.chunk.get().pos()
    } else {
        ptr == bump.chunk.get().pos()
    }
}

#[inline(always)]
pub unsafe fn grow<A, S>(
    bump: &RawBump<A, S>,
    old_ptr: NonNull<u8>,
    old_layout: Layout,
    new_layout: Layout,
) -> Result<NonNull<[u8]>, AllocError>
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    debug_assert!(
        new_layout.size() >= old_layout.size(),
        "`new_layout.size()` must be greater than or equal to `old_layout.size()`"
    );

    unsafe {
        if S::UP {
            if is_last(bump, old_ptr, old_layout) & align_fits(old_ptr, old_layout, new_layout) {
                // We may be able to grow in place! Just need to check if there is enough space.

                // `is_last` returned true, which guarantees a non-dummy
                let chunk = bump.chunk.get().as_non_dummy_unchecked();
                let chunk_end = chunk.content_end();
                let remaining = chunk_end.addr().get() - old_ptr.addr().get();

                if new_layout.size() <= remaining {
                    // There is enough space! We will grow in place. Just need to update the bump pointer.

                    let old_addr = old_ptr.addr();

                    // Up-aligning a pointer inside a chunks content by `MIN_ALIGN` never overflows.
                    let new_pos = up_align_usize_unchecked(old_addr.get() + new_layout.size(), S::MIN_ALIGN);

                    chunk.set_pos_addr(new_pos);

                    Ok(NonNull::slice_from_raw_parts(old_ptr, new_layout.size()))
                } else {
                    // The current chunk doesn't have enough space to allocate this layout. We need to allocate in another chunk.
                    let new_ptr = bump.alloc_in_another_chunk::<AllocError>(new_layout)?;
                    old_ptr.copy_to_nonoverlapping(new_ptr, old_layout.size());
                    Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
                }
            } else {
                // We can't grow in place. We have to make a new allocation.
                let new_ptr = bump.alloc::<AllocError>(new_layout)?;
                old_ptr.copy_to_nonoverlapping(new_ptr, old_layout.size());
                Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
            }
        } else {
            if is_last(bump, old_ptr, old_layout) {
                // `is_last` returned true, which guarantees a non-dummy
                let chunk = bump.chunk.get().as_non_dummy_unchecked();

                // We may be able to reuse the currently allocated space. Just need to check if the current chunk has enough space for that.
                let additional_size = new_layout.size() - old_layout.size();

                let old_addr = old_ptr.addr();
                let new_addr = bump_down(old_addr, additional_size, new_layout.align().max(S::MIN_ALIGN));

                let very_start = chunk.content_start().addr();

                if new_addr >= very_start.get() {
                    // There is enough space in the current chunk! We will reuse the allocated space.

                    let new_addr = NonZeroUsize::new_unchecked(new_addr);
                    let new_addr_end = new_addr.get() + new_layout.size();

                    let new_ptr = old_ptr.with_addr(new_addr);

                    // Check if the regions don't overlap so we may use the faster `copy_nonoverlapping`.
                    if new_addr_end < old_addr.get() {
                        old_ptr.copy_to_nonoverlapping(new_ptr, old_layout.size());
                    } else {
                        old_ptr.copy_to(new_ptr, old_layout.size());
                    }

                    // `is_last_and_allocated` returned true
                    chunk.set_pos_addr(new_addr.get());

                    Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
                } else {
                    // The current chunk doesn't have enough space to allocate this layout. We need to allocate in another chunk.
                    let new_ptr = bump.alloc_in_another_chunk::<AllocError>(new_layout)?;
                    old_ptr.copy_to_nonoverlapping(new_ptr, old_layout.size());
                    Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
                }
            } else {
                // We can't reuse the allocated space. We have to make a new allocation.
                let new_ptr = bump.alloc::<AllocError>(new_layout)?;
                old_ptr.copy_to_nonoverlapping(new_ptr, old_layout.size());
                Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
            }
        }
    }
}

#[inline(always)]
pub unsafe fn grow_zeroed<A, S>(
    bump: &RawBump<A, S>,
    old_ptr: NonNull<u8>,
    old_layout: Layout,
    new_layout: Layout,
) -> Result<NonNull<[u8]>, AllocError>
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    unsafe {
        let new_ptr = grow(bump, old_ptr, old_layout, new_layout)?;

        let delta = new_layout.size() - old_layout.size();
        new_ptr.cast::<u8>().add(old_layout.size()).write_bytes(0, delta);

        Ok(new_ptr)
    }
}

/// This shrink implementation always tries to reuse old memory if it can.
///
/// That's different to bumpalo's shrink implementation, which only shrinks if it can do so with `copy_nonoverlapping`
/// and doesn't attempt to recover memory if the alignment doesn't fit.
#[inline(always)]
pub unsafe fn shrink<A, S>(
    bump: &RawBump<A, S>,
    old_ptr: NonNull<u8>,
    old_layout: Layout,
    new_layout: Layout,
) -> Result<NonNull<[u8]>, AllocError>
where
    A: BaseAllocator<S::GuaranteedAllocated>,
    S: BumpAllocatorSettings,
{
    /// Called when `new_layout` doesn't fit alignment.
    #[cold]
    #[inline(never)]
    unsafe fn shrink_unfit<A, S>(
        bump: &RawBump<A, S>,
        old_ptr: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError>
    where
        A: BaseAllocator<S::GuaranteedAllocated>,
        S: BumpAllocatorSettings,
    {
        unsafe {
            if S::SHRINKS && is_last(bump, old_ptr, old_layout) {
                let old_pos = bump.chunk.get().pos();
                deallocate_assume_last(bump, old_ptr, old_layout);

                let overlaps;
                let new_ptr;

                if let Some(in_chunk) = bump.chunk.get().alloc(CustomLayout(new_layout)) {
                    new_ptr = in_chunk;
                    overlaps = if S::UP {
                        let old_ptr_end = old_ptr.add(new_layout.size());
                        old_ptr_end > new_ptr
                    } else {
                        let new_ptr_end = new_ptr.add(new_layout.size());
                        new_ptr_end > old_ptr
                    }
                } else {
                    // `is_last` returned true, which guarantees a non-dummy chunk
                    bump.chunk.get().as_non_dummy_unchecked().set_pos(old_pos);

                    new_ptr = bump.alloc_in_another_chunk::<AllocError>(new_layout)?;
                    overlaps = false;
                }

                if overlaps {
                    old_ptr.copy_to(new_ptr, new_layout.size());
                } else {
                    old_ptr.copy_to_nonoverlapping(new_ptr, new_layout.size());
                }

                Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
            } else {
                let new_ptr = bump.alloc::<AllocError>(new_layout)?;
                old_ptr.copy_to_nonoverlapping(new_ptr, new_layout.size());
                Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
            }
        }
    }

    debug_assert!(
        new_layout.size() <= old_layout.size(),
        "`new_layout.size()` must be smaller than or equal to `old_layout.size()`"
    );

    unsafe {
        if !align_fits(old_ptr, old_layout, new_layout) {
            return shrink_unfit(bump, old_ptr, old_layout, new_layout);
        }

        // If this is not the last allocation, then there's nothing we can do
        if !S::SHRINKS || !is_last(bump, old_ptr, old_layout) {
            // We can't shrink this allocation, so we return it as-is.
            return Ok(NonNull::slice_from_raw_parts(old_ptr, old_layout.size()));
        }

        if S::UP {
            let end = old_ptr.addr().get() + new_layout.size();

            // Up-aligning a pointer inside a chunk by `MIN_ALIGN` never overflows.
            let new_pos = up_align_usize_unchecked(end, S::MIN_ALIGN);

            // `is_last` returned true, which guarantees a non-dummy
            bump.chunk.get().as_non_dummy_unchecked().set_pos_addr(new_pos);

            Ok(NonNull::slice_from_raw_parts(old_ptr, new_layout.size()))
        } else {
            let old_addr = old_ptr.addr();
            let old_end_addr = NonZeroUsize::new_unchecked(old_addr.get() + old_layout.size());

            // The resulting `new_addr` will always be between `old_addr` and `old_end_addr` because:
            // - `new_layout`'s size must be less than `old_layout`'s size per `shrink`'s safety invariant
            // - downwards aligning can't create an address lower than `old_addr` since we checked in `align_fits` that `old_addr` is aligned to `new_layout`
            let new_addr = bump_down(old_end_addr, new_layout.size(), new_layout.align().max(S::MIN_ALIGN));
            let new_addr = NonZeroUsize::new_unchecked(new_addr);
            let new_ptr = old_ptr.with_addr(new_addr);

            let copy_src_end = NonZeroUsize::new_unchecked(old_addr.get() + new_layout.size());
            let copy_dst_start = new_addr;
            let overlaps = copy_src_end > copy_dst_start;

            if overlaps {
                old_ptr.copy_to(new_ptr, new_layout.size());
            } else {
                old_ptr.copy_to_nonoverlapping(new_ptr, new_layout.size());
            }

            // `is_last` returned true, which guarantees a non-dummy
            bump.chunk.get().as_non_dummy_unchecked().set_pos(new_ptr);

            Ok(NonNull::slice_from_raw_parts(new_ptr, new_layout.size()))
        }
    }
}

#[inline(always)]
fn align_fits(old_ptr: NonNull<u8>, _old_layout: Layout, new_layout: Layout) -> bool {
    non_null::is_aligned_to(old_ptr, new_layout.align())
}