moa_memblock 0.1.3

//! 内存区域及排序区域数组

use crate::{MemblockError, MemblockFlags};

/// 每种类型的最大区域数
pub const INIT_MEMBLOCK_REGIONS: usize = 32;

/// 单个连续内存区域
#[derive(Clone, Copy, Debug)]
pub struct MemblockRegion {
    /// 起始物理地址
    pub base: usize,
    /// 区域大小（字节）
    pub size: usize,
    /// 属性标志
    pub flags: MemblockFlags,
}

impl MemblockRegion {
    const EMPTY: Self = Self { base: 0, size: 0, flags: MemblockFlags::empty() };

    /// 区域结束地址（不包含）
    #[inline(always)]
    pub fn end(&self) -> usize {
        self.base + self.size
    }
}

/// 排序、无重叠的区域数组
///
/// 维护不变量：
/// - 区域按 `base` 升序排列
/// - 同 flags 的相邻区域自动合并
/// - 区域之间无重叠
pub struct MemblockType {
    cnt: usize,
    total_size: usize,
    regions: [MemblockRegion; INIT_MEMBLOCK_REGIONS],
    name: &'static str,
}

impl MemblockType {
    /// 创建空的区域数组
    pub const fn new(name: &'static str) -> Self {
        Self {
            cnt: 0,
            total_size: 0,
            regions: [MemblockRegion::EMPTY; INIT_MEMBLOCK_REGIONS],
            name,
        }
    }

    /// 活跃区域数量
    #[inline(always)]
    pub fn cnt(&self) -> usize {
        self.cnt
    }

    /// 所有区域大小之和
    #[inline(always)]
    pub fn total_size(&self) -> usize {
        self.total_size
    }

    /// 名称
    #[inline(always)]
    pub fn name(&self) -> &str {
        self.name
    }

    /// 获取区域切片
    #[inline(always)]
    pub fn regions(&self) -> &[MemblockRegion] {
        &self.regions[..self.cnt]
    }

    /// 添加区域，处理重叠并自动合并
    #[moa_sec_macros::init]
    pub(crate) fn add_range(
        &mut self,
        base: usize,
        size: usize,
        flags: MemblockFlags,
    ) -> Result<(), MemblockError> {
        if size == 0 {
            return Err(MemblockError::InvalidRange);
        }
        let end = base.checked_add(size).ok_or(MemblockError::InvalidRange)?;

        if self.cnt == 0 {
            return self.insert_region(0, base, size, flags);
        }

        let mut cur_base = base;
        let mut i = 0;
        while i < self.cnt && cur_base < end {
            let r_base = self.regions[i].base;
            let r_end = self.regions[i].end();

            if cur_base < r_base {
                let insert_end = end.min(r_base);
                self.insert_region(i, cur_base, insert_end - cur_base, flags)?;
                i += 1; // 跳过刚插入的
            }

            if cur_base < r_end {
                cur_base = r_end;
            }
            i += 1;
        }

        if cur_base < end {
            self.insert_region(self.cnt, cur_base, end - cur_base, flags)?;
        }

        self.merge();
        Ok(())
    }

    /// 移除指定范围
    #[moa_sec_macros::init]
    pub(crate) fn remove_range(&mut self, base: usize, size: usize) -> Result<(), MemblockError> {
        if size == 0 {
            return Err(MemblockError::InvalidRange);
        }
        let end = base.checked_add(size).ok_or(MemblockError::InvalidRange)?;

        let (start_idx, end_idx) = self.isolate_range(base, end)?;

        for i in (start_idx..end_idx).rev() {
            self.remove_region(i);
        }

        Ok(())
    }

