lancer / kernel / src / mem / phys.rs

use lancer_core::bitmap::bitmap_seal::Sealed;
use lancer_core::bitmap::{BitmapAllocator as BitmapCore, BitmapBacking};
use limine::memory_map::{Entry, EntryType};
use x86_64::PhysAddr;
use x86_64::structures::paging::{FrameAllocator, PhysFrame, Size4KiB};

use super::RawSlice;
use super::addr;
use super::frame::OwnedFrame;
use crate::sync::IrqMutex;

const PAGE_SIZE: u64 = 4096;

static BITMAP_PHYS_BASE: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);
static BITMAP_BYTE_COUNT: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);

#[allow(dead_code)]
pub fn bitmap_region() -> (u64, u64) {
    (
        BITMAP_PHYS_BASE.load(core::sync::atomic::Ordering::Relaxed),
        BITMAP_BYTE_COUNT.load(core::sync::atomic::Ordering::Relaxed),
    )
}

impl Sealed for RawSlice<u64> {}
impl BitmapBacking for RawSlice<u64> {
    fn chunks(&self) -> &[u64] {
        self.as_slice()
    }

    fn chunks_mut(&mut self) -> &mut [u64] {
        self.as_slice_mut()
    }
}

static BITMAP: IrqMutex<BitmapCore<RawSlice<u64>>, 3> =
    IrqMutex::new(BitmapCore::from_backing(RawSlice::empty()));

pub struct BitmapFrameAllocator;

impl BitmapFrameAllocator {
    pub fn init(memory_map: &[&Entry]) {
        let max_addr = memory_map
            .iter()
            .filter(|entry| entry.entry_type == EntryType::USABLE)
            .map(|entry| entry.base + entry.length)
            .fold(0u64, u64::max);

        let total_frames = (max_addr / PAGE_SIZE) as usize;
        let chunks_needed = total_frames.div_ceil(64);
        let bitmap_bytes = chunks_needed * 8;
        let bitmap_frames = bitmap_bytes.div_ceil(PAGE_SIZE as usize);

        let bitmap_region = memory_map
            .iter()
            .filter(|entry| entry.entry_type == EntryType::USABLE)
            .find(|entry| entry.length >= (bitmap_frames as u64) * PAGE_SIZE)
            .expect("[phys] no usable region large enough for bitmap");

        let bitmap_phys_base = bitmap_region.base;
        BITMAP_PHYS_BASE.store(bitmap_phys_base, core::sync::atomic::Ordering::Relaxed);
        BITMAP_BYTE_COUNT.store(
            (bitmap_frames as u64) * PAGE_SIZE,
            core::sync::atomic::Ordering::Relaxed,
        );
        let bitmap_virt = addr::phys_to_virt(PhysAddr::new(bitmap_phys_base));
        assert!(
            bitmap_virt.is_aligned(8u64),
            "bitmap backing memory not 8-byte aligned"
        );
        let bitmap_ptr = bitmap_virt.as_u64() as *mut u64;

        let mut bm = BITMAP.lock();
        let mut backing = RawSlice::empty();
        backing.init(bitmap_ptr, chunks_needed);
        *bm = BitmapCore::from_backing(backing);
        bm.init(total_frames);

        memory_map
            .iter()
            .filter(|entry| entry.entry_type == EntryType::USABLE)
            .for_each(|entry| {
                let start_frame = (entry.base / PAGE_SIZE) as usize;
                let end_frame = ((entry.base + entry.length) / PAGE_SIZE) as usize;
                bm.mark_range_free(start_frame..end_frame);
            });

        let bitmap_start_frame = (bitmap_phys_base / PAGE_SIZE) as usize;
        (bitmap_start_frame..bitmap_start_frame + bitmap_frames).for_each(|idx| bm.mark_used(idx));

        bm.mark_used(0);
    }

    pub fn allocate_contiguous(&self, count: usize) -> Option<PhysAddr> {
        BITMAP
            .lock()
            .allocate_contiguous(count)
            .map(|idx| PhysAddr::new((idx as u64) * PAGE_SIZE))
    }

