kernel/src/physical_memory.rs

use crate::{
    bootinfo::BOOTINFO,
    println,
    virtual_memory::{AsVirt, PHYS_OFFSET},
};
use bootloader_api::info::MemoryRegionKind;
use core::{alloc::Layout, ptr::NonNull};
use derive_try_from_primitive::TryFromPrimitive;
use linked_list_allocator::hole::HoleList;
use spin::{Lazy, Mutex};
use x86_64::{
    structures::paging::{FrameAllocator, FrameDeallocator, PhysFrame, Size4KiB},
    PhysAddr,
};

#[repr(u32)]
#[derive(Copy, Clone, Debug, TryFromPrimitive)]
enum EfiMemoryTypes {
    Reserved,
    LoaderCode,
    LoaderData,
    BootServicesCode,
    BootServicesData,
    RuntimeServicesCode,
    RuntimeServicesData,
    Conventional,
    Unusable,
    ACPIReclaim,
    ACPIMemoryNVS,
    MMIO,
    MMIOPortSpace,
    PalCode,
    Persistent,
}

pub struct PhysicalMemory {
    alloc: HoleList,
    frames_allocated: usize,
    frames_freed: usize,
}

unsafe impl Send for PhysicalMemory {}
unsafe impl Sync for PhysicalMemory {}

fn format_byte_count(count: usize) -> (f64, &'static str) {
    let mut count = count as f64;
    let mut prefix = 0;
    while count >= 1024.0 {
        count /= 1024.0;
        prefix += 1;
    }
    let prefix = ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei"][prefix];
    (count, prefix)
}

impl PhysicalMemory {
    pub fn print_stats(&self) {
        let (fmtd_alloced, alloced_pfx) = format_byte_count(self.frames_allocated * 4096);
        let (fmtd_freed, freed_pfx) = format_byte_count(self.frames_freed * 4096);
        let (fmtd_total, total_pfx) =
            format_byte_count((self.frames_allocated - self.frames_freed) * 4096);
        println!(
            "[PMM] {:2} {}B allocated, {:2} {}B freed ({:2} {}B total)",
            fmtd_alloced, alloced_pfx, fmtd_freed, freed_pfx, fmtd_total, total_pfx
        );
    }
}

const FRAME_LAYOUT: Layout = unsafe { Layout::from_size_align_unchecked(4096, 4096) };

unsafe impl FrameAllocator<Size4KiB> for PhysicalMemory {
    fn allocate_frame(&mut self) -> Option<PhysFrame> {
        self.frames_allocated += 1;
        self.alloc.allocate_first_fit(FRAME_LAYOUT).ok().map(|(ptr, _)| {
            PhysFrame::from_start_address(PhysAddr::new(
                (ptr.as_ptr() as u64) - PHYS_OFFSET.as_u64(),
            ))
            .unwrap()
        })
    }
}

impl FrameDeallocator<Size4KiB> for PhysicalMemory {
    unsafe fn deallocate_frame(&mut self, frame: PhysFrame) {
        self.frames_freed += 1;
        unsafe { self.alloc.deallocate(NonNull::new(frame.as_virt_ptr()).unwrap(), FRAME_LAYOUT) };
    }
}

pub static PHYSICAL_MEMORY: Lazy<Mutex<PhysicalMemory>> = Lazy::new(|| {
    println!("[PMM] Bootloader reports the following regions:");
    let mut region_iter = BOOTINFO.memory_regions.iter().peekable();
    let mut total_mem = 0;
    let mut usable_mem = 0;
    let mut alloc = HoleList::empty();
    loop {
        let mut region = if let Some(region) = region_iter.next() {
            region.clone()
        } else {
            break;
        };
        loop {
            if let Some(next_region) = region_iter.peek() {
                if (next_region.kind == region.kind) && (next_region.start == region.end) {
                    region.end = next_region.end;
                    region_iter.next();
                } else {
                    break;
                }
            } else {
                break;
            }
        }
        let (fmtd_size, pfx) = format_byte_count((region.end - region.start) as usize);
        if let MemoryRegionKind::UnknownUefi(efi_type) = region.kind {
            println!(
                "[PMM] Efi{:?}: {:#x} - {:#x} ({:2} {}B)",
                EfiMemoryTypes::try_from(efi_type).unwrap(),
                region.start,
                region.end,
                fmtd_size,
                pfx
            );
        } else {
            println!(
                "[PMM] {:?}: {:#x} - {:#x} ({:2} {}B)",
                region.kind, region.start, region.end, fmtd_size, pfx
            );
        }
        total_mem += region.end - region.start;
        if region.kind == MemoryRegionKind::Usable {
            region.end = region.end & !(0xFFF);
            if region.start & 0xFFF != 0 {
                region.start = (region.start & !(0xFFF)) + 0x1000;
            }
            usable_mem += region.end - region.start;
            unsafe {
                alloc.deallocate(
                    NonNull::new(PhysAddr::new(region.start).as_virt_ptr()).unwrap(),
                    Layout::from_size_align((region.end - region.start) as usize, 4096).unwrap(),
                );
            }
        }
    }
    let (fmtd_usable, usable_pfx) = format_byte_count(usable_mem as usize);
    let (fmtd_total, total_pfx) = format_byte_count(total_mem as usize);
    println!(
        "[PMM] Initialized, found {:.2} {}B of usable memory, {:2} {}B total",
        fmtd_usable, usable_pfx, fmtd_total, total_pfx
    );
    Mutex::new(PhysicalMemory { alloc, frames_allocated: 0, frames_freed: 0 })
});