feat: enhance address conversion functions and introduce address types

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

4 files changed, 72 insertions, 32 deletions

diff --git a/na/include/lume/addr.h b/na/include/lume/addr.h
index a873c1b..9c20d66 100644
--- a/na/include/lume/addr.h
+++ b/na/include/lume/addr.h
@@ -14,58 +14,61 @@
 
 #include <lume/types.h>
 #include <lume/config.h>
+#include <lume/addr_types.h>
 #include <arch/config.h>
 
-/* Convert physical address to kernel virtual address */
-inline uint64 pa_to_va(uint64 pa)
+/* Convert physical address to kernel virtual address (direct map). */
+inline VirtAddr pa_to_va(PhysAddr pa)
 {
-    return pa + arch::kVAOffset;
+    return VirtAddr{pa.raw + arch::g_direct_map_base};
 }
 
-/* Convert kernel virtual address to physical address */
-inline uint64 va_to_pa(uint64 va)
+/* Convert kernel virtual address to physical address (direct map). */
+inline PhysAddr va_to_pa(VirtAddr va)
 {
-    return va - arch::kVAOffset;
+    return PhysAddr{va.raw - arch::g_direct_map_base};
 }
 
-/* Check if an address is a kernel virtual address */
-inline bool is_kernel_va(uint64 addr)
+/* Align up to page boundary */
+inline uint64 page_align_up(uint64 addr)
 {
-    return addr >= arch::kVAOffset;
+    return (addr + kPageSize - 1) & ~(kPageSize - 1);
 }
 
-/* Convert an address that might be VA or PA to PA */
-inline uint64 ensure_pa(uint64 addr)
+/* Align down to page boundary */
+inline uint64 page_align_down(uint64 addr)
 {
-    return is_kernel_va(addr) ? va_to_pa(addr) : addr;
+    return addr & ~(kPageSize - 1);
 }
 
-/* Convert an address that might be VA or PA to VA */
-inline uint64 ensure_va(uint64 addr)
+inline PhysAddr page_align_up(PhysAddr addr)
 {
-    return is_kernel_va(addr) ? addr : pa_to_va(addr);
+    return PhysAddr{page_align_up(addr.raw)};
 }
 
-/* Align up to page boundary */
-inline uint64 page_align_up(uint64 addr)
+inline PhysAddr page_align_down(PhysAddr addr)
 {
-    return (addr + kPageSize - 1) & ~(kPageSize - 1);
+    return PhysAddr{page_align_down(addr.raw)};
 }
 
-/* Align down to page boundary */
-inline uint64 page_align_down(uint64 addr)
+inline VirtAddr page_align_up(VirtAddr addr)
 {
-    return addr & ~(kPageSize - 1);
+    return VirtAddr{page_align_up(addr.raw)};
+}
+
+inline VirtAddr page_align_down(VirtAddr addr)
+{
+    return VirtAddr{page_align_down(addr.raw)};
 }
 
 /* Convert physical address to page frame number (PFN) */
-inline uint64 pa_to_pfn(uint64 pa)
+inline Pfn pa_to_pfn(PhysAddr pa)
 {
-    return pa / kPageSize;
+    return Pfn{pa.raw / kPageSize};
 }
 
 /* Convert page frame number (PFN) to physical address */
-inline uint64 pfn_to_pa(uint64 pfn)
+inline PhysAddr pfn_to_pa(Pfn pfn)
 {
-    return pfn * kPageSize;
+    return PhysAddr{pfn.raw * kPageSize};
 }
