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, 308 insertions, 334 deletions

diff --git a/na/mm/pmm.cc b/na/mm/pmm.cc
index 00f998c..1400448 100644
--- a/na/mm/pmm.cc
+++ b/na/mm/pmm.cc
@@ -19,43 +19,42 @@
  * Reference: docs/specs/pmm.md
  */
 
-#include <lume/pmm.h>
+#include <arch/config.h>
+#include <arch/cpu.h>
+#include <kernel/sync/spinlock.h>
 #include <lume/addr.h>
 #include <lume/config.h>
-#include <lume/fdt.h>
-#include <lume/new.h>
 #include <lume/klog.h>
+#include <lume/kprintf.h>
+#include <lume/memblock.h>
+#include <lume/new.h>
+#include <lume/pmm.h>
 #include <lume/types.h>
-#include <arch/cpu.h>
-#include <kernel/sync/spinlock.h>
-
-/* ------------------------------------------------------------------ */
-/*  Linker symbols                                                    */
-/* ------------------------------------------------------------------ */
-
-extern "C" char _kernel_end[];
+#include <lume/vmm.h>
 
 /* ------------------------------------------------------------------ */
 /*  Per-CPU page cache                                                */
 /* ------------------------------------------------------------------ */
 
 struct PerCpuCache {
-    list_node free_list;   // Single-frame free list
-    uint32 count;          // Number of cached frames
+  list_node free_list; // Single-frame free list
+  uint32 count;        // Number of cached frames
 };
 
 /* ------------------------------------------------------------------ */
 /*  Module-private state                                              */
 /* ------------------------------------------------------------------ */
 
-static Spinlock        g_pmm_lock;
-static list_node       g_free_areas[kMaxOrder];   // Buddy free lists [0..10]
-static PerCpuCache     g_pcp[kMaxCpus];
+static Spinlock g_pmm_lock;
+static list_node g_free_areas[kMaxOrder]; // Buddy free lists [0..10]
+static PerCpuCache g_pcp[kMaxCpus];
 
-static Frame*          g_frame_map = nullptr;      // Flat array of Frame descriptors
-static uint64          g_mem_base  = 0;            // Physical memory base (from FDT)
-static uint64          g_mem_size  = 0;            // Total physical memory size
-static uint64          g_num_frames = 0;           // Total managed frames
+static Frame *g_frame_map = nullptr; // Flat array of Frame descriptors
+static uint64 g_frame_map_pa = 0;    // Physical base of frame_map storage
+static uint64 g_frame_map_bytes = 0; // Byte size of frame_map storage
+static uint64 g_mem_base = 0;        // Physical memory base (from FDT)
+static uint64 g_mem_size = 0;        // Total physical memory size
+static uint64 g_num_frames = 0;      // Total managed frames
 
 uint32 g_pcp_high_watermark = 0;
 uint32 g_pcp_batch_size = 0;
@@ -68,33 +67,28 @@ uint32 g_pcp_batch_size = 0;
  * Physical Frame Number: index into the frame_map array.
  * PFN = (pa - g_mem_base) / kPageSize
  */