    /// 在边界处分割区域，返回 [base, end) 覆盖的区域索引范围
    #[moa_sec_macros::init]
    fn isolate_range(&mut self, base: usize, end: usize) -> Result<(usize, usize), MemblockError> {
        let mut start_idx: Option<usize> = None;
        let mut end_idx = 0;

        let mut i = 0;
        while i < self.cnt {
            let r_base = self.regions[i].base;
            let r_end = self.regions[i].end();

            if r_end <= base {
                i += 1;
                continue;
            }
            if r_base >= end {
                break;
            }

            if r_base < base {
                let flags = self.regions[i].flags;
                let new_size = base - r_base;
                self.total_size -= self.regions[i].size - new_size;
                self.regions[i].size = new_size;
                self.insert_region(i + 1, base, r_end - base, flags)?;
                i += 1;
            }

            let r_end = self.regions[i].end();
            if r_end > end {
                let flags = self.regions[i].flags;
                let new_size = end - self.regions[i].base;
                self.total_size -= self.regions[i].size - new_size;
                self.regions[i].size = new_size;
                self.insert_region(i + 1, end, r_end - end, flags)?;
            }

            if start_idx.is_none() {
                start_idx = Some(i);
            }
            end_idx = i + 1;
            i += 1;
        }

        Ok((start_idx.unwrap_or(end_idx), end_idx))
    }

    /// 设置指定范围的 flags
    #[moa_sec_macros::init]
    pub(crate) fn set_flags(
        &mut self,
        base: usize,
        size: usize,
        flags: MemblockFlags,
    ) -> Result<(), MemblockError> {
        if size == 0 {
            return Err(MemblockError::InvalidRange);
        }
        let end = base.checked_add(size).ok_or(MemblockError::InvalidRange)?;

        let (start_idx, end_idx) = self.isolate_range(base, end)?;
        for r in &mut self.regions[start_idx..end_idx] {
            r.flags |= flags;
        }
        self.merge();
        Ok(())
    }

    /// 清除指定范围的 flags
    #[moa_sec_macros::init]
    pub(crate) fn clear_flags(
        &mut self,
        base: usize,
        size: usize,
        flags: MemblockFlags,
    ) -> Result<(), MemblockError> {
        if size == 0 {
            return Err(MemblockError::InvalidRange);
        }
        let end = base.checked_add(size).ok_or(MemblockError::InvalidRange)?;

        let (start_idx, end_idx) = self.isolate_range(base, end)?;
        for r in &mut self.regions[start_idx..end_idx] {
            r.flags &= !flags;
        }
        self.merge();
        Ok(())
    }

    /// 二分查找包含指定地址的区域索引
    pub fn find(&self, addr: usize) -> Option<usize> {
        if self.cnt == 0 {
            return None;
        }
        let mut lo = 0;
        let mut hi = self.cnt;
        while lo < hi {
            let mid = lo + (hi - lo) / 2;
            if self.regions[mid].base <= addr {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        if lo == 0 {
            return None;
        }
        let idx = lo - 1;
        if addr < self.regions[idx].end() { Some(idx) } else { None }
    }

    /// 在 index 处插入区域
    #[moa_sec_macros::init]
    fn insert_region(
        &mut self,
        index: usize,
        base: usize,
        size: usize,
        flags: MemblockFlags,
    ) -> Result<(), MemblockError> {
        if self.cnt >= INIT_MEMBLOCK_REGIONS {
            return Err(MemblockError::RegionsFull);
        }
        self.regions.copy_within(index..self.cnt, index + 1);
        self.regions[index] = MemblockRegion { base, size, flags };
        self.cnt += 1;
        self.total_size += size;
        Ok(())
    }

    /// 删除 index 处的区域
    #[moa_sec_macros::init]
    fn remove_region(&mut self, index: usize) {
        self.total_size -= self.regions[index].size;
        self.regions.copy_within(index + 1..self.cnt, index);
        self.cnt -= 1;
        self.regions[self.cnt] = MemblockRegion::EMPTY;
    }

    /// 合并相邻且 flags 相同的区域
    #[moa_sec_macros::init]
    fn merge(&mut self) {
        let mut i = 0;
        while i + 1 < self.cnt {
            let this = &self.regions[i];
            let next = &self.regions[i + 1];

            if this.end() == next.base && this.flags == next.flags {
                self.regions[i].size += next.size;
                self.remove_region(i + 1);
            } else {
                i += 1;
            }
        }
    }
}