moa_memblock 0.1.3

//! 顶层 Memblock 分配器

use crate::{
    MemblockError, MemblockFlags,
    iter::{FreeRegionIter, FreeRegionRevIter},
    region::{MemblockRegion, MemblockType},
};

/// Linux 风格 memblock 早期引导内存分配器
///
/// 管理两个区域列表：
/// - `memory`：可用物理内存
/// - `reserved`：已保留/已分配区域
///
/// free 区域 = memory ∩ ¬reserved
pub struct Memblock {
    memory: MemblockType,
    reserved: MemblockType,
    bottom_up: bool,
    current_limit: usize,
}

impl Default for Memblock {
    #[moa_sec_macros::init]
    fn default() -> Self {
        Self::new()
    }
}

impl Memblock {
    /// 创建空的 memblock
    pub const fn new() -> Self {
        Self {
            memory: MemblockType::new("memory"),
            reserved: MemblockType::new("reserved"),
            bottom_up: false,
            current_limit: usize::MAX,
        }
    }

    /// 添加可用内存区域
    #[moa_sec_macros::init]
    pub fn add(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        self.memory.add_range(base, size, MemblockFlags::empty())
    }

    /// 添加可用内存区域（带 flags）
    #[moa_sec_macros::init]
    pub fn add_flags(
        &mut self,
        base: usize,
        size: usize,
        flags: MemblockFlags,
    ) -> Result<(), MemblockError> {
        self.memory.add_range(base, size, flags)
    }

    /// 从可用内存中移除区域
    #[moa_sec_macros::init]
    pub fn remove(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        self.memory.remove_range(base, size)
    }

    /// 标记区域为已保留
    #[moa_sec_macros::init]
    pub fn reserve(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        self.reserved.add_range(base, size, MemblockFlags::empty())
    }

    /// 释放已保留区域
    #[moa_sec_macros::init]
    pub fn free(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        self.reserved.remove_range(base, size)
    }

    /// 标记 NOMAP
    #[moa_sec_macros::init]
    pub fn mark_nomap(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        self.memory.set_flags(base, size, MemblockFlags::NOMAP)
    }

    /// 清除 NOMAP
    #[moa_sec_macros::init]
    pub fn clear_nomap(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        self.memory.clear_flags(base, size, MemblockFlags::NOMAP)
    }

    /// 分配 `size` 字节，对齐到 `align`
    ///
    /// 在 `[0, current_limit)` 范围内查找 free 区域。
    /// 返回分配的物理地址，同时自动 reserve。
    #[moa_sec_macros::init]
    pub fn alloc(&mut self, size: usize, align: usize) -> Option<usize> {
        self.alloc_range(size, align, 0, self.current_limit)
    }

    /// 在 `[min_addr, max_addr)` 范围内分配
    #[moa_sec_macros::init]
    pub fn alloc_range(
        &mut self,
        size: usize,
        align: usize,
        min_addr: usize,
        max_addr: usize,
    ) -> Option<usize> {
        if size == 0 {
            return None;
        }
        let align = if align == 0 { 1 } else { align };
        let max_addr = max_addr.min(self.current_limit);

        let found = if self.bottom_up {
            self.find_bottom_up(size, align, min_addr, max_addr)
                .or_else(|| self.find_top_down(size, align, min_addr, max_addr))
        } else {
            self.find_top_down(size, align, min_addr, max_addr)
        };

        if let Some(addr) = found {
            self.reserve(addr, size).ok()?;
        }

        found
    }

    /// 从低地址向高地址查找
    #[moa_sec_macros::init]
    fn find_bottom_up(
        &self,
        size: usize,
        align: usize,
        min_addr: usize,
        max_addr: usize,
    ) -> Option<usize> {
        let skip = MemblockFlags::NOMAP;
        for (base, region_size) in FreeRegionIter::new(&self.memory, &self.reserved, skip) {
            let start = base.max(min_addr);
            let aligned = align_up(start, align);
            let end = base + region_size;

            if aligned >= end {
                continue;
            }
            if end.min(max_addr) - aligned >= size {
                return Some(aligned);
            }
        }
        None
    }

    /// 从高地址向低地址查找
    #[moa_sec_macros::init]
    fn find_top_down(
        &self,
        size: usize,
        align: usize,
        min_addr: usize,
        max_addr: usize,
    ) -> Option<usize> {
        let skip = MemblockFlags::NOMAP;
        for (base, region_size) in FreeRegionRevIter::new(&self.memory, &self.reserved, skip) {
            let end = (base + region_size).min(max_addr);
            if end < size {
                continue;
            }
            let aligned = align_down(end - size, align);

            if aligned < base || aligned < min_addr {
                continue;
            }
            return Some(aligned);
        }
        None
    }

    /// 地址是否在 memory 区域内
    #[inline(always)]
    pub fn is_memory(&self, addr: usize) -> bool {
        self.memory.find(addr).is_some()
    }

    /// 地址是否已 reserved
    #[inline(always)]
    pub fn is_reserved(&self, addr: usize) -> bool {
        self.reserved.find(addr).is_some()
    }

    /// 迭代 free 区域（memory - reserved，跳过 NOMAP）
    #[moa_sec_macros::init]
    pub fn free_regions(&self) -> FreeRegionIter<'_> {
        FreeRegionIter::new(&self.memory, &self.reserved, MemblockFlags::NOMAP)
    }

    /// 反向迭代 free 区域（从高地址到低地址，跳过 NOMAP）
    #[moa_sec_macros::init]
    pub fn free_regions_rev(&self) -> FreeRegionRevIter<'_> {
        FreeRegionRevIter::new(&self.memory, &self.reserved, MemblockFlags::NOMAP)
    }

    /// 第一个物理内存区域的起始地址
    #[inline(always)]
    pub fn phys_mem_start(&self) -> Option<usize> {
        self.memory.regions().first().map(|r| r.base)
    }

    /// 最后一个物理内存区域的结束地址
    #[inline(always)]
    pub fn phys_mem_end(&self) -> Option<usize> {
        self.memory.regions().last().map(MemblockRegion::end)
    }

    /// 物理内存总大小
    #[inline(always)]
    pub fn phys_mem_size(&self) -> usize {
        self.memory.total_size()
    }

    /// 设置分配方向
    #[moa_sec_macros::init]
    pub fn set_bottom_up(&mut self, bottom_up: bool) {
        self.bottom_up = bottom_up;
    }

    /// 获取分配方向
    #[moa_sec_macros::init]
    pub fn is_bottom_up(&self) -> bool {
        self.bottom_up
    }

    /// 设置分配上限
    #[moa_sec_macros::init]
    pub fn set_current_limit(&mut self, limit: usize) {
        self.current_limit = limit;
    }

    /// 获取分配上限
    #[moa_sec_macros::init]
    pub fn current_limit(&self) -> usize {
        self.current_limit
    }

    /// 访问 memory 区域
    #[inline(always)]
    pub fn memory(&self) -> &MemblockType {
        &self.memory
    }

    /// 访问 reserved 区域
    #[inline(always)]
    pub fn reserved(&self) -> &MemblockType {
        &self.reserved
    }
}

#[inline(always)]
fn align_up(addr: usize, align: usize) -> usize {
    debug_assert!(align.is_power_of_two());
    (addr + align - 1) & !(align - 1)
}

#[inline(always)]
fn align_down(addr: usize, align: usize) -> usize {
    debug_assert!(align.is_power_of_two());
    addr & !(align - 1)
}