mnemosyne-local 0.5.0

use crate::alloc::small_path_class;
use crate::free_helpers::resolve_owner_slot;
use crate::usable_size;
use crate::validation::{init_bytes, poison_bytes};
use crate::{
    LocalAllocatorSelector, ThreadAllocator, initialize_allocated_bytes, poison_freed_bytes,
    thread_alloc_layout, thread_free,
};
use core::alloc::Layout;
use core::ptr::NonNull;
use mnemosyne_arena::HasSegmentPool;
use mnemosyne_core::constants::{MAX_SMALL_ALLOC_SIZE, MIN_BLOCK_SIZE};
use mnemosyne_core::policy::AllocPolicy;
use mnemosyne_core::size_class::round_up_size;
use mnemosyne_core::types::Segment;
use mnemosyne_core::types::{Block, locate_segment};

/// Non-generic SSOT for the in-place realloc byte-mutation step.
///
/// When `realloc_can_reuse` returns true, the block stays in place but the
/// delta region must be initialized (grow) or poisoned (shrink) per policy.
/// All policy flags are passed as plain booleans so this body compiles once
/// regardless of how many `(P, B)` pairs call `thread_realloc`.
///
/// # Safety
///
/// `ptr` must be a live allocation whose usable size covers `new_size`.
#[inline(always)]
unsafe fn realloc_delta_init(
    ptr: *mut u8,
    old_size: usize,
    new_size: usize,
    zero_init: bool,
    poison: bool,
    alloc_byte: u8,
    free_byte: u8,
) {
    if new_size > old_size {
        // Grow: initialize the newly accessible region.
        // SAFETY: `ptr.add(old_size)` is within the backing block (caller
        // contract: usable size >= new_size > old_size).
        let delta_ptr = unsafe { ptr.add(old_size) };
        let delta = new_size - old_size;
        unsafe { init_bytes(delta_ptr, delta, zero_init, poison, alloc_byte) };
    } else if new_size < old_size {
        // Shrink: poison the truncated tail.
        // SAFETY: `ptr.add(new_size)` through `old_size - new_size` bytes is
        // within the backing block (new_size < old_size <= usable size).
        let tail_ptr = unsafe { ptr.add(new_size) };
        let truncated = old_size - new_size;
        unsafe { poison_bytes(tail_ptr, truncated, poison, free_byte) };
    }
}

/// Whether a small reallocation can stay in its current size class.
///
/// True when the new size still rounds to the same class stride, so the
/// existing block already has room and `realloc` can return the same
/// pointer without copying. Alignments above `MIN_BLOCK_SIZE` are
/// rejected outright, since those are not served by the small path.
#[inline(always)]
pub fn small_realloc_fits_existing_class(layout: Layout, new_size: usize) -> bool {
    if layout.align() > MIN_BLOCK_SIZE {
        return false;
    }

    // The old allocation occupies the block for size class
    // `size_to_class(old_adjusted_size)`, whose stride is `round_up_size` of
    // that size. `new_size` fits in place iff it does not exceed that block
    // stride. `round_up_size` is the core size-class SSOT (the const-fn stride
    // schedule), so route through it instead of re-encoding the 128/512/2048
    // breakpoints and their round-up masks here. A size past
    // `MAX_SMALL_ALLOC_SIZE` yields `None`, which correctly reports "does not
    // fit a small class in place".
    let old_adjusted_size = core::cmp::max(layout.size(), layout.align());
    match round_up_size(old_adjusted_size) {
        Some(block_stride) => new_size <= block_stride,
        None => false,
    }
}