-static inline uint64 pa_to_frame_index(uint64 pa)
-{
-    return (pa - g_mem_base) / kPageSize;
+static inline uint64 pa_to_frame_index(uint64 pa) {
+  return (pa - g_mem_base) / kPageSize;
 }
 
-static inline uint64 frame_index_to_pa(uint64 idx)
-{
-    return g_mem_base + idx * kPageSize;
+static inline uint64 frame_index_to_pa(uint64 idx) {
+  return g_mem_base + idx * kPageSize;
 }
 
-uint64 frame_to_pa(const Frame* frame)
-{
-    uint64 idx = static_cast<uint64>(frame - g_frame_map);
-    if (frame < g_frame_map || idx >= g_num_frames)
-        kernel_panic("frame_to_pa: frame outside frame_map bounds");
-    return frame_index_to_pa(idx);
+PhysAddr frame_to_pa(const Frame *frame) {
+  uint64 idx = static_cast<uint64>(frame - g_frame_map);
+  if (frame < g_frame_map || idx >= g_num_frames)
+    kernel_panic("frame_to_pa: frame outside frame_map bounds");
+  return phys_addr(frame_index_to_pa(idx));
 }
 
-Frame* pa_to_frame(uint64 pa)
-{
-    uint64 idx = pa_to_frame_index(pa);
-    return &g_frame_map[idx];
+Frame *pa_to_frame(PhysAddr pa) {
+  uint64 idx = pa_to_frame_index(pa.raw);
+  return &g_frame_map[idx];
 }
 
-uint64 frame_to_va(const Frame* frame)
-{
-    return pa_to_va(frame_to_pa(frame));
+VirtAddr frame_to_va(const Frame *frame) {
+  return pa_to_va(frame_to_pa(frame));
 }
 
 /* ------------------------------------------------------------------ */
@@ -105,40 +99,37 @@ uint64 frame_to_va(const Frame* frame)
  * Find a frame's buddy at a given order.
  * Buddy PFN = PFN ^ (1 << order).
  */
-static inline Frame* buddy_of(Frame* frame, uint8 order)
-{
-    uint64 pfn = static_cast<uint64>(frame - g_frame_map);
-    uint64 buddy_pfn = pfn ^ (1ULL << order);
-    if (buddy_pfn >= g_num_frames)
-        return nullptr;
-    return &g_frame_map[buddy_pfn];
+static inline Frame *buddy_of(Frame *frame, uint8 order) {
+  uint64 pfn = static_cast<uint64>(frame - g_frame_map);
+  uint64 buddy_pfn = pfn ^ (1ULL << order);
+  if (buddy_pfn >= g_num_frames)
+    return nullptr;
+  return &g_frame_map[buddy_pfn];
 }
 
 /*
  * Insert a free block into the buddy free list at the given order.
  */
-static void buddy_list_add(Frame* frame, uint8 order)
-{
-    frame->state = FrameState::Free;
-    frame->order = order;
-    list_node* head = &g_free_areas[order];
-    list_node* node = &frame->free.free_link;
-    node->next = head->next;
-    node->prev = head;
-    head->next->prev = node;
-    head->next = node;
+static void buddy_list_add(Frame *frame, uint8 order) {
+  frame->state = FrameState::Free;
+  frame->order = order;
+  list_node *head = &g_free_areas[order];
+  list_node *node = &frame->free.free_link;
+  node->next = head->next;
+  node->prev = head;
+  head->next->prev = node;
+  head->next = node;
 }
 
 /*
  * Remove a block from whatever buddy free list it is on.
  */
-static void buddy_list_del(Frame* frame)
-{
-    list_node* node = &frame->free.free_link;
-    node->prev->next = node->next;
-    node->next->prev = node->prev;
-    node->next = node;
-    node->prev = node;
+static void buddy_list_del(Frame *frame) {
+  list_node *node = &frame->free.free_link;
+  node->prev->next = node->next;
+  node->next->prev = node->prev;
+  node->next = node;
+  node->prev = node;
 }
 
 /*
@@ -146,66 +137,64 @@ static void buddy_list_del(Frame* frame)
  * Splits higher-order blocks if no exact match is available.
  * Returns nullptr if no memory available.
  */
-static Frame* buddy_alloc(uint8 order)
-{
-    for (int cur = order; cur < kMaxOrder; cur++) {
-        list_node* head = &g_free_areas[cur];
-        if (head->next == head)
-            continue; // Empty at this order
-
-        // Pop the first block from this order's free list
-        list_node* node = head->next;
-        // Compute Frame* from the free_link member offset
-        Frame* frame = reinterpret_cast<Frame*>(
-            reinterpret_cast<uintptr>(node) -
-            __builtin_offsetof(Frame, free.free_link));
-        buddy_list_del(frame);
-
-        // Split down to the requested order
-        while (cur > order) {
-            cur--;
-            // The upper half becomes a free buddy
-            Frame* split = &g_frame_map[(frame - g_frame_map) + (1ULL << cur)];
-            buddy_list_add(split, static_cast<uint8>(cur));
-        }
-
-        frame->state = FrameState::Allocated;
-        frame->order = order;
-        frame->refcount.store(1, __ATOMIC_RELAXED);
-        return frame;
+static Frame *buddy_alloc(uint8 order) {
+  for (int cur = order; cur < kMaxOrder; cur++) {
+    list_node *head = &g_free_areas[cur];
+    if (head->next == head)
+      continue; // Empty at this order
+
+    // Pop the first block from this order's free list
+    list_node *node = head->next;
+    // Compute Frame* from the free_link member offset
+    Frame *frame =
+        reinterpret_cast<Frame *>(reinterpret_cast<uintptr>(node) -
+                                  __builtin_offsetof(Frame, free.free_link));
+    buddy_list_del(frame);
+
+    // Split down to the requested order
+    while (cur > order) {
+      cur--;
+      // The upper half becomes a free buddy
+      Frame *split = &g_frame_map[(frame - g_frame_map) + (1ULL << cur)];
+      buddy_list_add(split, static_cast<uint8>(cur));
     }
-    return nullptr;
+
+    frame->state = FrameState::Allocated;
+    frame->order = order;
+    frame->refcount.store(1, __ATOMIC_RELAXED);
+    return frame;
+  }
+  return nullptr;
 }
 
 /*
  * Return a block of 2^order frames to the buddy system.
  * Coalesces with its buddy recursively up to kMaxOrder-1.
  */
-static void buddy_free(Frame* frame, uint8 order)
-{
-    uint64 pfn = static_cast<uint64>(frame - g_frame_map);
-
-    while (order < kMaxOrder - 1) {
-        Frame* buddy = buddy_of(frame, order);
-        if (!buddy)
-            break;
-        // Buddy must be free AND at the same order to coalesce
-        if (buddy->state != FrameState::Free || buddy->order != order)
-            break;
-
-        // Remove buddy from its free list and merge
-        buddy_list_del(buddy);
-
-        // The merged block starts at the lower PFN
-        uint64 buddy_pfn = static_cast<uint64>(buddy - g_frame_map);
-        if (buddy_pfn < pfn) {
-            frame = buddy;
-            pfn = buddy_pfn;
-        }
-        order++;
+static void buddy_free(Frame *frame, uint8 order) {
+  uint64 pfn = static_cast<uint64>(frame - g_frame_map);
+
+  while (order < kMaxOrder - 1) {
+    Frame *buddy = buddy_of(frame, order);
+    if (!buddy)
+      break;
+    // Buddy must be free AND at the same order to coalesce
+    if (buddy->state != FrameState::Free || buddy->order != order)
+      break;
+
+    // Remove buddy from its free list and merge
+    buddy_list_del(buddy);
+
+    // The merged block starts at the lower PFN
+    uint64 buddy_pfn = static_cast<uint64>(buddy - g_frame_map);
+    if (buddy_pfn < pfn) {
+      frame = buddy;
+      pfn = buddy_pfn;
     }
+    order++;
+  }
 
-    buddy_list_add(frame, order);
+  buddy_list_add(frame, order);
 }
 
 /* ------------------------------------------------------------------ */
@@ -216,166 +205,157 @@ static void buddy_free(Frame* frame, uint8 order)
  * Refill the PCP from the global buddy system.
  * Acquires g_pmm_lock internally.
  */
-static void refill_pcp(PerCpuCache* pcp)
-{
-    LockGuard guard(g_pmm_lock);
-    for (uint32 i = 0; i < g_pcp_batch_size; i++) {
-        Frame* f = buddy_alloc(0);
-        if (!f)
-            break;
-        // Put onto PCP list (reuse the free_link for PCP chain)
-        f->state = FrameState::PcpCached;
-        list_node* node = &f->free.free_link;
-        node->next = pcp->free_list.next;
-        node->prev = &pcp->free_list;
-        pcp->free_list.next->prev = node;
-        pcp->free_list.next = node;
-        pcp->count++;
-    }
+static void refill_pcp(PerCpuCache *pcp) {
+  LockGuard guard(g_pmm_lock);
+  for (uint32 i = 0; i < g_pcp_batch_size; i++) {
+    Frame *f = buddy_alloc(0);
+    if (!f)
+      break;
+    // Put onto PCP list (reuse the free_link for PCP chain)
+    f->state = FrameState::PcpCached;
+    list_node *node = &f->free.free_link;
+    node->next = pcp->free_list.next;
+    node->prev = &pcp->free_list;
+    pcp->free_list.next->prev = node;
+    pcp->free_list.next = node;
+    pcp->count++;
+  }
 }
 
 /*
  * Drain half the PCP back to the global buddy system.
  * Acquires g_pmm_lock internally.
  */
-static void drain_pcp(PerCpuCache* pcp)
-{
-    LockGuard guard(g_pmm_lock);
-    int to_drain = pcp->count / 2;
-    for (int i = 0; i < to_drain; i++) {
-        list_node* head = &pcp->free_list;
-        if (head->next == head)
-            break;
-        list_node* node = head->next;
-        // Unlink from PCP
-        node->prev->next = node->next;
-        node->next->prev = node->prev;
-        pcp->count--;
-
-        Frame* f = reinterpret_cast<Frame*>(
-            reinterpret_cast<uintptr>(node) -
-            __builtin_offsetof(Frame, free.free_link));
-        buddy_free(f, 0);
-    }
+static void drain_pcp(PerCpuCache *pcp) {
+  LockGuard guard(g_pmm_lock);
+  int to_drain = pcp->count / 2;
+  for (int i = 0; i < to_drain; i++) {
+    list_node *head = &pcp->free_list;
+    if (head->next == head)
+      break;
+    list_node *node = head->next;
+    // Unlink from PCP
+    node->prev->next = node->next;
+    node->next->prev = node->prev;
+    pcp->count--;
+
+    Frame *f =
+        reinterpret_cast<Frame *>(reinterpret_cast<uintptr>(node) -
+                                  __builtin_offsetof(Frame, free.free_link));
+    buddy_free(f, 0);
+  }
 }
 
 /* ------------------------------------------------------------------ */
 /*  Public API: single-frame (fast path via PCP)                      */
 /* ------------------------------------------------------------------ */
 
-Frame* pmm_alloc_frame()
-{
-    /*
-     * Disable interrupts to prevent reentry on the same CPU
-     * (e.g. timer IRQ handler calling pmm_alloc_frame while we
-     * are mid-way through manipulating the PCP list).
-     */
-    bool was_on = arch::cpu::intr_enabled();
-    arch::cpu::intr_off();
-
-    uint64 cpu = arch::cpu::id();
-    PerCpuCache* pcp = &g_pcp[cpu];
-
-    if (pcp->count == 0)
-        refill_pcp(pcp);
-
-    Frame* frame = nullptr;
-    if (pcp->count > 0) {
-        list_node* node = pcp->free_list.next;
-        // Unlink
-        node->prev->next = node->next;
-        node->next->prev = node->prev;
-        pcp->count--;
-
-        frame = reinterpret_cast<Frame*>(
-            reinterpret_cast<uintptr>(node) -
-            __builtin_offsetof(Frame, free.free_link));
-
-        // Validate the frame is within the frame_map bounds
-        if (frame < g_frame_map || frame >= g_frame_map + g_num_frames) {
-            kernel_panic("pmm_alloc_frame: PCP returned out-of-bounds frame");
-        }
+Frame *pmm_alloc_frame() {
+  /*
+   * Disable interrupts to prevent reentry on the same CPU
+   * (e.g. timer IRQ handler calling pmm_alloc_frame while we
+   * are mid-way through manipulating the PCP list).
+   */
+  bool was_on = arch::cpu::intr_enabled();
+  arch::cpu::intr_off();
+
+  uint64 cpu = arch::cpu::id();
+  PerCpuCache *pcp = &g_pcp[cpu];
+
+  if (pcp->count == 0)
+    refill_pcp(pcp);
+
+  Frame *frame = nullptr;
+  if (pcp->count > 0) {
+    list_node *node = pcp->free_list.next;
+    // Unlink
+    node->prev->next = node->next;
+    node->next->prev = node->prev;
+    pcp->count--;
+
+    frame =
+        reinterpret_cast<Frame *>(reinterpret_cast<uintptr>(node) -
+                                  __builtin_offsetof(Frame, free.free_link));
 
-        frame->state = FrameState::Allocated;
-        frame->order = 0;
-        frame->refcount.store(1, __ATOMIC_RELAXED);
+    // Validate the frame is within the frame_map bounds
+    if (frame < g_frame_map || frame >= g_frame_map + g_num_frames) {
+      kernel_panic("pmm_alloc_frame: PCP returned out-of-bounds frame");
     }
 
-    if (was_on)
-        arch::cpu::intr_on();
-    return frame;
+    frame->state = FrameState::Allocated;
+    frame->order = 0;
+    frame->refcount.store(1, __ATOMIC_RELAXED);
+  }
+
+  if (was_on)
+    arch::cpu::intr_on();
+  return frame;
 }
 
-void pmm_free_frame(Frame* frame)
-{
-    if (!frame)
-        kernel_panic("pmm_free_frame: null frame");
+void pmm_free_frame(Frame *frame) {
+  if (!frame)
+    kernel_panic("pmm_free_frame: null frame");
 
-    bool was_on = arch::cpu::intr_enabled();
-    arch::cpu::intr_off();
+  bool was_on = arch::cpu::intr_enabled();
+  arch::cpu::intr_off();
 
-    uint64 cpu = arch::cpu::id();
-    PerCpuCache* pcp = &g_pcp[cpu];
+  uint64 cpu = arch::cpu::id();
+  PerCpuCache *pcp = &g_pcp[cpu];
 
-    // Push onto PCP
-    frame->state = FrameState::PcpCached;
-    frame->order = 0;
-    list_node* node = &frame->free.free_link;
-    node->next = pcp->free_list.next;
-    node->prev = &pcp->free_list;
-    pcp->free_list.next->prev = node;
-    pcp->free_list.next = node;
-    pcp->count++;
+  // Push onto PCP
+  frame->state = FrameState::PcpCached;
+  frame->order = 0;
+  list_node *node = &frame->free.free_link;
+  node->next = pcp->free_list.next;
+  node->prev = &pcp->free_list;
+  pcp->free_list.next->prev = node;
+  pcp->free_list.next = node;
+  pcp->count++;
 
-    // Drain if over high watermark
-    if (pcp->count >= g_pcp_high_watermark)
-        drain_pcp(pcp);
+  // Drain if over high watermark
+  if (pcp->count >= g_pcp_high_watermark)
+    drain_pcp(pcp);
 
-    if (was_on)
-        arch::cpu::intr_on();
+  if (was_on)
+    arch::cpu::intr_on();
 }
 
 /* ------------------------------------------------------------------ */
 /*  Public API: multi-frame (slow path, direct buddy)                 */
 /* ------------------------------------------------------------------ */
 
-Frame* pmm_alloc_frames(uint8 order)
-{
-    if (order >= kMaxOrder)
-        return nullptr;
-    LockGuard guard(g_pmm_lock);
-    return buddy_alloc(order);
+Frame *pmm_alloc_frames(uint8 order) {
+  if (order >= kMaxOrder)
+    return nullptr;
+  LockGuard guard(g_pmm_lock);
+  return buddy_alloc(order);
 }
 
-void pmm_free_frames(Frame* frame, uint8 order)
-{
-    if (!frame)
-        kernel_panic("pmm_free_frames: null frame");
-    if (order >= kMaxOrder)
-        kernel_panic("pmm_free_frames: order out of range");
-    LockGuard guard(g_pmm_lock);
-    buddy_free(frame, order);
+void pmm_free_frames(Frame *frame, uint8 order) {
+  if (!frame)
+    kernel_panic("pmm_free_frames: null frame");
+  if (order >= kMaxOrder)
+    kernel_panic("pmm_free_frames: order out of range");
+  LockGuard guard(g_pmm_lock);
+  buddy_free(frame, order);
 }
 
 /* ------------------------------------------------------------------ */
 /*  Public API: reference counting                                    */
 /* ------------------------------------------------------------------ */
 
-void frame_incref(Frame* frame)
-{
-    frame->refcount.fetch_add(1);
-}
-
-void frame_decref(Frame* frame)
-{
-    uint32 old = frame->refcount.fetch_sub(1);
-    if (old == 1) {
-        // Refcount dropped to 0 — return to pool based on originally allocated order
-        if (frame->order == 0)
-            pmm_free_frame(frame);
-        else
-            pmm_free_frames(frame, frame->order);
-    }
+void frame_incref(Frame *frame) { frame->refcount.fetch_add(1); }
+
+void frame_decref(Frame *frame) {
+  uint32 old = frame->refcount.fetch_sub(1);
+  if (old == 1) {
+    // Refcount dropped to 0 — return to pool based on originally allocated
+    // order
+    if (frame->order == 0)
+      pmm_free_frame(frame);
+    else
+      pmm_free_frames(frame, frame->order);
+  }
 }
 
 /* ------------------------------------------------------------------ */
@@ -386,110 +366,104 @@ void frame_decref(Frame* frame)
  * Helper: zero out a range of bytes.  We cannot use memset here
  * because lib/string may not be linked yet.
  */
-static void zero_range(void* start, uint64 bytes)
-{
-    uint8* p = static_cast<uint8*>(start);
-    for (uint64 i = 0; i < bytes; i++)
-        p[i] = 0;
+static void zero_range(void *start, uint64 bytes) {
+  uint8 *p = static_cast<uint8 *>(start);
+  for (uint64 i = 0; i < bytes; i++)
+    p[i] = 0;
 }
 
-void pmm_init()
-{
-    /* 1. Query FDT for physical memory layout */
-    uint64 mem_base = 0, mem_size = 0;
-    fdt_early_get_mem_info(&mem_base, &mem_size);
-
-    if (mem_size == 0)
-        kernel_panic("pmm_init: FDT reports zero memory");
-
-    g_mem_base = mem_base;
-    g_mem_size = mem_size;
-
-    uint64 mem_end = mem_base + mem_size;
-    g_num_frames = mem_size / kPageSize;
-
-    /*
-     * Dynamically calculate PCP parameters based on memory size.
-     * Rule of thumb: batch_size = max(16, total_pages / (cpu_count * 1024))
-     * This avoids depleting global memory on small RAM systems.
-     */
-    uint32 calc_batch = static_cast<uint32>(g_num_frames / (kMaxCpus * 1024));
-    g_pcp_batch_size = (calc_batch > 16) ? calc_batch : 16;
-    g_pcp_high_watermark = g_pcp_batch_size * 2;
-
-    /*
-     * 2. Place the frame_map array.
-     *
-     * The frame_map lives right after the kernel image (_kernel_end),
-     * page-aligned.  We compute its size and mark that region as used.
-     *
-     * _kernel_end is a virtual address; convert to physical for arithmetic,
-     * then back to virtual for the actual pointer.
-     */
-    /*
-     * _kernel_end is linked at a virtual address (0xFFFFFFC0...),
-     * but if the kernel is still running via identity mapping (before
-     * vmm_init removes it), PC-relative `la` resolves to the physical
-     * address instead.  Detect which case we're in.
-     */
-    uint64 kernel_end_raw = reinterpret_cast<uint64>(_kernel_end);
-    uint64 kernel_end_pa = ensure_pa(kernel_end_raw);
-    uint64 frame_map_pa = (kernel_end_pa + kPageSize - 1) & ~(kPageSize - 1);
-    uint64 frame_map_bytes = g_num_frames * sizeof(Frame);
-    uint64 frame_map_end_pa = (frame_map_pa + frame_map_bytes + kPageSize - 1)
-                              & ~(kPageSize - 1);
-
-    if (frame_map_end_pa > mem_end)
-        kernel_panic("pmm_init: not enough memory for frame_map");
-
-    g_frame_map = reinterpret_cast<Frame*>(pa_to_va(frame_map_pa));
-
-    /* 3. Zero out the entire frame_map */
-    zero_range(g_frame_map, frame_map_bytes);
-
-    /* 4. Initialize buddy free list heads */
-    g_pmm_lock.init("pmm");
-    for (int i = 0; i < kMaxOrder; i++)
-        g_free_areas[i].init();
-
-    /* 5. Initialize PCP list heads */
-    for (int i = 0; i < kMaxCpus; i++) {
-        g_pcp[i].free_list.init();
-        g_pcp[i].count = 0;
-    }
-
-    /*
-     * 6. Feed all free pages into the buddy system.
-     *
-     * Pages before frame_map_end (kernel + frame_map) are reserved.
-     * We iterate from the first free page to the end of memory,
-     * inserting at the highest possible order for efficiency.
-     */
-    uint64 free_start_pfn = pa_to_frame_index(frame_map_end_pa);
-    uint64 end_pfn = g_num_frames;
-
-    uint64 pfn = free_start_pfn;
-    while (pfn < end_pfn) {
-        // Find the largest order block that is naturally aligned at this PFN
-        uint8 order = 0;
-        for (int o = kMaxOrder - 1; o >= 0; o--) {
+void pmm_init() {
+  /* 1. Get memory layout from memblock (already scanned FDT) */
+  g_mem_base = memblock_get_mem_base();
+  g_mem_size = memblock_get_mem_size();
+  g_num_frames = vmm_get_num_frames();
+
+  if (g_mem_size == 0) {
+    kernel_panic("pmm_init: memblock reports zero memory");
+  }
+
+  /* 2. PCP parameters */
+  uint32 calc_batch = static_cast<uint32>(g_num_frames / (kMaxCpus * 1024));
+  g_pcp_batch_size = (calc_batch > 16) ? calc_batch : 16;
+  g_pcp_high_watermark = g_pcp_batch_size * 2;
+
+  /* 3. frame_map was allocated by vmm_init at vmemmap VA.
+   *    No relocation needed — frame_map starts life at the right VA. */
+  g_frame_map_pa = vmm_get_frame_map_pa();
+  g_frame_map_bytes = vmm_get_frame_map_size();
+  g_frame_map = reinterpret_cast<Frame *>(arch::kVmemmapBase);
+
+  /* 4. Zero out the entire frame_map */
+  zero_range(g_frame_map, g_frame_map_bytes);
+
+  /* 5. Initialize buddy free list heads */
+  g_pmm_lock.init("pmm");
+  for (int i = 0; i < kMaxOrder; i++) {
+    g_free_areas[i].init();
+  }
+
+  /* 6. Initialize PCP list heads */
+  for (int i = 0; i < kMaxCpus; i++) {
+    g_pcp[i].free_list.init();
+    g_pcp[i].count = 0;
+  }
+
+  /* 7. Feed free regions from memblock into the buddy system.
+   *    memblock knows about all reserved regions (firmware, kernel,
+   *    page table pages, frame_map).  Everything else is free. */
+  struct FeedCtx {
+    uint64 total_pages;
+  };
+  FeedCtx ctx{0};
+
+  memblock_for_each_free(
+      [](uint64 base, uint64 size, void *arg) {
+        auto *c = static_cast<FeedCtx *>(arg);
+        uint64 start_pfn = (base - memblock_get_mem_base()) / kPageSize;
+        uint64 end_pfn = start_pfn + size / kPageSize;
+
+        /* Access frame_map via the extern g_frame_map (already at vmemmap VA)
+         */
+        extern Frame *g_frame_map;
+        extern list_node g_free_areas[];
+
+        uint64 pfn = start_pfn;
+        while (pfn < end_pfn) {
+          uint8 order = 0;
+          for (int o = kMaxOrder - 1; o >= 0; o--) {
             uint64 block_size = 1ULL << o;
             if ((pfn & (block_size - 1)) == 0 && pfn + block_size <= end_pfn) {
-                order = static_cast<uint8>(o);
-                break;
+              order = static_cast<uint8>(o);
+              break;
             }
+          }
+
+          Frame *f = &g_frame_map[pfn];
+          f->state = FrameState::Free;
+          f->order = order;
+          f->refcount.store(0, __ATOMIC_RELAXED);
+          f->free.free_link.init();
+
+          /* Inline buddy_list_add since we're inside a lambda */
+          list_node *head = &g_free_areas[order];
+          list_node *node = &f->free.free_link;
+          node->next = head->next;
+          node->prev = head;
+          head->next->prev = node;
+          head->next = node;
+
+          pfn += (1ULL << order);
+          c->total_pages += (1ULL << order);
         }
+      },
+      &ctx);
 
-        Frame* f = &g_frame_map[pfn];
-        f->state = FrameState::Free;
-        f->order = order;
-        f->refcount.store(0, __ATOMIC_RELAXED);
-        f->free.free_link.init();
-        buddy_list_add(f, order);
+  kprintf("[pmm] buddy: %llu free pages (%llu MB)\n", ctx.total_pages,
+          (ctx.total_pages * kPageSize) / (1024 * 1024));
 
-        pfn += (1ULL << order);
-    }
+  /* 8. Retire memblock — all future allocations go through PMM */
+  memblock_retire();
 }
 
-uint64 pmm_get_mem_base() { return g_mem_base; }
+PhysAddr pmm_get_mem_base() { return phys_addr(g_mem_base); }
 uint64 pmm_get_mem_size() { return g_mem_size; }