mnemosyne-local 0.6.0

use crate::local_alloc::ThreadAllocator;
use crate::local_alloc::page::{push_page_front_raw, unlink_page_from_list_raw};
use core::ptr::NonNull;
use mnemosyne_arena::{HasSegmentPool, deallocate_segment, try_deallocate_segment};
use mnemosyne_core::constants::NUM_SIZE_CLASSES;
use mnemosyne_core::types::{OccupiedPageBits, Page, Segment};

const MIN_RETAINED_OWNED_SEGMENTS: usize = 3;
const RECLAIM_THRESHOLD_SEGMENTS: usize = MIN_RETAINED_OWNED_SEGMENTS + 1;

use crate::local_alloc::segment::deferred_chain::DeferredChain;
impl<B: HasSegmentPool> ThreadAllocator<B> {
    /// Reclaims every segment owned by this thread cache back to the global
    /// pools, then clears the owned-segment chain so the operation is
    /// idempotent.
    ///
    /// # Waiting
    ///
    /// Every caller of this function is a destructor —
    /// `ThreadAllocator::drop` and, on the `#[thread_local]` fast path,
    /// `ThreadExitReclaim::drop` — so it must not block. An empty segment is
    /// offered to [`try_deallocate_segment`], whose every wait is bounded; one
    /// still holding live allocations cannot be unmapped and so goes directly
    /// to a local deferred chain, as does any segment the bounded attempt
    /// declines.
    ///
    /// The chain is placed once at the end, so the worst case for the whole
    /// teardown is a single orphan-pool acquisition rather than one contended
    /// segment-pool acquisition per owned segment. That final placement is the
    /// one point where a wait remains possible, because a segment that cannot
    /// be unmapped has no sink but the orphan pool and the alternative is to
    /// leak it. It is preceded by a bounded attempt, and the orphan pool is
    /// touched only by teardown and by the allocation slow path, so it is far
    /// less contended than the segment pool this path used to wait on per
    /// segment.
    pub fn reclaim_owned_segments(&mut self) {
        // We must clear the current segment first before deallocating any segments,
        // to avoid a use-after-free if the current segment gets deallocated.
        // SAFETY: `set_current_segment(None)` only clears this allocator's own
        // `current_segment`/`is_current` state; no segment pointer is read.
        unsafe { self.set_current_segment(None) };

        // Segments no sink would take without waiting, linked through
        // `next_free_segment` (free of the owned-chain links cleared below) and
        // placed in one acquisition after the walk.
        let mut deferred = DeferredChain::new();

        let mut curr = self.owned_segments_head;
        while !curr.is_null() {
            // SAFETY: `curr` walks this thread's own intrusive owned-segments
            // chain (`owned_segments_head` then each `next_owned_segment`), so
            // every node is a live segment owned exclusively by this allocator.
            // `next` is captured before `curr` is deallocated or pushed to the
            // orphan pool, so the walk never dereferences a freed segment. All
            // `pages[i]` reads use `i` drawn from `page_occupied_mask` bits,
            // which index valid entries of the segment's page array.
            unsafe {
                let next = (*curr).next_owned_segment;

                let dynamic_encrypted = (*curr).free_list_encrypted;
                let mut total_allocations = 0;
                // SAFETY: `curr` is live/exclusive; `i` indexes valid occupied pages.
                for i in OccupiedPageBits::new((*curr).page_occupied_mask) {
                    // Raw pointer, not `&mut`: remote threads read page metadata
                    // during cross-thread frees — a `&mut` would alias those reads.
                    let page = reclaim_and_record(
                        curr,
                        i,
                        dynamic_encrypted,
                        &mut self.cross_thread_reclaimed,
                    );
                    total_allocations += (*page).alloc_count;
                }

                // Clear the allocator cache before the owner token so a remote
                // thread cannot observe a stale, non-null allocator on a segment
                // that has already been handed back to the global pool.
                // SAFETY: `curr` is exclusively owned by this teardown sweep.
                Segment::clear_ownership(curr);
                Segment::set_current(curr, false);
                (*curr).next_owned_segment = core::ptr::null_mut();
                (*curr).prev_owned_segment = core::ptr::null_mut();

                // An empty segment may be cached or unmapped, both bounded. One
                // still holding live allocations cannot be unmapped, so the
                // orphan pool is its only sink and it goes straight to the
                // deferred chain.
                if total_allocations > 0 || !try_deallocate_segment::<B>(curr) {
                    deferred.push(curr);
                }

                curr = next;
            }
        }

        // SAFETY: every deferred node was detached from this thread's owned
        // chain and had its owner identity cleared above, so the chain is
        // exclusively owned here and transfers wholesale to the orphan pool.
        unsafe { deferred.place_in_orphan_pool::<B>() };

        self.next_page_index = 0;
        self.owned_segments_head = core::ptr::null_mut();
        self.owned_segment_count = 0;
        self.active_pages = [None; NUM_SIZE_CLASSES];
        self.full_pages = [None; NUM_SIZE_CLASSES];
        self.empty_pages = None;
    }

