//! # Per-Actor Buddy Allocator  (`src/mem/buddy.rs`)
//!
//! A power-of-two page-block allocator for a fixed, pre-committed physical
//! region.  Designed for **single-consumer** use inside a [`crate::mem::pm_manages::PMActor`]:
//! the owning actor is the only entity that ever mutates its buddy allocator, so
//! no locking is needed — the entire subsystem is inherently lock-free from the
//! consumer's perspective.
//!
//! ## Block Layout
//!
//! The region is subdivided into blocks of `2^order` pages (order 0 … MAX_ORDER).
//! Every block at order *k* is aligned to `2^k` pages within the region.
//!
//! ```text
//! order  0 →      1 page  =     4 KiB
//! order  1 →      2 pages =     8 KiB
//! …
//! order 11 →  2 048 pages =     8 MiB   (MAX_ORDER)
//! ```
//!
//! ## Complexity
//!
//! | Operation | Amortised | Worst-case |
//! |-----------|-----------|------------|
//! | `alloc`   | O(log N)  | O(MAX_ORDER · Lₖ) |
//! | `free`    | O(log N)  | O(MAX_ORDER · Lₖ) |
//!
//! where Lₖ = `free_lists[k].len()` ≤ `total_pages / 2^k`.
//!
//! ## Production note
//!
//! For managed ranges > 1 GiB (> 256 K pages), replace the per-order `Vec<usize>`
//! free lists with a radix tree or interval tree to bound Lₖ.  For current
//! PMActor capital sizes (typically ≤ 128 MiB = 32 K pages at order 0), the
//! `Vec`-based implementation is fully adequate.

extern crate alloc;
use alloc::vec::Vec;

// Constants & helpers
/// Maximum allocation order.
/// `2^11 × 4 096 bytes = 8 MiB` per single allocation.
pub const MAX_ORDER: usize = 11;

/// ⌈log2(n)⌉ — the minimum order whose block size covers `page_count` pages.
///
/// ```text
/// order_for(1)    = 0   (2^0 = 1)
/// order_for(2)    = 1   (2^1 = 2)
/// order_for(3)    = 2   (2^2 = 4 ≥ 3)
/// order_for(2048) = 11
/// ```
#[inline]
pub fn order_for(page_count: usize) -> usize {
    if page_count <= 1 {
        0
    } else {
        // Number of bits needed to represent (page_count - 1).
        usize::BITS as usize - (page_count - 1).leading_zeros() as usize
    }
}

// BuddyAllocator
/// Buddy allocator over a contiguous, pre-committed range of physical pages.
///
/// All page indices stored in free lists are **relative to the start of the
/// managed range**.  The owner (`PMActor`) translates to absolute `PhysAddr`
/// by adding `managed_range.0`.
///
/// # Invariants
/// - `free_pages ≤ total_pages` at all times.
/// - Every block in `free_lists[k]` is aligned to `2^k` pages (i.e.
///   `idx % (1 << k) == 0`).
/// - No block appears in more than one order's free list simultaneously.
pub struct BuddyAllocator {
    /// `free_lists[k]` = relative page indices of free 2^k-page blocks.
    /// Ordering within each list is irrelevant; `pop()` / `swap_remove()` are used.
    free_lists: [Vec<usize>; MAX_ORDER + 1],
    /// Total pages in the managed range (need not be a power of two).
    total_pages: usize,
    /// Running count of free pages.  Always equals `Σ (2^k × free_lists[k].len())`.
    free_pages: usize,
}

impl BuddyAllocator {
    //Construction
    /// Create a new allocator over `total_pages` pages, **all initially free**.
    ///
    /// Uses a greedy largest-first decomposition to build the initial free lists
    /// in O(total_pages / 2^MAX_ORDER) iterations — essentially O(1) for
    /// power-of-two sizes.
    ///
    /// Example: 7 pages → blocks [4, 2, 1] → free_lists[2]=[0], [1]=[4], [0]=[6].
    pub fn new(total_pages: usize) -> Self {
        // core::array::from_fn is the idiomatic way to init a non-Copy array.
        let free_lists: [Vec<usize>; MAX_ORDER + 1] =
            core::array::from_fn(|_| Vec::new());

        let mut this = Self {
            free_lists,
            total_pages,
            free_pages: 0,
        };

        if total_pages == 0 {
            return this;
        }

        let mut idx = 0usize;
        while idx < total_pages {
            let remaining = total_pages - idx;

            // Alignment constraint: block at `idx` must be aligned to 2^order.
            // trailing_zeros(0) is u32::MAX, so we clamp to MAX_ORDER.
            let align_order = if idx == 0 {
                MAX_ORDER
            } else {
                (idx.trailing_zeros() as usize).min(MAX_ORDER)
            };

            // Size constraint: 2^order ≤ remaining  →  order ≤ ⌊log₂(remaining)⌋.
            let size_order = (usize::BITS as usize - 1)
                - remaining.leading_zeros() as usize; // ⌊log2(remaining)⌋

            let order      = MAX_ORDER.min(align_order).min(size_order);
            let block_size = 1usize << order;

            this.free_lists[order].push(idx);
            this.free_pages += block_size;
            idx += block_size;
        }

        this
    }

