Elyz/kernel/src/mem/pmm.rs

use crate::mem::address::PhysAddr;
use crate::mem::allocator::Locked;

pub const PAGE_SIZE: u64 = 4096;

pub struct BitmapPMM {
    bitmap:      &'static mut [u8],
    ref_counts:  &'static mut [u16],
    total_pages: usize,
    used_pages:  usize,
    last_byte:   usize,
}

pub static PMM: Locked<Option<BitmapPMM>> = Locked::new(None);

impl BitmapPMM {
    #[allow(dead_code)]
    pub fn used_pages(&self)  -> usize { self.used_pages }
    #[allow(dead_code)]
    pub fn total_pages(&self) -> usize { self.total_pages }
    #[allow(dead_code)]
    pub fn free_pages(&self)  -> usize { self.total_pages.saturating_sub(self.used_pages) }

    pub unsafe fn init(mmap: &limine::response::MemoryMapResponse, hhdm_offset: u64) {
        let max_addr = mmap.entries().iter()
            .map(|e| e.base + e.length)
            .max()
            .unwrap_or(0);

        let total_pages = (max_addr / PAGE_SIZE) as usize;

        let bitmap_size = total_pages.div_ceil(8);
        let ref_counts_size = total_pages * core::mem::size_of::<u16>();
        let total_meta_size = bitmap_size + ref_counts_size;

        let meta_phys = mmap.entries().iter()
            .find(|e| {
                e.entry_type == limine::memory_map::EntryType::USABLE
                    && e.length >= total_meta_size as u64
            })
            .map(|e| e.base)
            .expect("PMM: no usable region large enough for metadata");

        let bitmap_ptr = (meta_phys + hhdm_offset) as *mut u8;
        let ref_counts_ptr = (meta_phys + hhdm_offset + bitmap_size as u64) as *mut u16;

        let bitmap = unsafe { core::slice::from_raw_parts_mut(bitmap_ptr, bitmap_size) };
        bitmap.fill(0xFF);

        let ref_counts = unsafe { core::slice::from_raw_parts_mut(ref_counts_ptr, total_pages) };
        ref_counts.fill(1);

        let mut pmm = Self {
            bitmap,
            ref_counts,
            total_pages,
            used_pages: total_pages,
            last_byte:  0,
        };

        // Free all USABLE pages (gap pages not covered by any entry stay locked).
        for entry in mmap.entries() {
            if entry.entry_type == limine::memory_map::EntryType::USABLE {
                for addr in (entry.base..entry.base + entry.length).step_by(PAGE_SIZE as usize) {
                    pmm.free_frame(PhysAddr(addr));
                }
            }
        }

        // Metadata pages sit inside a USABLE region — lock them back.
        let meta_end = (meta_phys + total_meta_size as u64 + PAGE_SIZE - 1) & !(PAGE_SIZE - 1);
        for addr in (meta_phys..meta_end).step_by(PAGE_SIZE as usize) {
            pmm.lock_frame(PhysAddr(addr));
        }

        pmm.lock_frame(PhysAddr(0));

        *PMM.lock() = Some(pmm);
    }

    //Core operations

    pub fn free_frame(&mut self, phys_addr: PhysAddr) {
        let idx = (phys_addr.0 / PAGE_SIZE) as usize;
        if idx >= self.total_pages { return; }

        let byte = idx / 8;
        let bit  = idx % 8;

        if self.bitmap[byte] & (1 << bit) != 0 {
            self.ref_counts[idx] = self.ref_counts[idx].saturating_sub(1);

            if self.ref_counts[idx] == 0 {
                self.bitmap[byte] &= !(1 << bit);
                self.used_pages  -= 1;
                if byte < self.last_byte { self.last_byte = byte; }
            }
        }
    }

    pub fn lock_frame(&mut self, phys_addr: PhysAddr) {
        let idx = (phys_addr.0 / PAGE_SIZE) as usize;
        if idx >= self.total_pages { return; }

        let byte = idx / 8;
        let bit  = idx % 8;

        if self.bitmap[byte] & (1 << bit) == 0 {
            self.bitmap[byte] |= 1 << bit;
            self.ref_counts[idx] = 1;
            self.used_pages   += 1;
        } else if self.ref_counts[idx] == 0 {
            self.ref_counts[idx] = 1;
        }
    }

    pub fn inc_ref_frame(&mut self, phys_addr: PhysAddr) {
        let idx = (phys_addr.0 / PAGE_SIZE) as usize;
        if idx >= self.total_pages { return; }

        let byte = idx / 8;
        let bit  = idx % 8;

        if self.bitmap[byte] & (1 << bit) != 0 {
            self.ref_counts[idx] = self.ref_counts[idx].saturating_add(1);
        }
    }

    pub fn alloc_frame(&mut self) -> Option<PhysAddr> {
        let len = self.bitmap.len();

        for pass in 0..2usize {
            let (from, to) = if pass == 0 {
                (self.last_byte, len)
            } else {
                (0, self.last_byte)
            };

            for byte_idx in from..to {
                if self.bitmap[byte_idx] == 0xFF { continue; }

                for bit in 0..8u8 {
                    if self.bitmap[byte_idx] & (1 << bit) == 0 {
                        let page_idx = byte_idx * 8 + bit as usize;
                        if page_idx >= self.total_pages { return None; }

                        self.bitmap[byte_idx] |= 1 << bit;
                        self.ref_counts[page_idx] = 1;
                        self.used_pages       += 1;
                        self.last_byte         = byte_idx;

                        return Some(PhysAddr(page_idx as u64 * PAGE_SIZE));
                    }
                }
            }
        }

        None
    }

    pub fn alloc_contiguous(&mut self, count: usize) -> Option<PhysAddr> {
        if count == 0 { return None; }

        let mut run_start = 0usize;
        let mut run_len   = 0usize;

        for page_idx in 0..self.total_pages {
            let byte = page_idx / 8;
            let bit  = page_idx % 8;

            if self.bitmap[byte] & (1 << bit) == 0 {
                if run_len == 0 { run_start = page_idx; }
                run_len += 1;

                if run_len == count {
                    for i in run_start..run_start + count {
                        self.bitmap[i / 8] |= 1 << (i % 8);
                        self.ref_counts[i] = 1;
                    }
                    self.used_pages += count;
                    self.last_byte = run_start / 8;
                    return Some(PhysAddr(run_start as u64 * PAGE_SIZE));
                }
            } else {
                run_len = 0;
            }
        }

        None
    }
}

//Module-level convenience functions

pub fn alloc_frame() -> Option<PhysAddr> {
    PMM.lock().as_mut()?.alloc_frame()
}

pub fn alloc_contiguous(count: usize) -> Option<PhysAddr> {
    PMM.lock().as_mut()?.alloc_contiguous(count)
}

pub fn free_frame(addr: PhysAddr) {
    if let Some(pmm) = PMM.lock().as_mut() {
        pmm.free_frame(addr);
    }
}

pub fn inc_ref_frame(addr: PhysAddr) {
    if let Some(pmm) = PMM.lock().as_mut() {
        pmm.inc_ref_frame(addr);
    }
}

pub fn get_stats() -> (usize, usize) {
    if let Some(pmm) = PMM.lock().as_ref() {
        (pmm.used_pages(), pmm.total_pages())
    } else {
        (0, 0)
    }
}