kernel/src/physical_memory.rs

use crate::{
    bootinfo::BOOTINFO,
    println,
    virtual_memory::{AsVirt, PHYS_OFFSET},
};
use bootloader_api::info::MemoryRegionKind;
use cast::{u64, usize};
use core::{alloc::Layout, ptr::NonNull};
use derive_try_from_primitive::TryFromPrimitive;
use humansize::{SizeFormatter, BINARY};
use linked_list_allocator::hole::HoleList;
use spin::{Lazy, Mutex};
use x86_64::{
    structures::paging::{
        frame::PhysFrameRange, 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,
    Invalid = u32::MAX,
}

pub struct PhysicalMemory {
    alloc: HoleList,
    bytes_used: usize,
    peak_used: usize,
}

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

impl PhysicalMemory {
    pub fn print_stats(&self) {
        println!(
            "[PMM] {} currently allocated ({})",
            SizeFormatter::new(self.bytes_used, BINARY),
            self.bytes_used / 4096,
        );
        println!(
            "[PMM] {} allocated at peak ({})",
            SizeFormatter::new(self.peak_used, BINARY),
            self.peak_used / 4096,
        );
    }
}

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.bytes_used = self.bytes_used.saturating_add(4096);
        if self.bytes_used > self.peak_used {
            self.peak_used = self.bytes_used;
        }
        self.alloc.allocate_first_fit(FRAME_LAYOUT).ok().map(|(ptr, _)| {
            #[expect(clippy::unwrap_used, reason = "PhysFrame requires its argument to be 4k aligned, which is guaranteed by the allocator")]
            #[expect(clippy::arithmetic_side_effects, reason = "All addresses passed to the allocator are > PHYS_OFFSET, so this cannot underflow")]
            PhysFrame::from_start_address(PhysAddr::new(
                (u64(ptr.as_ptr().expose_provenance())) - PHYS_OFFSET.as_u64(),
            ))
            .unwrap()
        })
    }
}

impl FrameDeallocator<Size4KiB> for PhysicalMemory {
    unsafe fn deallocate_frame(&mut self, frame: PhysFrame) {
        self.bytes_used = self.bytes_used.saturating_sub(4096);
        #[expect(
            clippy::unwrap_used,
            reason = "The virtual mapping for the frame will never be null as the physical mapping offset is in kernel space"
        )]
        unsafe {
            self.alloc.deallocate(NonNull::new(frame.as_virt_ptr()).unwrap(), FRAME_LAYOUT)
        };
    }
}

impl PhysicalMemory {
    pub fn allocate_frame_range(&mut self, len: usize) -> Option<PhysFrameRange> {
        self.bytes_used = self.bytes_used.saturating_add(4096 * len);
        if self.bytes_used > self.peak_used {
            self.peak_used = self.bytes_used;
        }
        if len >= 0x10_0000_0000_0000 {
            return None;
        }
        #[expect(clippy::arithmetic_side_effects, reason = "The check above makes sure the multiplication cannot overflow.")]
        #[expect(clippy::unwrap_used, reason = "4096 is a nonzero power of two, and len is already a multiple of 4096 so rounding cannot overflow.")]
        self.alloc.allocate_first_fit(Layout::from_size_align(4096 * len, 4096).unwrap()).ok().map(|(ptr, _)| {
            #[expect(clippy::unwrap_used, reason = "PhysFrame requires its argument to be 4k aligned, which is guaranteed by the allocator")]
            #[expect(clippy::arithmetic_side_effects, reason = "All addresses passed to the allocator are > PHYS_OFFSET, so this cannot underflow")]
            let start = PhysFrame::from_start_address(PhysAddr::new(
                (u64(ptr.as_ptr().expose_provenance())) - PHYS_OFFSET.as_u64(),
            ))
            .unwrap();
            #[expect(clippy::unwrap_used, reason = "PhysFrame requires its argument to be 4k aligned, which is guaranteed by the allocator")]
            #[expect(clippy::arithmetic_side_effects, reason = "All addresses passed to the allocator are > PHYS_OFFSET, so this cannot underflow")]
            let end = PhysFrame::from_start_address(PhysAddr::new(
                (u64(ptr.as_ptr().expose_provenance()) + (u64(len) * 4096)) - PHYS_OFFSET.as_u64(),
            ))
            .unwrap();
            PhysFrameRange {
                start,
                end
            }
        })
    }
}

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();
    while let Some(&(mut region)) = region_iter.next() {
        while 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;
            }
        }
        let fmtd_size = SizeFormatter::new(region.end - region.start, BINARY);
        if let MemoryRegionKind::UnknownUefi(efi_type) = region.kind {
            println!(
                "[PMM] Efi{:?}: {:#x} - {:#x} ({})",
                EfiMemoryTypes::try_from(efi_type).unwrap_or(EfiMemoryTypes::Invalid),
                region.start,
                region.end,
                fmtd_size,
            );
        } else {
            println!(
                "[PMM] {:?}: {:#x} - {:#x} ({})",
                region.kind, region.start, region.end, fmtd_size
            );
        }
        total_mem += region.end - region.start;
        if region.kind == MemoryRegionKind::Usable {
            region.end &= !(0xFFF);
            if region.start & 0xFFF != 0 {
                region.start = (region.start & !(0xFFF)) + 0x1000;
            }
            usable_mem += region.end - region.start;
            unsafe {
                alloc.deallocate(
                    #[expect(clippy::unwrap_used, reason = "The virtual mapping for the frame will never be null as the physical mapping offset is in kernel space")]
                    NonNull::new(PhysAddr::new(region.start).as_virt_ptr()).unwrap(),
                    #[expect(clippy::unwrap_used, reason = "
                        from_size_align requires align to be a nonzero power of two, which it is.
                        Also, size must be less than isize when rounded up to a multiple of align. 
                        Since size is already a multiple of align due to start and end being page aligned, no overflow will occur.
                    ")]
                    Layout::from_size_align(usize(region.end - region.start), 4096).unwrap(),
                );
            }
        }
    }
    println!(
        "[PMM] Initialized, found {} of usable memory, {} total",
        SizeFormatter::new(usable_mem, BINARY),
        SizeFormatter::new(total_mem, BINARY),
    );
    Mutex::new(PhysicalMemory { alloc, bytes_used: 0, peak_used: 0 })
});