/// Determines whether a reallocation can be served in-place without copying.
///
/// This is a **non-generic** helper: it uses only `MAX_SMALL_ALLOC_SIZE`,
/// `MIN_BLOCK_SIZE`, and `usable_size` (all non-generic), so it compiles
/// once and is shared across every `(P, B)` monomorphization of
/// `thread_realloc`, reducing the amount of code specialized per policy.
///
/// Returns `true` when the existing block is large enough for `new_size` and
/// no copy is required.
///
/// # Safety
///
/// `ptr` must be a non-null allocation of `layout` returned by this allocator;
/// `usable_size(ptr)` reads the originating segment metadata.
#[inline]
unsafe fn realloc_can_reuse(ptr: *mut u8, layout: Layout, new_size: usize) -> bool {
    let is_small = layout.size() <= MAX_SMALL_ALLOC_SIZE && layout.align() <= MIN_BLOCK_SIZE;
    if new_size <= layout.size() {
        if is_small {
            // Small shrink: the existing class block already holds `new_size`;
            // reuse when the shrink stays within the half-capacity threshold
            // that prevents excessive block fragmentation.
            return new_size >= layout.size() / 2;
        }
        // Large/huge shrink: fast-path reuse for modest shrinks; page-rounded
        // comparison for larger ones.
        let new_adjusted = core::cmp::max(new_size, layout.align());
        if new_size >= layout.size() / 2 {
            return true;
        }
        if new_adjusted > MAX_SMALL_ALLOC_SIZE || layout.align() > MIN_BLOCK_SIZE {
            // SAFETY: `ptr` is a live allocation per the caller's contract;
            // `usable_size` reads only segment/page metadata.
            let current_usable = unsafe { usable_size(ptr) };
            let page_size = mnemosyne_core::constants::PAGE_SIZE;
            let new_page_rounded = (new_adjusted + page_size - 1) & !(page_size - 1);
            return new_page_rounded >= current_usable;
        }
        false
    } else {
        // Grow: reuse in place if the allocation already has capacity.
        if is_small {
            small_realloc_fits_existing_class(layout, new_size)
        } else {
            // SAFETY: same contract as above.
            let current_usable = unsafe { usable_size(ptr) };
            new_size <= current_usable
        }
    }
}

/// Reallocates a memory block, optimizing performance and memory footprint by avoiding redundant
/// allocation-deallocation cycles, reusing existing size-class blocks in place, and reducing TLS
/// lookup overhead.
///
/// # Safety
///
/// Same contract as `GlobalAlloc::realloc`.
/// # Examples
///
/// ```
/// use core::alloc::Layout;
/// use mnemosyne_local::{thread_alloc_layout, thread_free, thread_realloc};
/// use mnemosyne_core::StandardPolicy;
/// use mnemosyne_backend::MemoryBackendWrapper as Backend;
///
/// let layout = Layout::from_size_align(64, 8).expect("64/8 is a valid layout");
///
/// // SAFETY: `layout` describes the live allocation being resized, and the
/// // returned pointer replaces it — the old one must not be used or freed
/// // again once `thread_realloc` returns non-null.
/// unsafe {
///     let p = thread_alloc_layout::<StandardPolicy, Backend>(layout.size(), layout.align());
///     assert!(!p.is_null());
///     p.write_bytes(0x3C, layout.size());
///
///     let grown = thread_realloc::<StandardPolicy, Backend>(p, layout, 256);
///     assert!(!grown.is_null());
///
///     // Growing preserves every byte of the original contents.
///     for off in 0..layout.size() {
///         assert_eq!(*grown.add(off), 0x3C, "byte {off} did not survive the move");
///     }
///
///     thread_free::<StandardPolicy, Backend>(grown);
/// }
/// ```
#[inline]
pub unsafe fn thread_realloc<
    P: AllocPolicy + crate::tls_slot::PolicySlotSelection<B>,
    B: HasSegmentPool + LocalAllocatorSelector<B>,
