feat: refactor memory management and integrate memblock

- Updated physical memory management (PMM) to utilize memblock for memory layout and frame allocation.
- Enhanced the virtual memory manager (VMM) to allocate the frame map from memblock, ensuring it starts at the correct virtual address without relocation.
- Improved frame allocation and deallocation functions to handle physical addresses more consistently.
- Adjusted the slab allocator to correctly convert between physical and virtual addresses.
- Refined page fault handling to ensure proper frame allocation and address translation.
- Streamlined the initialization process for PMM and VMM, ensuring a more robust setup of memory structures.

Co-authored-by: Copilot <copilot@github.com>

1 files changed, 150 insertions, 84 deletions

diff --git a/na/mm/vmm.cc b/na/mm/vmm.cc
index 50c7920..041413d 100644
--- a/na/mm/vmm.cc
+++ b/na/mm/vmm.cc
@@ -4,24 +4,30 @@
 /*
  * vmm.cc — Virtual Memory Manager (Phase 1: Kernel Page Table)
  *
- * Builds a fine-grained kernel page table (4KB pages, precise
- * permissions per section), activates it, and retires the early_pgdir
- * identity mapping from entry.S.
+ * Builds a fine-grained kernel page table, activates it, and retires
+ * the early_pgdir identity mapping from entry.S.
  *
- * This module is machine-independent (MI).  All architecture-specific
- * page table manipulation goes through the pmap:: namespace.
+ * With memblock integration, the initialization order is:
+ *   memblock_init → vmm_init (allocates frame_map from memblock,
+ *   maps vmemmap, activates page table) → pmm_init (buddy system
+ *   starts with frame_map already at vmemmap VA — no relocation).
  *
  * Reference: docs/specs/vmm.md
  */
 
-#include <lume/types.h>
+#include <arch/config.h>
+#include <arch/mmu.h>
+#include <arch/pmap.h>
 #include <lume/addr.h>
-#include <lume/pmm.h>
-#include <lume/klog.h>
+#include <lume/console.h>
 #include <lume/fdt.h>
+#include <lume/frame.h>
+#include <lume/klog.h>
+#include <lume/kprintf.h>
+#include <lume/memblock.h>
+#include <lume/pmm.h>
+#include <lume/types.h>
 #include <lume/vmm.h>
-#include <arch/pmap.h>
-#include <arch/mmu.h>
 
 /* Linker-exported section boundaries (virtual addresses) */
 extern "C" char _stext[], _etext[];
@@ -33,92 +39,152 @@ extern "C" char _kernel_end[];
 /* Global kernel page table root PA */
 static uint64 g_kernel_pgtbl = 0;
 