    /// Tries to reclaim a segment if it has zero active allocations.
    ///
    /// # Safety
    ///
    /// Accesses and modifies page and segment lists.
    pub unsafe fn try_reclaim_segment(&mut self, segment: *mut Segment) -> bool {
        if self
            .current_segment
            .is_some_and(|current| current.as_ptr() == segment)
        {
            return false;
        }

        if self.owned_segment_count < RECLAIM_THRESHOLD_SEGMENTS {
            return false;
        }

        // SAFETY: `segment` is a live segment owned by this allocator (the
        // caller's precondition; it is not `current_segment`, checked above).
        // `OccupiedPageBits` skips bit 0; each remaining `i` is a valid index.
        unsafe {
            let dynamic_encrypted = (*segment).free_list_encrypted;
            for i in OccupiedPageBits::new((*segment).page_occupied_mask) {
                let pg = Page::page_in_segment(segment, i);
                if (*pg).alloc_count > 0 {
                    // SAFETY: live segment, valid occupied-page index, pre-read encrypted mode.
                    let pg = reclaim_and_record(
                        segment,
                        i,
                        dynamic_encrypted,
                        &mut self.cross_thread_reclaimed,
                    );
                    if (*pg).alloc_count > 0 {
                        return false;
                    }
                }
            }
        }

        // SAFETY: every occupied page of `segment` was just confirmed to have
        // zero live allocations, so detaching them from this allocator's page
        // lists and unlinking the segment from the owned chain leaves no live
        // references; both helpers operate only on this thread's own structures.
        unsafe {
            unlink_segment_pages(self, segment);
            self.unlink_owned_segment(segment);
        }

        // SAFETY: `segment` is now fully detached (no page-list or owned-list
        // membership; `unlink_owned_segment` already cleared both link fields),
        // so clearing its owner identity and returning it hands a segment with no
        // live references back to the pool exactly once.
        unsafe { detach_and_release_segment::<B>(segment) };
        true
    }

    /// Performs a defragmentation sweep over all owned segments (excluding `current_segment`),
    /// consolidating cross-thread frees, identifying empty pages, and reclaiming empty segments.
    ///
    /// # Safety
    /// Sweeps all owned segments, reclaims empty ones, and moves zero-alloc
    /// pages to the empty list.
    ///
    /// Non-generic: no `P::` const is read anywhere in this function or its
    /// callees (`unlink_segment_pages`, `detach_and_release_segment`). The
    /// `<P: AllocPolicy>` parameter was vestigial — its removal compiles this
    /// 140-line sweep once per `B` rather than once per `(P, B)`.
    ///
    /// The caller must ensure that the allocator is in a safe, non-reentrant state.
    pub unsafe fn periodic_defragmentation_sweep(&mut self) {
        let mut curr = self.owned_segments_head;
        while !curr.is_null() {
            let segment = curr;
            // SAFETY: `segment` is a live node of this thread's own
            // owned-segments chain; reading its `next_owned_segment` before any
            // teardown advances the walk without dereferencing a freed segment.
            curr = unsafe { (*segment).next_owned_segment };

            if self.is_current_segment(segment) {
                continue;
            }

            // SAFETY: `segment` is a live segment owned exclusively by this
            // thread, so reading its `free_list_encrypted` flag and
            // `page_occupied_mask` is a valid, unaliased load.
            let dynamic_encrypted = unsafe { (*segment).free_list_encrypted };
            let mut total_allocations = 0;

            // SAFETY: `segment` is owned by this thread; `OccupiedPageBits`
            // skips bit 0 and yields only set-bit indices of valid occupied pages.
            unsafe {
                for i in OccupiedPageBits::new((*segment).page_occupied_mask) {
                    let pg = &raw mut (*segment).pages[i];
                    // SAFETY: live segment, valid occupied-page index, pre-read encrypted.
                    reclaim_and_record(
                        segment,
                        i,
                        dynamic_encrypted,
                        &mut self.cross_thread_reclaimed,
                    );
                    total_allocations += (*pg).alloc_count;

                    if (*pg).alloc_count == 0 && ((*pg).list_state == 1 || (*pg).list_state == 2) {
                        let class = (*pg).size_class as usize;
                        let is_only_active =
                            crate::free_helpers::is_sole_active_page(self.active_pages[class], pg);
                        if !is_only_active {
                            let pg_ptr = NonNull::new_unchecked(pg);
                            if (*pg).list_state == 1 {
                                unlink_page_from_list_raw(
                                    pg_ptr,
                                    self.active_pages.get_unchecked_mut(class),
                                );
                            } else {
                                unlink_page_from_list_raw(
                                    pg_ptr,
                                    self.full_pages.get_unchecked_mut(class),
                                );
                            }
                            push_page_front_raw(pg_ptr, &mut self.empty_pages, 3);
                        }
                    }
                }
            }

            if total_allocations == 0 && self.owned_segment_count >= RECLAIM_THRESHOLD_SEGMENTS {
                // SAFETY: the sweep above observed zero live allocations across
                // every occupied page of `segment`, so detaching its pages and
                // unlinking it from the owned chain leaves no live references;
                // `detach_and_release_segment` then clears the owner identity and
                // returns the fully detached segment to the pool exactly once.
                // The next node was captured into `curr` before this teardown.
                unsafe {
                    unlink_segment_pages(self, segment);
                    self.unlink_owned_segment(segment);
                    detach_and_release_segment::<B>(segment);
                }
            }
        }
    }
}