>(
    ptr: *mut u8,
    layout: Layout,
    new_size: usize,
) -> *mut u8 {
    if !ptr.is_null() && new_size != 0 {
        // SAFETY: `ptr` is non-null and allocator-owned per the caller's
        // `# Safety` contract; `realloc_can_reuse` only reads metadata.
        let can_reuse = unsafe { realloc_can_reuse(ptr, layout, new_size) };

        if can_reuse {
            // SAFETY: `ptr` is non-null and its backing block covers `new_size`
            // (proven by `can_reuse`).
            unsafe {
                realloc_delta_init(
                    ptr,
                    layout.size(),
                    new_size,
                    P::ZERO_INITIALIZE,
                    P::ENABLE_POISONING,
                    P::POISON_ALLOC_BYTE,
                    P::POISON_FREE_BYTE,
                )
            };
            return ptr;
        }
    } else {
        if ptr.is_null() {
            if new_size == 0 {
                return core::ptr::null_mut();
            }
            // SAFETY: `new_size != 0` and `layout.align()` is a power of two
            // (valid Layout invariant); thread_alloc_layout's contract is met.
            return unsafe { thread_alloc_layout::<P, B>(new_size, layout.align()) };
        }
        // new_size == 0 && !ptr.is_null()
        // SAFETY: `ptr` is a live non-null allocation from this allocator;
        // thread_free's contract is met.
        unsafe { thread_free::<P, B>(ptr) };
        return core::ptr::null_mut();
    }

    let new_adjusted = core::cmp::max(new_size, layout.align());
    // Use the shared routing decision so the in-place small-realloc target class
    // honours the requested alignment. `size_to_class_nonzero(new_adjusted)`
    // alone could pick a class whose stride does not carry `align` (e.g. class
    // 224 for a 64-byte-aligned 200-byte request), yielding a misaligned block.
    // `None` falls through to the `thread_alloc_layout` path below, which routes
    // correctly (small or huge) for the alignment.
    let new_class = small_path_class(new_size, layout.align());

    // SAFETY: `ptr` is non-null and allocator-owned per the `# Safety`
    // contract, satisfying `locate_segment`'s precondition; it recovers the live
    // segment header and the bounded page index.
    let (segment, page_index) = unsafe { locate_segment(ptr) };

    // SAFETY: `segment`/`page_index` come from `locate_segment` on an
    // allocator-owned `ptr`, so the segment header is live and the index is in
    // bounds of its `pages` array. Keep this as a raw pointer until the old
    // block has been replaced: `alloc_class` may read the same segment while
    // selecting its fresh page, and an outstanding `&mut Page` would violate
    // Stacked Borrows before the replacement free begins.
    let page = unsafe { &raw mut (*segment).pages[page_index] };
    // SAFETY: `page` is the live page recovered from the allocator-owned
    // segment and `block_size` is initialized for every allocated small block.
    let is_old_small = page_index > 0 && unsafe { (*page).block_size > 0 };

    let mut new_ptr = core::ptr::null_mut();
    let mut local_free_done = false;

    if is_old_small && let Some(class) = new_class {
        let slot_ptr = B::get_allocator_ptr_raw_for_policy::<P>();
        let owner = unsafe { Segment::owner(segment) };
        // SAFETY: `segment` is the live header from `locate_segment`; the platform
        // helper only reads ownership metadata through raw-pointer projections.
        let (is_owner, owner_slot) = unsafe { resolve_owner_slot(segment, owner, slot_ptr) };

        if is_owner && !owner_slot.is_null() {
            // SAFETY: `owner_slot` is the live allocator slot that owns this
            // segment, even when the caller uses a different policy slot for the
            // same thread. The holder of that slot remains the authoritative owner
            // for the segment-local free-list mutation below.
            if !unsafe { crate::tls_slot::LocalAllocatorSlot::<B>::is_allocating(owner_slot) } {
                let alloc = unsafe { &mut *(owner_slot as *mut ThreadAllocator<B>) };
                unsafe {
                    crate::tls_slot::LocalAllocatorSlot::<B>::set_allocating(owner_slot, true)
                };
                // SAFETY: `segment` owns the free-list encoding and page
                // randomization bits, so a standard-policy caller may still
                // need to allocate/free under the owning hardened mode.
                let encrypted = unsafe { Segment::free_list_encrypted(segment) };
                let allocated = if encrypted {
                    unsafe { alloc.alloc_class::<mnemosyne_core::policy::HardenedPolicy>(class) }
                } else {
                    unsafe { alloc.alloc_class::<P>(class) }
                };
                new_ptr = allocated;
                if !new_ptr.is_null() {
                    crate::bin_stats::record_alloc_with_size(class, new_adjusted);
                    unsafe {
                        // SAFETY: `new_ptr` is a fresh block of at least
                        // `new_adjusted` bytes; init writes only within it.
                        initialize_allocated_bytes::<P>(new_ptr, new_adjusted);
                        // SAFETY: `ptr` (old, valid for `layout.size()`)
                        // and `new_ptr` (fresh, distinct block) are
                        // non-overlapping; copy length is the smaller size.
                        core::ptr::copy_nonoverlapping(
                            ptr,
                            new_ptr,
                            core::cmp::min(layout.size(), new_size),
                        );
                        // SAFETY: `segment`/`page_index` identify the
                        // live page and its key slot. Read the segment
                        // metadata before materializing `page_ref`,
                        // because the exclusive page borrow must not
                        // overlap this shared parent-segment access.
                        let cookie = Segment::cookie_for_dynamic(segment, encrypted, page_index);
                        // SAFETY: `page` is the exclusively-borrowed page
                        // owning the old block; reborrowing yields the sole
                        // live `&mut` for the free bookkeeping below.
                        let page_ref = &mut *page;
                        if P::ENABLE_POISONING {
                            // SAFETY: `ptr` is the old block, valid for the
                            // page's `block_size` bytes being poisoned.
                            poison_freed_bytes::<P>(ptr, page_ref.block_size as usize);
                        }
                        let block = ptr as *mut Block;
                        let page_free = page_ref.free;
                        let page_alloc_count = page_ref.alloc_count as usize;
                        let randomized = (P::RANDOMIZE_ALLOCATION && encrypted)
                            || page_ref.secondary_free.is_some();
                        if page_ref.alloc_count == 0 {
                            std::process::abort();
                        }
                        // SAFETY: `block` is the old user pointer, non-null
                        // by the allocator invariant; `new_unchecked` is
                        // sound and equality with `page_free` is the
                        // double-free guard.
                        if Some(NonNull::new_unchecked(block)) == page_free
                            || (randomized
                                && Some(NonNull::new_unchecked(block)) == page_ref.secondary_free)
                        {
                            std::process::abort();
                        }
                        if page_free.is_some()
                            && (page_alloc_count != 1 || alloc.is_current_segment(segment))
                        {
                            // SAFETY: in-place free — `block` is the guarded
                            // old block, `page_free`/`cookie` are `page_ref`'s
                            // current head and cookie, and `page_alloc_count`
                            // is its live count (`>= 1`), so the shared commit
                            // stays inside this owned page.
                            crate::free_helpers::commit_in_place_free(
                                block,
                                page_ref,
                                page_free,
                                cookie,
                                encrypted,
                                page_alloc_count,
                                randomized,
                            );
                        } else {
                            // SAFETY: `block` belongs to `page_ref` in
                            // `segment` at `page_index`, and `alloc` owns
                            // them — exactly `do_local_free_internal`'s
                            // contract for the page-list transition path.
                            let _became_empty = if encrypted {
                                crate::do_local_free_internal_policy::<
                                    mnemosyne_core::policy::HardenedPolicy,
                                    B,
                                >(
                                    alloc, block, page_ref, segment, page_index
                                )
                            } else {
                                crate::do_local_free_internal_policy::<P, B>(
                                    alloc, block, page_ref, segment, page_index,
                                )
                            };
                        }
                    }
                    local_free_done = true;
                }
                // SAFETY: `owner_slot` is the owning allocator slot for the
                // segment; clearing the gate on that slot is the matching
                // release of the borrow above.
                unsafe {
                    crate::tls_slot::LocalAllocatorSlot::<B>::set_allocating(owner_slot, false)
                };
            }
        }
    }

    if new_ptr.is_null() {
        // SAFETY: `new_size != 0` and `layout.align()` is a valid power-of-two
        // alignment; thread_alloc_layout's contract is met.
        new_ptr = unsafe { thread_alloc_layout::<P, B>(new_size, layout.align()) };
        if new_ptr.is_null() {
            return core::ptr::null_mut();
        }
    }

    if !local_free_done {
        unsafe {
            core::ptr::copy_nonoverlapping(ptr, new_ptr, core::cmp::min(layout.size(), new_size));
            thread_free::<P, B>(ptr);
        }
    }

    new_ptr
}