moa_memblock 0.1.4

//! Free 区域迭代器

use crate::{
    flags::MemblockFlags,
    region::{MemblockRegion, MemblockType},
};

/// 迭代 free 区域（memory 减去 reserved）
///
/// 使用双指针遍历 memory 和 reserved 数组，产出差集。
/// 可通过 `skip_flags` 过滤掉带有指定 flags 的 memory 区域。
pub struct FreeRegionIter<'a> {
    memory: &'a MemblockType,
    reserved: &'a MemblockType,
    skip_flags: MemblockFlags,
    mem_idx: usize,
    rsv_idx: usize,
    cur_start: usize,
}

impl<'a> FreeRegionIter<'a> {
    #[moa_sec_macros::init]
    pub(crate) fn new(
        memory: &'a MemblockType,
        reserved: &'a MemblockType,
        skip_flags: MemblockFlags,
    ) -> Self {
        let cur_start = memory.regions().first().map_or(0, |r| r.base);
        Self { memory, reserved, skip_flags, mem_idx: 0, rsv_idx: 0, cur_start }
    }
}

impl Iterator for FreeRegionIter<'_> {
    /// (base, size)
    type Item = (usize, usize);

    #[moa_sec_macros::init]
    fn next(&mut self) -> Option<Self::Item> {
        let mem_regions = self.memory.regions();
        let rsv_regions = self.reserved.regions();

        while self.mem_idx < mem_regions.len() {
            let m = &mem_regions[self.mem_idx];
            let m_end = m.end();

            if !self.skip_flags.is_empty() && m.flags.intersects(self.skip_flags) {
                self.mem_idx += 1;
                if self.mem_idx < mem_regions.len() {
                    self.cur_start = mem_regions[self.mem_idx].base;
                }
                continue;
            }

            if self.cur_start < m.base {
                self.cur_start = m.base;
            }

            if self.cur_start >= m_end {
                self.mem_idx += 1;
                if self.mem_idx < mem_regions.len() {
                    self.cur_start = mem_regions[self.mem_idx].base;
                }
                continue;
            }

            while self.rsv_idx < rsv_regions.len()
                && rsv_regions[self.rsv_idx].end() <= self.cur_start
            {
                self.rsv_idx += 1;
            }

            if self.rsv_idx >= rsv_regions.len() || rsv_regions[self.rsv_idx].base >= m_end {
                let result = (self.cur_start, m_end - self.cur_start);
                self.mem_idx += 1;
                if self.mem_idx < mem_regions.len() {
                    self.cur_start = mem_regions[self.mem_idx].base;
                }
                return Some(result);
            }

            let r = &rsv_regions[self.rsv_idx];

            if self.cur_start < r.base {
                let gap_end = r.base.min(m_end);
                let result = (self.cur_start, gap_end - self.cur_start);
                self.cur_start = r.end();
                return Some(result);
            }

            self.cur_start = r.end();
            self.rsv_idx += 1;
        }

        None
    }
}

/// 反向迭代 free 区域（从高地址到低地址）
pub struct FreeRegionRevIter<'a> {
    memory: &'a MemblockType,
    reserved: &'a MemblockType,
    skip_flags: MemblockFlags,
    mem_idx: Option<usize>,
    rsv_idx: Option<usize>,
    cur_end: usize,
}

impl<'a> FreeRegionRevIter<'a> {
    #[moa_sec_macros::init]
    pub(crate) fn new(
        memory: &'a MemblockType,
        reserved: &'a MemblockType,
        skip_flags: MemblockFlags,
    ) -> Self {
        let mem_idx = memory.cnt().checked_sub(1);
        let rsv_idx = reserved.cnt().checked_sub(1);
        let cur_end = memory.regions().last().map_or(0, MemblockRegion::end);
        Self { memory, reserved, skip_flags, mem_idx, rsv_idx, cur_end }
    }
}

impl Iterator for FreeRegionRevIter<'_> {
    /// (base, size)
    type Item = (usize, usize);

    #[moa_sec_macros::init]
    fn next(&mut self) -> Option<Self::Item> {
        let mem_regions = self.memory.regions();
        let rsv_regions = self.reserved.regions();

        while let Some(mi) = self.mem_idx {
            let m = &mem_regions[mi];
            let m_base = m.base;

            if !self.skip_flags.is_empty() && m.flags.intersects(self.skip_flags) {
                self.mem_idx = mi.checked_sub(1);
                if let Some(prev) = self.mem_idx {
                    self.cur_end = mem_regions[prev].end();
                }
                continue;
            }

            if self.cur_end > m.end() {
                self.cur_end = m.end();
            }

            if self.cur_end <= m_base {
                self.mem_idx = mi.checked_sub(1);
                if let Some(prev) = self.mem_idx {
                    self.cur_end = mem_regions[prev].end();
                }
                continue;
            }

            while let Some(ri) = self.rsv_idx {
                if rsv_regions[ri].base < self.cur_end {
                    break;
                }
                self.rsv_idx = ri.checked_sub(1);
            }

            let rsv_end_before = self.rsv_idx.is_none_or(|ri| rsv_regions[ri].end() <= m_base);

            if rsv_end_before {
                let result = (m_base, self.cur_end - m_base);
                self.mem_idx = mi.checked_sub(1);
                if let Some(prev) = self.mem_idx {
                    self.cur_end = mem_regions[prev].end();
                }
                return Some(result);
            }

            let ri = self.rsv_idx.unwrap();
            let r = &rsv_regions[ri];

            if r.end() < self.cur_end {
                let gap_start = r.end().max(m_base);
                let result = (gap_start, self.cur_end - gap_start);
                self.cur_end = r.base;
                self.rsv_idx = ri.checked_sub(1);
                return Some(result);
            }

            self.cur_end = r.base;
            self.rsv_idx = ri.checked_sub(1);
        }

        None
    }
}