    //Allocation
    /// Allocate a 2^`order`-page block.
    ///
    /// Returns the **relative** page index of the block's first page, or `None`
    /// if insufficient contiguous memory remains.
    ///
    /// The returned index can be converted to a physical address:
    /// ```text
    /// phys = actor.managed_range.0.0 + (idx as u64) * PAGE_SIZE
    /// ```
    pub fn alloc(&mut self, order: usize) -> Option<usize> {
        if order > MAX_ORDER {
            return None;
        }

        // Find the smallest available order ≥ requested.
        let found_order = (order..=MAX_ORDER)
            .find(|&o| !self.free_lists[o].is_empty())?;

        // Consume one block from the found order.
        let block_idx = self.free_lists[found_order]
            .pop()
            .expect("buddy: free_list non-empty but pop() returned None");

        // Debit: we removed a 2^found_order block from free.
        self.free_pages -= 1 << found_order;

        // Split down to the requested order, putting buddies back into free lists.
        //
        // Invariant at each iteration:
        //   `block_idx` is the lower half of a 2^cur_order block.
        //   The upper half (= `block_idx + 2^(cur_order-1)`) is returned to the
        //   free list as an order-(cur_order-1) block.
        let mut cur_order = found_order;
        while cur_order > order {
            cur_order -= 1;
            let buddy_idx = block_idx + (1 << cur_order);
            self.free_lists[cur_order].push(buddy_idx);
            self.free_pages += 1 << cur_order;
        }

        // Net accounting: removed 2^found_order, added (2^found_order − 2^order).
        // Result: free_pages decreased by exactly 2^order. ✓
        Some(block_idx)
    }

    /// Allocate `page_count` pages, rounding up to the nearest power of two.
    ///
    /// Returns `(relative_page_idx, actual_order)`.  The caller **must** pass the
    /// same `order` to [`BuddyAllocator::free`] — mismatched orders corrupt the
    /// allocator.
    ///
    /// The wasted "rounding up" pages remain unusable until the block is freed.
    /// For tightly-packed allocations, callers should use `alloc(order)` directly.
    pub fn alloc_pages(&mut self, page_count: usize) -> Option<(usize, usize)> {
        if page_count == 0 {
            return None;
        }
        let order = order_for(page_count);
        if order > MAX_ORDER {
            return None;
        }
        self.alloc(order).map(|idx| (idx, order))
    }

    //Deallocation

    /// Return a 2^`order`-page block at **relative** index `block_idx` to the
    /// free pool, coalescing with free buddies up the order chain.
    ///
    /// # Panics (debug builds only)
    /// - `order > MAX_ORDER`
    /// - `block_idx + 2^order > total_pages`
    ///
    /// In release builds the checks are elided for performance; passing incorrect
    /// arguments causes undefined bookkeeping, not UB (no unsafe indexing).
    pub fn free(&mut self, mut block_idx: usize, mut order: usize) {
        debug_assert!(
            order <= MAX_ORDER,
            "buddy::free: order {} > MAX_ORDER {}",
            order,
            MAX_ORDER
        );
        debug_assert!(
            block_idx.checked_add(1 << order).map_or(false, |e| e <= self.total_pages),
            "buddy::free: block [{}, +{}] out of range (total_pages={})",
            block_idx,
            1usize << order,
            self.total_pages
        );

        // Try to coalesce with our buddy at each order.
        //
        // The buddy of a block at relative index `i` of order `k` is at:
        //   buddy_idx = i XOR 2^k
        // This works because aligned buddy pairs always differ in exactly bit k.
        while order < MAX_ORDER {
            let buddy_idx = block_idx ^ (1 << order);

            // Buddy must lie entirely within the managed range.
            let buddy_end = match buddy_idx.checked_add(1 << order) {
                Some(e) => e,
                None    => break,
            };
            if buddy_end > self.total_pages {
                break;
            }

            // Search for the buddy in the free list for this order.
            // `swap_remove` is O(1) and order-preserving is not required.
            if let Some(pos) = self.free_lists[order].iter().position(|&b| b == buddy_idx) {
                self.free_lists[order].swap_remove(pos);
                self.free_pages -= 1 << order; // buddy was in free count; remove it

                // Merged block starts at the lower address of the two.
                block_idx = block_idx.min(buddy_idx);
                order += 1;
                // Continue trying to merge at the next level.
            } else {
                break; // buddy is allocated or out-of-range; stop coalescing
            }
        }

        // Push the (possibly merged) block onto the appropriate free list.
        self.free_pages += 1 << order;
        self.free_lists[order].push(block_idx);
    }

    //Introspection

    /// Number of pages currently available for allocation.
    #[inline]
    pub fn free_pages(&self) -> usize {
        self.free_pages
    }

    /// Total pages in the managed range.
    #[inline]
    pub fn total_pages(&self) -> usize {
        self.total_pages
    }

    /// `true` if the allocator has no free pages.
    #[inline]
    pub fn is_exhausted(&self) -> bool {
        self.free_pages == 0
    }

    /// Dump free-list statistics for each order (useful in panic handlers).
    ///
    /// Returns an array `[(order, free_block_count); MAX_ORDER + 1]`.
    pub fn stats(&self) -> [(usize, usize); MAX_ORDER + 1] {
        core::array::from_fn(|o| (o, self.free_lists[o].len()))
    }
}