-/* Convert a linker symbol to PA. Handles both VA and PA cases. */
-static inline uint64 sym_to_pa(const void *sym)
-{
-    return ensure_pa(reinterpret_cast<uint64>(sym));
+/* VMM exports for PMM: frame_map location after vmemmap mapping. */
+static uint64 g_vmemmap_frame_map_pa = 0;
+static uint64 g_vmemmap_frame_map_size = 0;
+static uint64 g_vmemmap_num_frames = 0;
+
+uint64 vmm_get_frame_map_pa() { return g_vmemmap_frame_map_pa; }
+uint64 vmm_get_frame_map_size() { return g_vmemmap_frame_map_size; }
+uint64 vmm_get_num_frames() { return g_vmemmap_num_frames; }
+
+/* Convert a linker symbol to PA (kernel is executing in higher-half VA). */
+static inline uint64 sym_to_pa(const void *sym) {
+  return va_to_pa(virt_addr(reinterpret_cast<uint64>(sym))).raw;
 }
 
-/* Map a PA range into the kernel page table with given permissions.
- * All addresses must be page-aligned. Panics on failure. */
-static void map_range(uint64 root, uint64 pa_start, uint64 pa_end, uint64 perm)
-{
-    for (uint64 pa = pa_start; pa < pa_end; pa += kPageSize) {
-        uint64 va = pa_to_va(pa);
-        int ret = pmap::map(root, va, pa, perm | pmap::PTE_G);
-        if (ret != 0)
-            kernel_panic("vmm_init: pmap::map failed");
+/* Map a PA range into the kernel page table via the direct map. */
+static void map_range(uint64 root, uint64 pa_start, uint64 pa_end,
+                      uint64 perm) {
+  for (uint64 pa = pa_start; pa < pa_end; pa += kPageSize) {
+    uint64 va = pa_to_va(phys_addr(pa)).raw;
+    int ret = pmap::map(root, va, pa, perm | pmap::PTE_G);
+    if (ret != 0) {
+      kernel_panic("vmm_init: pmap::map failed");
     }
+  }
 }
 
-
-
-void vmm_init()
-{
-    /* 1. Allocate new root page table */
-    uint64 root = pmap::create();
-    if (!root)
-        kernel_panic("vmm_init: failed to create root page table");
-
-    /* 2. Compute PA boundaries of each kernel section */
-    uint64 text_start  = page_align_down(sym_to_pa(_stext));
-    uint64 text_end    = page_align_up(sym_to_pa(_etext));
-    uint64 ro_start    = page_align_down(sym_to_pa(_srodata));
-    uint64 ro_end      = page_align_up(sym_to_pa(_erodata));
-    uint64 data_start  = page_align_down(sym_to_pa(_sdata));
-    uint64 bss_end     = page_align_up(sym_to_pa(_ebss));
-
-    /* Memory extent from PMM */
-    uint64 mem_base = pmm_get_mem_base();
-    uint64 mem_end  = mem_base + pmm_get_mem_size();
-
-    /* 3. Map kernel .text (RX) */
-    map_range(root, text_start, text_end,
-              pmap::PTE_R | pmap::PTE_X);
-
-    /* 4. Map .rodata (R) */
-    map_range(root, ro_start, ro_end,
-              pmap::PTE_R);
-
-    /* 5. Map .data + .bss (RW) */
-    map_range(root, data_start, bss_end,
-              pmap::PTE_R | pmap::PTE_W);
-
-    /* 6. Map remaining physical memory after kernel (RW)
-     * This covers frame_map, slab objects, and all free pages. */
-    uint64 after_kernel = page_align_up(sym_to_pa(_kernel_end));
-    if (after_kernel < mem_end) {
-        map_range(root, after_kernel, mem_end,
-                  pmap::PTE_R | pmap::PTE_W);
+/* Map a PA range to an explicit VA base, using 2MB pages where possible. */
+static void map_range_at_huge(uint64 root, uint64 va_base, uint64 pa_start,
+                              uint64 pa_end, uint64 perm) {
+  constexpr uint64 k2MB = 2ULL * 1024 * 1024;
+  uint64 pa = pa_start;
+  uint64 va = va_base;
+
+  /* Map 2MB-aligned chunks with mega pages */
+  while (pa + k2MB <= pa_end) {
+    if ((va & (k2MB - 1)) == 0 && (pa & (k2MB - 1)) == 0) {
+      int ret = pmap::map_2mb(root, va, pa, perm | pmap::PTE_G);
+      if (ret != 0) {
+        kernel_panic("vmm_init: pmap::map_2mb failed");
+      }
+      pa += k2MB;
+      va += k2MB;
+    } else {
+      break;
     }
-    /* 7. MMIO regions will be mapped later via vmm_map_kernel_mmio by architecture/drivers */
-
-    /* 8. Activate the new page table */
-    pmap::activate(root);
+  }
 
-    g_kernel_pgtbl = root;
+  /* Map remaining 4KB tail */
+  while (pa < pa_end) {
+    int ret = pmap::map(root, va, pa, perm | pmap::PTE_G);
+    if (ret != 0) {
+      kernel_panic("vmm_init: vmemmap 4KB pmap::map failed");
+    }
+    pa += kPageSize;
+    va += kPageSize;
+  }
 }
 
-void vmm_init_ap()
-{
-    /* AP just switches to the kernel page table built by BSP */
-    if (!g_kernel_pgtbl)
-        kernel_panic("vmm_init_ap: kernel page table not ready");
-    pmap::activate(g_kernel_pgtbl);
+void vmm_init() {
+  /* 1. Allocate new root page table (from memblock) */
+  uint64 root = pmap::create();
+  if (!root) {
+    kernel_panic("vmm_init: failed to create root page table");
+  }
+
+  /* 2. Compute PA boundaries of each kernel section */
+  uint64 text_start = page_align_down(sym_to_pa(_stext));
+  uint64 text_end = page_align_up(sym_to_pa(_etext));
+  uint64 ro_start = page_align_down(sym_to_pa(_srodata));
+  uint64 ro_end = page_align_up(sym_to_pa(_erodata));
+  uint64 data_start = page_align_down(sym_to_pa(_sdata));
+  uint64 bss_end = page_align_up(sym_to_pa(_ebss));
+
+  /* Memory extent from memblock */
+  uint64 mem_base = memblock_get_mem_base();
+  uint64 mem_size = memblock_get_mem_size();
+  uint64 mem_end = mem_base + mem_size;
+
+  /* 3. Map kernel .text (RX) */
+  map_range(root, text_start, text_end, pmap::PTE_R | pmap::PTE_X);
+
+  /* 4. Map .rodata (R) */
+  map_range(root, ro_start, ro_end, pmap::PTE_R);
+
+  /* 5. Map .data + .bss (RW) */
+  map_range(root, data_start, bss_end, pmap::PTE_R | pmap::PTE_W);
+
+  /* 6. Map remaining physical memory after kernel (RW) */
+  uint64 after_kernel = page_align_up(sym_to_pa(_kernel_end));
+  if (after_kernel < mem_end) {
+    map_range(root, after_kernel, mem_end, pmap::PTE_R | pmap::PTE_W);
+  }
+
+  /* 7. Allocate frame_map from memblock (top of RAM) and map at VMEMMAP.
+   *
+   * Since memblock is active (PMM not yet initialized), pmap::map
+   * allocates page table pages from memblock too.  frame_map starts
+   * life at the vmemmap VA — no relocation ever needed.
+   */
+  uint64 num_frames = mem_size / kPageSize;
+  uint64 frame_map_bytes = num_frames * sizeof(Frame);
+  uint64 frame_map_span = page_align_up(frame_map_bytes);
+  uint64 frame_map_pa = memblock_alloc_top(frame_map_span, kPageSize);
+
+  if (!frame_map_pa) {
+    kernel_panic("vmm_init: memblock failed to allocate frame_map");
+  }
+
+  kprintf("[vmm] vmemmap: %llu frames, %llu KB at PA 0x%llx\n", num_frames,
+          frame_map_span / 1024, frame_map_pa);
+
+  /* Map frame_map at kVmemmapBase, using 2MB pages where possible */
+  map_range_at_huge(root, arch::kVmemmapBase, frame_map_pa,
+                    frame_map_pa + frame_map_span, pmap::PTE_R | pmap::PTE_W);
+
+  /* Export for pmm_init() */
+  g_vmemmap_frame_map_pa = frame_map_pa;
+  g_vmemmap_frame_map_size = frame_map_bytes;
+  g_vmemmap_num_frames = num_frames;
+
+  /* 8. Pre-map console MMIO before page table switch */
+  g_kernel_pgtbl = root;
+  console_init();
+
+  /* 9. Activate the new page table */
+  pmap::activate(root);
+
+  kprintf("[vmm] vmm_init complete (vmemmap at 0x%llx)\n", arch::kVmemmapBase);
 }
 
-void vmm_map_kernel_mmio(uint64 pa, uint64 size)
-{
-    if (!g_kernel_pgtbl)
-        kernel_panic("vmm_map_kernel_mmio: kernel page table not ready");
-    
-    uint64 aligned_pa = page_align_down(pa);
-    uint64 aligned_size = page_align_up(pa + size) - aligned_pa;
-    
-    map_range(g_kernel_pgtbl, aligned_pa, aligned_pa + aligned_size,
-              pmap::PTE_R | pmap::PTE_W | pmap::PTE_G);
-    arch::mmu::flush_tlb_all();
+void vmm_init_ap() {
+  if (!g_kernel_pgtbl) {
+    kernel_panic("vmm_init_ap: kernel page table not ready");
+  }
+  pmap::activate(g_kernel_pgtbl);
 }
 
+void vmm_map_kernel_mmio(uint64 pa, uint64 size) {
+  if (!g_kernel_pgtbl) {
+    kernel_panic("vmm_map_kernel_mmio: kernel page table not ready");
+  }
 
+  uint64 aligned_pa = page_align_down(pa);
+  uint64 aligned_size = page_align_up(pa + size) - aligned_pa;
 
+  map_range(g_kernel_pgtbl, aligned_pa, aligned_pa + aligned_size,
+            pmap::PTE_R | pmap::PTE_W | pmap::PTE_G);
+  arch::mmu::flush_tlb_all();
+}