changed page allocation to happen on page fault for performance reasons

libk/src/drivers/kalloc/allocator.rs

@@ -2,10 +2,16 @@ use linked_list_allocator::LockedHeap;
 use x86_64::{
     VirtAddr,
     structures::paging::{
-        FrameAllocator, Mapper, Page, PageTableFlags, Size4KiB, mapper::MapToError,
+        FrameAllocator, Mapper, Page, PageTableFlags, Size4KiB,
+        mapper::{MapToError, MapperFlushAll},
     },
 };
 
+use crate::{
+    drivers::memory::{FRAME_ALLOCATOR, OFFSET_PAGE_TABLE},
+    serial_print, serial_println,
+};
+
 /// We are currently using a linked list heap allocator which uses our underlying page allocator.
 pub type FoundryAllocator = LockedHeap;
 
@@ -14,28 +20,50 @@ pub type FoundryAllocator = LockedHeap;
 static ALLOCATOR: FoundryAllocator = FoundryAllocator::empty();
 
 pub const HEAP_START: usize = 0x4444_4444_0000;
-pub const HEAP_SIZE: usize = 1000 * 1024;
+pub const HEAP_SIZE: usize = 1024 * 1024 * 1024;
 
 /// Sets up the heap using the backing page frame allocator.
-pub fn init_heap(
-    mapper: &mut impl Mapper<Size4KiB>,
-    frame_allocator: &mut impl FrameAllocator<Size4KiB>,
-) -> Result<(), MapToError<Size4KiB>> {
-    let range = {
-        let heap_start = VirtAddr::new(HEAP_START as u64);
-        let heap_end = heap_start + HEAP_SIZE as u64 - 1u64;
-        let heap_start_page = Page::<Size4KiB>::containing_address(heap_start);
-        let heap_end_page = Page::<Size4KiB>::containing_address(heap_end);
-        Page::range_inclusive(heap_start_page, heap_end_page)
-    };
+pub fn init_heap() -> Result<(), MapToError<Size4KiB>> {
+    // let mut frame_allocator = if let Some(f) = FRAME_ALLOCATOR.get() {
+    //     f.lock()
+    // } else {
+    //     return Err(MapToError::FrameAllocationFailed);
+    // };
 
-    for page in range {
-        let frame = frame_allocator
-            .allocate_frame()
-            .ok_or(MapToError::FrameAllocationFailed)?;
-        let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
-        unsafe { mapper.map_to(page, frame, flags, frame_allocator)?.flush() };
-    }
+    // let mut mapper = if let Some(m) = OFFSET_PAGE_TABLE.get() {
+    //     m.lock()
+    // } else {
+    //     return Err(MapToError::FrameAllocationFailed);
+    // };
+
+    // let range = {
+    //     let heap_start = VirtAddr::new(HEAP_START as u64);
+    //     let heap_end = heap_start + HEAP_SIZE as u64 - 1u64;
+    //     let heap_start_page = Page::<Size4KiB>::containing_address(heap_start);
+    //     let heap_end_page = Page::<Size4KiB>::containing_address(heap_end);
+    //     Page::range_inclusive(heap_start_page, heap_end_page)
+    // };
+
+    // let usable_frames = frame_allocator.count_usable_frames();
+    // serial_println!("usable frames: {}", usable_frames);
+
+    // let mut i = 0;
+    // for page in range {
+    //     i += 1;
+    //     if i % 128 == 0 {
+    //         serial_println!("allocated {} pages", i);
+    //     }
+    //     let frame = frame_allocator
+    //         .allocate_frame()
+    //         .ok_or(MapToError::FrameAllocationFailed)?;
+    //     let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
+    //     unsafe {
+    //         // IMPORTANT: make sure to flush the mapper!!!!
+    //         let _ = mapper.map_to(page, frame, flags, &mut *frame_allocator)?;
+    //     }
+    // }
+
+    // MapperFlushAll::new().flush_all();
 
     unsafe {
         ALLOCATOR.lock().init(HEAP_START as *mut u8, HEAP_SIZE);

libk/src/drivers/memory.rs

@@ -0,0 +1,149 @@
+// use lib_alloc::allocator::FoundryAllocator;
+use limine::{memory_map::EntryType, response::MemoryMapResponse};
+use spin::{Mutex, Once};
+use x86_64::{
+    PhysAddr, VirtAddr,
+    registers::control::Cr3,
+    structures::paging::{
+        FrameAllocator, Mapper, OffsetPageTable, Page, PageTable, PhysFrame, Size4KiB,
+        page_table::FrameError,
+    },
+};
+
+pub static FRAME_ALLOCATOR: Once<Mutex<FoundryOSFrameAllocator>> = Once::new();
+pub static OFFSET_PAGE_TABLE: Once<Mutex<OffsetPageTable>> = Once::new();
+/// Returns a mutable reference to the current level 4 page table.
+///
+/// # Safety
+///
+/// The caller must ensure that the level 4 page table is not modified
+/// simultaneously. The caller must also ensure that the physical memory offset
+/// is correct, to ensure that the correct virtual address is constructed.
+unsafe fn active_l4_table(physical_memory_offset: VirtAddr) -> &'static mut PageTable {
+    let (level_4_frame, _) = Cr3::read();
+
+    let phys_addr = level_4_frame.start_address();
+    let virt = phys_addr.as_u64() + physical_memory_offset.as_u64();
+    unsafe { &mut *(virt as *mut PageTable) }
+}
+
+/// Initializes the `OffsetPageTable` for the current CPU architecture.
+///
+/// # Safety
+///
+/// This function must be called only once and should be called before any
+/// memory operations are performed that rely on virtual memory management.
+/// The provided `physical_memory_offset` must be accurate to ensure correct
+/// translation of physical addresses.
+///
+/// # Parameters
+///
+/// - `physical_memory_offset`: The offset to convert physical addresses to
+///   virtual addresses in the higher-half direct map.
+///
+/// # Returns
+///
+/// Returns an `OffsetPageTable` that allows for manipulation of the page
+/// tables for the current CPU architecture.
+pub fn init_page_table(physical_memory_offset: VirtAddr) {
+    unsafe {
+        let l4_table = active_l4_table(physical_memory_offset);
+        let offset_table = OffsetPageTable::new(l4_table, physical_memory_offset);
+        OFFSET_PAGE_TABLE.call_once(|| Mutex::new(offset_table));
+    }
+}
+
+pub struct FoundryOSFrameAllocator {
+    memory_map: &'static MemoryMapResponse,
+    next: usize,
+}
+
+pub fn init_frame_allocator(memory_map: &'static MemoryMapResponse) {
+    unsafe {
+        FRAME_ALLOCATOR.call_once(|| Mutex::new(FoundryOSFrameAllocator::init(memory_map)));
+    }
+}
+
+impl FoundryOSFrameAllocator {
+    /// Creates a new `FoundryOSFrameAllocator` from a memory map.
+    ///
+    /// This function takes a reference to a `MemoryMapResponse` and initializes a
+    /// `FoundryOSFrameAllocator` with it. The `next` field is set to 0, indicating that
+    /// the first frame to be allocated is the first frame in the memory map.
+    pub const unsafe fn init(memory_map: &'static MemoryMapResponse) -> Self {
+        Self {
+            memory_map,
+            next: 0,
+        }
+    }
+
+    pub fn count_usable_frames(&self) -> u32 {
+        self.usable_frames().count() as u32
+    }
+
+    /// An iterator over all usable frames in the memory map.
+    ///
+    /// Yields one `PhysFrame` for each available 4KiB frame in the memory map.
+    ///
+    /// This function is used to allocate frames for the pagemap.
+    fn usable_frames(&self) -> impl Iterator<Item = PhysFrame> + use<> {
+        let regions = self.memory_map.entries().iter();
+        let usable_regions = regions.filter(|region| region.entry_type == EntryType::USABLE);
+
+        let addr_ranges = usable_regions.map(|region| region.base..region.base + region.length);
+
+        let frame_addresses = addr_ranges.flat_map(|r| r.step_by(4096));
+
+        frame_addresses.map(|addr| PhysFrame::from_start_address(PhysAddr::new(addr)).unwrap())
+    }
+}
+
+unsafe impl FrameAllocator<Size4KiB> for FoundryOSFrameAllocator {
+    /// Allocates a frame from the list of usable frames.
+    ///
+    /// This function returns the next available `PhysFrame` from the memory map,
+    /// if one exists. Once a frame is allocated, the internal counter is incremented
+    /// to point to the next frame for future allocations.
+    ///
+    /// # Returns
+    ///
+    /// - `Some(PhysFrame)`: If a usable frame is available.
+    /// - `None`: If there are no more usable frames to allocate.
+    fn allocate_frame(&mut self) -> Option<PhysFrame> {
+        let frame = self.usable_frames().nth(self.next);
+        self.next += 1;
+        frame
+    }
+}
+
+// pub unsafe fn translate_addr(addr: VirtAddr, physical_memory_offset: VirtAddr) -> Option<PhysAddr> {
+//     translate_addr_inner(addr, physical_memory_offset)
+// }
+
+// fn translate_addr_inner(addr: VirtAddr, physical_memory_offset: VirtAddr) -> Option<PhysAddr> {
+//     let (l4_table_frame, _) = Cr3::read();
+
+//     let table_indexes = [
+//         addr.p4_index(),
+//         addr.p3_index(),
+//         addr.p2_index(),
+//         addr.p1_index(),
+//     ];
+//     let mut frame = l4_table_frame;
+
+//     for &i in &table_indexes {
+//         let virt = physical_memory_offset + frame.start_address().as_u64();
+//         let table_ptr: *const PageTable = virt.as_ptr();
+//         let table = unsafe { &*table_ptr };
+
+//         let entry = &table[i];
+
+//         frame = match entry.frame() {
+//             Ok(frame) => frame,
+//             Err(FrameError::FrameNotPresent) => return None,
+//             Err(FrameError::HugeFrame) => panic!("huge frames are not supported!"),
+//         };
+//     }
+
+//     Some(frame.start_address() + u64::from(addr.page_offset()))
+// }

libk/src/drivers/mod.rs

@@ -2,4 +2,5 @@ pub mod io;
 pub mod kalloc;
 
 pub mod async_io;
+pub mod memory;
 pub mod pic;