opencoreC/kernel/src/mm/pfa.rs

//! Physical Frame Allocator (PFA)
//! Bitmap-based 4KiB page frame allocator.

use crate::limine_requests::{LimineMemmapResponse, LimineMemoryType};
use crate::kprintln;

pub const PAGE_SIZE: usize = 4096;

pub struct FrameAllocator {
    pub bitmap: *mut u8,
    pub bitmap_size_bytes: usize,
    pub total_frames: usize,
    pub free_frames: usize,
    pub usable_ram_bytes: u64,
    pub hhdm_offset: u64,
}

unsafe impl Send for FrameAllocator {}
unsafe impl Sync for FrameAllocator {}

pub static mut PFA: FrameAllocator = FrameAllocator {
    bitmap: core::ptr::null_mut(),
    bitmap_size_bytes: 0,
    total_frames: 0,
    free_frames: 0,
    usable_ram_bytes: 0,
    hhdm_offset: 0,
};

impl FrameAllocator {
    pub unsafe fn init(&mut self, memmap: *const LimineMemmapResponse, hhdm: u64) {
        self.hhdm_offset = hhdm;
        if memmap.is_null() {
            kprintln!("[PFA] ERROR: memmap is null!");
            return;
        }

        let entry_count = (*memmap).entry_count as usize;
        let entries_ptr = (*memmap).entries;
        let mut max_paddr: u64 = 0;
        for i in 0..entry_count {
            let entry = *entries_ptr.add(i);
            let end = (*entry).base + (*entry).length;
            if end > max_paddr {
                max_paddr = end;
            }
        }

        self.total_frames = (max_paddr as usize) / PAGE_SIZE;
        self.bitmap_size_bytes = (self.total_frames + 7) / 8;

        // Find a usable region large enough to place the bitmap
        let mut bitmap_paddr: u64 = 0;
        for i in 0..entry_count {
            let entry = *entries_ptr.add(i);
            if (*entry).typ == LimineMemoryType::Usable as u64
                && (*entry).length >= self.bitmap_size_bytes as u64
            {
                bitmap_paddr = (*entry).base;
                break;
            }
        }

        if bitmap_paddr == 0 {
            kprintln!("[FATAL] Could not find memory for PFA bitmap!");
            return;
        }
        self.bitmap = (bitmap_paddr + hhdm) as *mut u8;
        core::ptr::write_bytes(self.bitmap, 0xFF, self.bitmap_size_bytes);

        self.free_frames = 0;
        let mut total_usable: u64 = 0;
        for i in 0..entry_count {
            let entry = *entries_ptr.add(i);
            if (*entry).typ == LimineMemoryType::Usable as u64 {
                total_usable += (*entry).length;
                let start_frame = ((*entry).base as usize) / PAGE_SIZE;
                let count = ((*entry).length as usize) / PAGE_SIZE;
                for f in start_frame..(start_frame + count) {
                    self.clear_bit(f);
                    self.free_frames += 1;
                }
            }
        }
        self.usable_ram_bytes = total_usable;

        let bitmap_start_frame = (bitmap_paddr as usize) / PAGE_SIZE;
        let bitmap_frame_count = (self.bitmap_size_bytes + PAGE_SIZE - 1) / PAGE_SIZE;
        for f in bitmap_start_frame..(bitmap_start_frame + bitmap_frame_count) {
            self.set_bit(f);
            if self.free_frames > 0 {
                self.free_frames -= 1;
            }
        }

        let low_1mb_frames = 0x100000 / PAGE_SIZE;
        for f in 0..low_1mb_frames {
            if !self.test_bit(f) {
                self.set_bit(f);
                if self.free_frames > 0 {
                    self.free_frames -= 1;
                }
            }
        }

        kprintln!("[PFA] Initialized: {} MiB total, {} MiB free",
            (self.total_frames * PAGE_SIZE) / (1024 * 1024),
            (self.free_frames * PAGE_SIZE) / (1024 * 1024)
        );
    }

    #[inline]
    fn test_bit(&self, frame: usize) -> bool {
        if frame >= self.total_frames {
            return true;
        }
        unsafe {
            let byte = *self.bitmap.add(frame / 8);
            (byte & (1 << (frame % 8))) != 0
        }
    }

    #[inline]
    fn set_bit(&self, frame: usize) {
        if frame < self.total_frames {
            unsafe {
                let ptr = self.bitmap.add(frame / 8);
                *ptr |= 1 << (frame % 8);
            }
        }
    }

    #[inline]
    fn clear_bit(&self, frame: usize) {
        if frame < self.total_frames {
            unsafe {
                let ptr = self.bitmap.add(frame / 8);
                *ptr &= !(1 << (frame % 8));
            }
        }
    }

    pub fn alloc_frame(&mut self) -> Option<u64> {
        if self.free_frames == 0 {
            return None;
        }
        for f in 0..self.total_frames {
            if !self.test_bit(f) {
                self.set_bit(f);
                self.free_frames -= 1;
                let paddr = (f * PAGE_SIZE) as u64;
                unsafe {
                    let vaddr = (paddr + self.hhdm_offset) as *mut u8;
                    core::ptr::write_bytes(vaddr, 0, PAGE_SIZE);
                }
                return Some(paddr);
            }
        }
        None
    }

    pub fn alloc_contiguous_frames(&mut self, count: usize) -> Option<u64> {
        if self.free_frames < count || count == 0 {
            return None;
        }
        let mut consecutive = 0;
        let mut start_frame = 0;

        for f in 0..self.total_frames {
            if !self.test_bit(f) {
                if consecutive == 0 {
                    start_frame = f;
                }
                consecutive += 1;
                if consecutive == count {
                    for i in 0..count {
                        self.set_bit(start_frame + i);
                        let paddr = ((start_frame + i) * PAGE_SIZE) as u64;
                        unsafe {
                            let vaddr = (paddr + self.hhdm_offset) as *mut u8;
                            core::ptr::write_bytes(vaddr, 0, PAGE_SIZE);
                        }
                    }
                    self.free_frames -= count;
                    return Some((start_frame * PAGE_SIZE) as u64);
                }
            } else {
                consecutive = 0;
            }
        }
        None
    }

    pub fn free_frame(&mut self, paddr: u64) {
        let frame = (paddr as usize) / PAGE_SIZE;
        if frame < self.total_frames && self.test_bit(frame) {
            self.clear_bit(frame);
            self.free_frames += 1;
        }
    }
}