diff --git a/na/include/lume/addr_types.h b/na/include/lume/addr_types.h
new file mode 100644
index 0000000..91a0779
--- /dev/null
+++ b/na/include/lume/addr_types.h
@@ -0,0 +1,29 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+#pragma once
+
+#include <lume/types.h>
+
+struct PhysAddr {
+    uint64 raw;
+};
+
+struct VirtAddr {
+    uint64 raw;
+};
+
+struct Pfn {
+    uint64 raw;
+};
+
+inline constexpr PhysAddr phys_addr(uint64 raw) { return PhysAddr{raw}; }
+inline constexpr VirtAddr virt_addr(uint64 raw) { return VirtAddr{raw}; }
+inline constexpr Pfn pfn(uint64 raw) { return Pfn{raw}; }
+
+inline constexpr bool operator==(PhysAddr a, PhysAddr b) { return a.raw == b.raw; }
+inline constexpr bool operator!=(PhysAddr a, PhysAddr b) { return a.raw != b.raw; }
+inline constexpr bool operator==(VirtAddr a, VirtAddr b) { return a.raw == b.raw; }
+inline constexpr bool operator!=(VirtAddr a, VirtAddr b) { return a.raw != b.raw; }
+inline constexpr bool operator==(Pfn a, Pfn b) { return a.raw == b.raw; }
+inline constexpr bool operator!=(Pfn a, Pfn b) { return a.raw != b.raw; }
diff --git a/na/include/lume/pmm.h b/na/include/lume/pmm.h
index b970404..a9249ac 100644
--- a/na/include/lume/pmm.h
+++ b/na/include/lume/pmm.h
@@ -21,6 +21,7 @@
 
 #include <lume/types.h>
 #include <lume/frame.h>
+#include <lume/addr_types.h>
 
 /* Buddy system constants */
 inline constexpr int kMaxOrder = 11;           // Orders 0..10 → 4KB..4MB
@@ -34,9 +35,10 @@ extern uint32 g_pcp_batch_size;
 /* ================================================================== */
 
 /*
- * Initialize the PMM.  Called by BSP after fdt_init().
- * Discovers memory via FDT, places frame_map, and populates buddy lists.
- * Panics on failure (no memory, bad FDT, etc.).
+ * Initialize the PMM.  Called by BSP after vmm_init() + memblock.
+ * frame_map is already mapped at VMEMMAP VA by vmm_init.
+ * Populates buddy lists from memblock's free regions, then retires memblock.
+ * Panics on failure.
  */
 void pmm_init();
 
@@ -64,14 +66,14 @@ void frame_decref(Frame* frame);
 /* --- Address conversion --- */
 
 // Convert a Frame* to its physical address.
-uint64 frame_to_pa(const Frame* frame);
+PhysAddr frame_to_pa(const Frame* frame);
 
 // Convert a physical address to its Frame*.
-Frame* pa_to_frame(uint64 pa);
+Frame* pa_to_frame(PhysAddr pa);
 
 // Convert a Frame* to a kernel virtual address.
-uint64 frame_to_va(const Frame* frame);
+VirtAddr frame_to_va(const Frame* frame);
 
 // Query PMM memory range (for VMM kernel mapping)
-uint64 pmm_get_mem_base();
+PhysAddr pmm_get_mem_base();
 uint64 pmm_get_mem_size();
diff --git a/na/include/lume/vmm.h b/na/include/lume/vmm.h
index ea32332..ec26d6a 100644
--- a/na/include/lume/vmm.h
+++ b/na/include/lume/vmm.h
@@ -72,6 +72,12 @@ void vmm_init_ap();
 /* Map MMIO device region to the active kernel page table */
 void vmm_map_kernel_mmio(uint64 pa, uint64 size);
 
+/* Query frame_map allocation (set by vmm_init, read by pmm_init).
+ * Returns the PA, byte size, and frame count of the vmemmap region. */
+uint64 vmm_get_frame_map_pa();
+uint64 vmm_get_frame_map_size();
+uint64 vmm_get_num_frames();
+
 /* User space memory management (VMA & Demand Paging) */
 int vmm_map_user(VmSpace* space, uint64 va, uint64 len, uint64 perm);
 int vmm_unmap_user(VmSpace* space, uint64 va, uint64 len);