    pub fn allocate(&self) -> Option<OwnedFrame> {
        let mut bm = BITMAP.lock();
        bm.allocate().map(|idx| {
            let phys = PhysAddr::new((idx as u64) * PAGE_SIZE);
            OwnedFrame::new(PhysFrame::containing_address(phys))
        })
    }

    pub(crate) fn free_frame_by_addr(phys: PhysAddr) {
        let mut bm = BITMAP.lock();
        let frame_idx = (phys.as_u64() / PAGE_SIZE) as usize;
        bm.deallocate(frame_idx).unwrap_or_else(|_| {
            panic!(
                "[phys] invariant violation: double-free of frame {:#x}",
                phys.as_u64()
            )
        });
    }

    #[allow(dead_code)]
    pub fn deallocate(&self, frame: OwnedFrame) {
        Self::free_frame_by_addr(frame.inner().start_address());
    }

    pub fn deallocate_frame(&self, frame: PhysFrame<Size4KiB>) {
        Self::free_frame_by_addr(frame.start_address());
    }

    #[allow(dead_code)]
    pub fn mark_used(idx: usize) {
        BITMAP.lock().mark_used(idx);
    }

    #[allow(dead_code)]
    pub fn is_used(idx: usize) -> bool {
        BITMAP.lock().is_used(idx)
    }

    pub fn total_frames() -> usize {
        BITMAP.lock().total_items()
    }

    #[allow(dead_code)]
    pub fn used_frames() -> usize {
        BITMAP.lock().used_items()
    }

    pub fn free_frames() -> usize {
        BITMAP.lock().free_items()
    }

    pub fn for_each_used_range<F: FnMut(usize, usize)>(f: &mut F) {
        BITMAP.lock().for_each_used_range(f);
    }
}

unsafe impl FrameAllocator<Size4KiB> for BitmapFrameAllocator {
    fn allocate_frame(&mut self) -> Option<PhysFrame<Size4KiB>> {
        self.allocate().map(|tf| tf.inner())
    }
}

#[allow(dead_code)]
pub fn log_memory_map(memory_map: &[&Entry]) {
    crate::kprintln!("  Memory map ({} entries):", memory_map.len());
    memory_map.iter().for_each(|entry| {
        let entry_type_str = match entry.entry_type {
            EntryType::USABLE => "Usable",
            EntryType::RESERVED => "Reserved",
            EntryType::ACPI_RECLAIMABLE => "ACPI Reclaimable",
            EntryType::ACPI_NVS => "ACPI NVS",
            EntryType::BAD_MEMORY => "Bad Memory",
            EntryType::BOOTLOADER_RECLAIMABLE => "Bootloader Reclaimable",
            EntryType::EXECUTABLE_AND_MODULES => "Kernel/Modules",
            EntryType::FRAMEBUFFER => "Framebuffer",
            _ => "Unknown",
        };
        let size_kb = entry.length / 1024;
        crate::kprintln!(
            "    {:#016x} - {:#016x} ({:>8} KB) {}",
            entry.base,
            entry.base + entry.length,
            size_kb,
            entry_type_str
        );
    });
}

pub fn frame_stats() -> (usize, usize) {
    (
        BitmapFrameAllocator::total_frames(),
        BitmapFrameAllocator::free_frames(),
    )
}

#[allow(dead_code)]
pub fn log_frame_stats() {
    crate::kprintln!("  Frame allocator stats:");
    crate::kprintln!("    Total frames: {}", BitmapFrameAllocator::total_frames());
    crate::kprintln!("    Used frames:  {}", BitmapFrameAllocator::used_frames());
    crate::kprintln!("    Free frames:  {}", BitmapFrameAllocator::free_frames());
    crate::kprintln!(
        "    Free memory:  {} MB",
        BitmapFrameAllocator::free_frames() * 4096 / (1024 * 1024)
    );
}