/// Detaches every page of `segment` from `alloc`'s active/full/empty page lists.
///
/// # Safety
///
/// `segment` must be a live segment owned exclusively by `alloc`, and its
/// `page_linked_mask` must accurately mark the pages currently linked into
/// `alloc`'s page lists. The caller must hold exclusive access to `alloc`.
unsafe fn unlink_segment_pages<B: HasSegmentPool>(
    alloc: &mut ThreadAllocator<B>,
    segment: *mut Segment,
) {
    // SAFETY: `segment` is a live segment owned by `alloc` (caller contract).
    // `OccupiedPageBits` skips bit 0 and yields only set-bit indices of
    // pages currently linked into `alloc`'s active/full/empty lists.
    for i in OccupiedPageBits::new(unsafe { (*segment).page_linked_mask }) {
        // SAFETY: `i` is a set bit of `page_linked_mask`, indexing a valid
        // linked page of `segment`; the segment is exclusive to `alloc`.
        let pg = unsafe { &raw mut (*segment).pages[i] };
        // SAFETY: `pg` addresses a live page of `segment`.
        let state = unsafe { (*pg).list_state };
        if state == 1 || state == 2 {
            // SAFETY: list_state 1/2 means `pg` is in the active/full list.
            let class = unsafe { (*pg).size_class } as usize;
            unsafe { alloc.unlink_page(pg, class) };
        } else if state == 3 {
            // SAFETY: list_state 3 means `pg` is in the empty-page list.
            unsafe { alloc.unlink_empty_page(pg) };
        }
    }
}

/// Clears a fully-detached segment's owner identity and returns it to the pool.
///
/// This is the shared teardown tail for `try_reclaim_segment` and
/// `periodic_defragmentation_sweep`: both call it only after
/// `unlink_owned_segment` has already spliced `segment` out of the owned chain
/// (clearing *both* `prev_owned_segment` and `next_owned_segment`) and
/// `unlink_segment_pages` has detached its pages, so the links need no
/// re-clearing here — the previous per-site code redundantly re-nulled
/// `next_owned_segment` while leaving `prev_owned_segment` untouched, an
/// asymmetry this consolidation removes.
///
/// # Safety
///
/// `segment` must be a live segment already unlinked from every owned-segment
/// and page list, so that clearing its owner identity and handing it to
/// `deallocate_segment` returns a segment with no live references exactly once.
#[inline]
unsafe fn detach_and_release_segment<B: HasSegmentPool>(segment: *mut Segment) {
    // SAFETY: `segment` is the fully-detached live segment per the contract.
    unsafe {
        Segment::clear_ownership(segment);
        deallocate_segment::<B>(segment);
    }
}

/// Drains remote frees from page `page_index` of `segment` and accumulates the
/// reclaimed count into `cross_thread_sink`.
///
/// This is the **SSOT** for the 6-line reclaim core repeated in
/// `reclaim_owned_segments`, `try_reclaim_segment`, and
/// `periodic_defragmentation_sweep`. Non-generic and non-method so it
/// compiles once regardless of the backend `B`.
///
/// # Safety
///
/// `segment` must be a live segment exclusively owned by the calling
/// allocator; `page_index` must be a non-zero set bit of
/// `segment.page_occupied_mask`; `encrypted` must equal
/// `segment.free_list_encrypted`.
///
/// Returns the raw page pointer so callers can read `alloc_count` or perform
/// list-management afterward.
#[inline(always)]
unsafe fn reclaim_and_record(
    segment: *mut Segment,
    page_index: usize,
    encrypted: bool,
    cross_thread_sink: &mut usize,
) -> *mut Page {
    // SAFETY: live segment + in-range index per contract.
    let page = unsafe { Page::page_in_segment(segment, page_index) };
    // SAFETY: `page` is the initialized metadata projected above.
    let randomized = unsafe { (*page).secondary_free.is_some() };
    // SAFETY: valid segment/index/encrypted triple per contract.
    let reclaimed = unsafe {
        Page::reclaim_thread_free_if_present_in_segment_with_randomized(
            segment, page_index, encrypted, randomized,
        )
    };
    if reclaimed > 0 {
        *cross_thread_sink += reclaimed;
    }
    page
}