From 9fe5e02cf615ad121c4460fbaaac2787d1e09c0e Mon Sep 17 00:00:00 2001
From: Minep <zelong56@gmail.com>
Date: Sun, 27 Feb 2022 23:46:51 +0000
Subject: [PATCH 1/1] vmm - allow remapping vmm,page - rename some symbols that
 seems make more sense boot.S - make things more predictable hal - more
 abstraction on the daily-use functionalities. dmm - a place for dynamic heap
 allocation

---
 lunaix-os/arch/x86/boot.S           |   2 +
 lunaix-os/arch/x86/hhk.c            |  20 +--
 lunaix-os/hal/cpu.c                 |  40 +-----
 lunaix-os/includes/hal/ahci.h       |   5 +
 lunaix-os/includes/hal/cpu.h        |  68 +++++++---
 lunaix-os/includes/hal/io.h         | 103 +++++++++-----
 lunaix-os/includes/hal/pic.h        |   5 +
 lunaix-os/includes/lunaix/mm/dmm.h  |  16 +++
 lunaix-os/includes/lunaix/mm/page.h |  38 ++++--
 lunaix-os/includes/lunaix/mm/vmm.h  |  68 +++++++---
 lunaix-os/kernel/k_init.c           |   4 +-
 lunaix-os/kernel/mm/dmm.c           |  19 +++
 lunaix-os/kernel/mm/vmm.c           | 200 ++++++++++++++++------------
 13 files changed, 370 insertions(+), 218 deletions(-)
 create mode 100644 lunaix-os/includes/hal/ahci.h
 create mode 100644 lunaix-os/includes/hal/pic.h
 create mode 100644 lunaix-os/includes/lunaix/mm/dmm.h
 create mode 100644 lunaix-os/kernel/mm/dmm.c

diff --git a/lunaix-os/arch/x86/boot.S b/lunaix-os/arch/x86/boot.S
index 5f832aa..ddcd891 100644
--- a/lunaix-os/arch/x86/boot.S
+++ b/lunaix-os/arch/x86/boot.S
@@ -41,6 +41,8 @@
                                  * linker more knowledge about the label 
                                  */
     start_: 
+        cld
+        cli
         movl $stack_top, %esp
 
         subl $16, %esp
diff --git a/lunaix-os/arch/x86/hhk.c b/lunaix-os/arch/x86/hhk.c
index c0e5f9a..ed3709f 100644
--- a/lunaix-os/arch/x86/hhk.c
+++ b/lunaix-os/arch/x86/hhk.c
@@ -29,26 +29,26 @@ extern uint8_t _k_stack;
 
 void 
 _init_page(ptd_t* ptd) {
-    SET_PDE(ptd, 0, PDE(PG_PRESENT, ptd + PG_MAX_ENTRIES))
+    SET_PDE(ptd, 0, NEW_L1_ENTRY(PG_PRESENT, ptd + PG_MAX_ENTRIES))
     
     // 对低1MiB空间进行对等映射（Identity mapping），也包括了我们的VGA，方便内核操作。
     for (uint32_t i = 0; i < 256; i++)
     {
-        SET_PTE(ptd, PG_TABLE_IDENTITY, i, PTE(PG_PREM_RW, (i << 12)))
+        SET_PTE(ptd, PG_TABLE_IDENTITY, i, NEW_L2_ENTRY(PG_PREM_RW, (i << PG_SIZE_BITS)))
     }
 
     // 对等映射我们的hhk_init，这样一来，当分页与地址转换开启后，我们依然能够照常执行最终的 jmp 指令来跳转至
     //  内核的入口点
     for (uint32_t i = 0; i < HHK_PAGE_COUNT; i++)
     {
-        SET_PTE(ptd, PG_TABLE_IDENTITY, 256 + i, PTE(PG_PREM_RW, 0x100000 + (i << 12)))
+        SET_PTE(ptd, PG_TABLE_IDENTITY, 256 + i, NEW_L2_ENTRY(PG_PREM_RW, 0x100000 + (i << PG_SIZE_BITS)))
     }
     
     // --- 将内核重映射至高半区 ---
     
     // 这里是一些计算，主要是计算应当映射进的 页目录 与 页表 的条目索引（Entry Index）
-    uint32_t kernel_pde_index = PD_INDEX(sym_val(__kernel_start));
-    uint32_t kernel_pte_index = PT_INDEX(sym_val(__kernel_start));
+    uint32_t kernel_pde_index = L1_INDEX(sym_val(__kernel_start));
+    uint32_t kernel_pte_index = L2_INDEX(sym_val(__kernel_start));
     uint32_t kernel_pg_counts = KERNEL_PAGE_COUNT;
     
     // 将内核所需要的页表注册进页目录
@@ -59,12 +59,12 @@ _init_page(ptd_t* ptd) {
         SET_PDE(
             ptd, 
             kernel_pde_index + i,   
-            PDE(PG_PREM_RW, PT_ADDR(ptd, PG_TABLE_KERNEL + i))
+            NEW_L1_ENTRY(PG_PREM_RW, PT_ADDR(ptd, PG_TABLE_KERNEL + i))
         )
     }
     
     // 首先，检查内核的大小是否可以fit进我们这几个表（12MiB）
-    if (kernel_pg_counts > (PG_TABLE_STACK - PG_TABLE_KERNEL) * 1024) {
+    if (kernel_pg_counts > (PG_TABLE_STACK - PG_TABLE_KERNEL) * PG_MAX_ENTRIES) {
         // ERROR: require more pages
         //  here should do something else other than head into blocking
         while (1);
@@ -80,15 +80,15 @@ _init_page(ptd_t* ptd) {
             ptd, 
             PG_TABLE_KERNEL, 
             kernel_pte_index + i, 
-            PTE(PG_PREM_RW, kernel_pm + (i << 12))
+            NEW_L2_ENTRY(PG_PREM_RW, kernel_pm + (i << PG_SIZE_BITS))
         )
     }
 
     // 最后一个entry用于循环映射
     SET_PDE(
         ptd,
-        1023,
-        PDE(T_SELF_REF_PERM, ptd)
+        PG_MAX_ENTRIES - 1,
+        NEW_L1_ENTRY(T_SELF_REF_PERM, ptd)
     );
 }
 
diff --git a/lunaix-os/hal/cpu.c b/lunaix-os/hal/cpu.c
index d57ae26..ad1b0a5 100644
--- a/lunaix-os/hal/cpu.c
+++ b/lunaix-os/hal/cpu.c
@@ -38,42 +38,4 @@ void cpu_get_brand(char* brand_out) {
         j+=4;
     }
     brand_out[48] = '\0';
-}
-
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wreturn-type"
-reg32 cpu_r_cr0() {
-    asm volatile ("mov %cr0, %eax");
-}
-
-reg32 cpu_r_cr2() {
-    asm volatile ("mov %cr2, %eax");
-}
-
-reg32 cpu_r_cr3() {
-    asm volatile ("mov %cr3, %eax");
-}
-#pragma GCC diagnostic push
-
-void cpu_w_cr0(reg32 v) {
-    asm volatile (
-        "mov %0, %%cr0"
-        :: "r"(v)
-    );
-}
-
-void cpu_w_cr2(reg32 v) {
-    asm volatile (
-        "mov %0, %%cr2"
-        :: "r"(v)
-    );
-}
-
-void cpu_w_cr3(reg32 v) {
-    asm volatile (
-        "mov %0, %%cr3"
-        :: "r"(v)
-    );
-}
-
-
+}
\ No newline at end of file
diff --git a/lunaix-os/includes/hal/ahci.h b/lunaix-os/includes/hal/ahci.h
new file mode 100644
index 0000000..fc457eb
--- /dev/null
+++ b/lunaix-os/includes/hal/ahci.h
@@ -0,0 +1,5 @@
+#ifndef __LUNAIX_AHCI_H
+#define __LUNAIX_AHCI_H
+// TODO: AHCI Drivers
+
+#endif /* __LUNAIX_AHCI_H */
diff --git a/lunaix-os/includes/hal/cpu.h b/lunaix-os/includes/hal/cpu.h
index af8863e..6381143 100644
--- a/lunaix-os/includes/hal/cpu.h
+++ b/lunaix-os/includes/hal/cpu.h
@@ -1,36 +1,72 @@
 #ifndef __LUNAIX_CPU_H
 #define __LUNAIX_CPU_H
 
+#include <stdint.h>
+
 typedef unsigned int reg32;
 typedef unsigned short reg16;
 
 typedef struct {
+    reg32 eax;
+    reg32 ebx;
+    reg32 ecx;
+    reg32 edx;
     reg32 edi;
-    reg32 esi;
     reg32 ebp;
+    reg32 esi;
     reg32 esp;
-    reg32 ebx;
-    reg32 edx;
-    reg32 ecx;
-    reg32 eax;
+    reg32 cs;
+    reg32 eip;
 } __attribute__((packed)) registers;
 
-reg32 cpu_r_cr0();
-
-reg32 cpu_r_cr2();
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wreturn-type"
+static inline reg32 cpu_rcr0() {
+    asm volatile ("mov %cr0, %eax");
+}
 
-reg32 cpu_r_cr3();
+static inline reg32 cpu_rcr2() {
+    asm volatile ("mov %cr2, %eax");
+}
 
-void cpu_w_cr0(reg32 v);
+static inline reg32 cpu_rcr3() {
+    asm volatile ("mov %cr3, %eax");
+}
+#pragma GCC diagnostic pop
 
-void cpu_w_cr2(reg32 v);
+static inline void cpu_lcr0(reg32 v) {
+    asm volatile (
+        "mov %0, %%cr0"
+        :: "r"(v)
+    );
+}
 
-void cpu_w_cr3(reg32 v);
+static inline void cpu_lcr2(reg32 v) {
+    asm volatile (
+        "mov %0, %%cr2"
+        :: "r"(v)
+    );
+}
 
-void cpu_get_model(char* model_out);
+static inline void cpu_lcr3(reg32 v) {
+    asm volatile (
+        "mov %0, %%cr3"
+        :: "r"(v)
+    );
+}
 
-int cpu_brand_string_supported();
+static inline void cpu_invplg(void* va) {
+    __asm__("invlpg (%0)" ::"r"((uintptr_t)va) : "memory");
+}
 
-void cpu_get_brand(char* brand_out);
+static inline void cpu_invtlb() {
+    reg32 interm;
+    __asm__(
+        "movl %%cr3, %0\n"
+        "movl %0, %%cr3"
+        :"=r"(interm)
+        :"r"(interm)
+    );
+}
 
-#endif /* __LUNAIX_CPU_H */
+#endif
\ No newline at end of file
diff --git a/lunaix-os/includes/hal/io.h b/lunaix-os/includes/hal/io.h
index 0f5e9d6..cbcf2bd 100644
--- a/lunaix-os/includes/hal/io.h
+++ b/lunaix-os/includes/hal/io.h
@@ -3,39 +3,76 @@
 
 #include <stdint.h>
 
-void io_port_wb(uint8_t port, uint8_t value) {
-    asm volatile (
-        "movb %0, %%al\n"
-        "movb %1, %%dx\n"
-        "out %%al, %%dx\n"
-        :: "r"(value) "r"(port)
-    );
-}
-
-void io_port_wl(uint8_t port, uint32_t value) {
-    asm volatile (
-        "movl %0, %%eax\n"
-        "movb %1, %%dx\n"
-        "out %%eax, %%dx\n"
-        :: "r"(value) "r"(port)
-    );
-}
-
-uint8_t io_port_rb(uint8_t port) {
-    asm volatile (
-        "movb $0, %%eax\n"
-        "movb %0, %%dx\n"
-        "in %%dx, %%al\n"
-        :: "r"(port)
-    );
-}
-
-uint32_t io_port_rl(uint8_t port) {
-    asm volatile (
-        "movb %0, %%dx\n"
-        "in %%dx, %%eax\n"
-        :: "r"(port)
-    );
+static inline uint8_t io_inb(int port) {
+    uint8_t data;
+    asm volatile("inb %w1,%0" : "=a" (data) : "d" (port));
+    return data;
+}
+
+static inline void io_insb(int port, void* addr, int cnt) {
+    asm volatile("cld\n\trepne\n\tinsb"
+                 : "=D" (addr), "=c" (cnt)
+                 : "d" (port), "0" (addr), "1" (cnt)
+                 : "memory", "cc");
+}
+
+static inline uint16_t io_inw(int port) {
+    uint16_t data;
+    asm volatile("inw %w1,%0" : "=a" (data) : "d" (port));
+    return data;
+}
+
+static inline void io_insw(int port, void* addr, int cnt) {
+    asm volatile("cld\n\trepne\n\tinsw"
+                 : "=D" (addr), "=c" (cnt)
+                 : "d" (port), "0" (addr), "1" (cnt)
+                 : "memory", "cc");
+}
+
+static inline uint32_t io_inl(int port) {
+    uint32_t data;
+    asm volatile("inl %w1,%0" : "=a" (data) : "d" (port));
+    return data;
+}
+
+static inline void io_insl(int port, void* addr, int cnt) {
+    asm volatile("cld\n\trepne\n\tinsl"
+                 : "=D" (addr), "=c" (cnt)
+                 : "d" (port), "0" (addr), "1" (cnt)
+                 : "memory", "cc");
+}
+
+static inline void io_outb(int port, uint8_t data) {
+    asm volatile("outb %0,%w1" : : "a" (data), "d" (port));
+}
+
+static inline void io_outsb(int port, const void* addr, int cnt) {
+    asm volatile("cld\n\trepne\n\toutsb"
+                 : "=S" (addr), "=c" (cnt)
+                 : "d" (port), "0" (addr), "1" (cnt)
+                 : "cc");
+}
+
+static inline void io_outw(int port, uint16_t data) {
+    asm volatile("outw %0,%w1" : : "a" (data), "d" (port));
+}
+
+static inline void io_outsw(int port, const void* addr, int cnt) {
+    asm volatile("cld\n\trepne\n\toutsw"
+                 : "=S" (addr), "=c" (cnt)
+                 : "d" (port), "0" (addr), "1" (cnt)
+                 : "cc");
+}
+
+static inline void io_outsl(int port, const void* addr, int cnt) {
+    asm volatile("cld\n\trepne\n\toutsl"
+                 : "=S" (addr), "=c" (cnt)
+                 : "d" (port), "0" (addr), "1" (cnt)
+                 : "cc");
+}
+
+static inline void io_outl(int port, uint32_t data) {
+    asm volatile("outl %0,%w1" : : "a" (data), "d" (port));
 }
 
 #endif /* __LUNAIX_IO_H */
diff --git a/lunaix-os/includes/hal/pic.h b/lunaix-os/includes/hal/pic.h
new file mode 100644
index 0000000..20ee67b
--- /dev/null
+++ b/lunaix-os/includes/hal/pic.h
@@ -0,0 +1,5 @@
+#ifndef __LUNAIX_PIC_H
+#define __LUNAIX_PIC_H
+// TODO: PIC
+
+#endif /* __LUNAIX_PIC_H */
diff --git a/lunaix-os/includes/lunaix/mm/dmm.h b/lunaix-os/includes/lunaix/mm/dmm.h
new file mode 100644
index 0000000..6ea878b
--- /dev/null
+++ b/lunaix-os/includes/lunaix/mm/dmm.h
@@ -0,0 +1,16 @@
+#ifndef __LUNAIX_DMM_H
+#define __LUNAIX_DMM_H
+// Dynamic Memory (i.e., heap) Manager
+
+#include <stddef.h>
+
+void
+lxsbrk(void* current, void* next);
+
+void
+lxmalloc(size_t size);
+
+void
+lxfree(size_t size);
+
+#endif /* __LUNAIX_DMM_H */
diff --git a/lunaix-os/includes/lunaix/mm/page.h b/lunaix-os/includes/lunaix/mm/page.h
index 2923258..e97de1d 100644
--- a/lunaix-os/includes/lunaix/mm/page.h
+++ b/lunaix-os/includes/lunaix/mm/page.h
@@ -3,6 +3,9 @@
 #include <stdint.h>
 #include <lunaix/constants.h>
 
+#define PG_SIZE_BITS              12
+#define PG_INDEX_BITS               10
+
 #define PG_MAX_ENTRIES              1024U
 #define PG_LAST_TABLE               PG_MAX_ENTRIES - 1
 #define PG_FIRST_TABLE              0
@@ -12,9 +15,9 @@
 
 #define PG_ALIGN(addr)      ((uintptr_t)(addr)   & 0xFFFFF000UL)
 
-#define PD_INDEX(vaddr)     (((uintptr_t)(vaddr) & 0xFFC00000UL) >> 22)
-#define PT_INDEX(vaddr)     (((uintptr_t)(vaddr) & 0x003FF000UL) >> 12)
-#define PG_OFFSET(vaddr)    ((uintptr_t)(vaddr)  & 0x00000FFFUL)
+#define L1_INDEX(vaddr)     (uint32_t)(((uintptr_t)(vaddr) & 0xFFC00000UL) >> 22)
+#define L2_INDEX(vaddr)     (uint32_t)(((uintptr_t)(vaddr) & 0x003FF000UL) >> 12)
+#define PG_OFFSET(vaddr)    (uint32_t)((uintptr_t)(vaddr)  & 0x00000FFFUL)
 
 #define GET_PT_ADDR(pde)    PG_ALIGN(pde)
 #define GET_PG_ADDR(pte)    PG_ALIGN(pte)
@@ -31,32 +34,49 @@
 #define PG_DISABLE_CACHE        (1 << 4)
 #define PG_PDE_4MB              (1 << 7)
 
-#define PDE(flags, pt_addr)     (PG_ALIGN(pt_addr) | ((flags) & 0xfff))
-#define PTE(flags, pg_addr)     (PG_ALIGN(pg_addr) | ((flags) & 0xfff))
+#define NEW_L1_ENTRY(flags, pt_addr)     (PG_ALIGN(pt_addr) | ((flags) & 0xfff))
+#define NEW_L2_ENTRY(flags, pg_addr)     (PG_ALIGN(pg_addr) | ((flags) & 0xfff))
 
 #define V_ADDR(pd, pt, offset)  ((pd) << 22 | (pt) << 12 | (offset))
 #define P_ADDR(ppn, offset)     ((ppn << 12) | (offset))
 
-// 用于对PD进行循环映射，因为我们可能需要对PD进行频繁操作，我们在这里禁用TLB缓存
+#define PG_ENTRY_FLAGS(entry)   (entry & 0xFFFU)
+#define PG_ENTRY_ADDR(entry)   (entry & ~0xFFFU)
+
 #define PG_PREM_R              PG_PRESENT
 #define PG_PREM_RW             PG_PRESENT | PG_WRITE
 #define PG_PREM_UR             PG_PRESENT | PG_ALLOW_USER
 #define PG_PREM_URW            PG_PRESENT | PG_WRITE | PG_ALLOW_USER
+
+// 用于对PD进行循环映射，因为我们可能需要对PD进行频繁操作，我们在这里禁用TLB缓存
 #define T_SELF_REF_PERM        PG_PREM_RW | PG_DISABLE_CACHE
 
 
 // 页目录的虚拟基地址，可以用来访问到各个PDE
-#define PTD_BASE_VADDR                0xFFFFF000U
+#define L1_BASE_VADDR                0xFFFFF000U
 
 // 页表的虚拟基地址，可以用来访问到各个PTE
-#define PT_BASE_VADDR                 0xFFC00000U
+#define L2_BASE_VADDR                 0xFFC00000U
 
 // 用来获取特定的页表的虚拟地址
-#define PT_VADDR(pd_offset)           (PT_BASE_VADDR | (pd_offset << 12))
+#define L2_VADDR(pd_offset)           (L2_BASE_VADDR | (pd_offset << 12))
 
 typedef unsigned long ptd_t;
 typedef unsigned long pt_t;
 typedef unsigned int pt_attr;
 
+/**
+ * @brief 虚拟映射属性
+ * 
+ */
+typedef struct {
+    // 物理页码（如果不存在映射，则为0）
+    uint32_t pn;
+    // 物理页地址（如果不存在映射，则为0）
+    uintptr_t pa;
+    // 映射的flags
+    uint16_t flags;
+} v_mapping;
+
 
 #endif /* __LUNAIX_PAGE_H */
diff --git a/lunaix-os/includes/lunaix/mm/vmm.h b/lunaix-os/includes/lunaix/mm/vmm.h
index 14b3a7f..950928c 100644
--- a/lunaix-os/includes/lunaix/mm/vmm.h
+++ b/lunaix-os/includes/lunaix/mm/vmm.h
@@ -1,60 +1,84 @@
 #ifndef __LUNAIX_VMM_H
 #define __LUNAIX_VMM_H
-#include <stdint.h>
-#include <stddef.h>
 #include <lunaix/mm/page.h>
+#include <stddef.h>
+#include <stdint.h>
 // Virtual memory manager
 
 /**
  * @brief 初始化虚拟内存管理器
- * 
+ *
  */
-void vmm_init();
+void
+vmm_init();
 
 /**
  * @brief 创建一个页目录
- * 
+ *
  * @return ptd_entry* 页目录的物理地址，随时可以加载进CR3
  */
-ptd_t* vmm_init_pd();
+ptd_t*
+vmm_init_pd();
 
 /**
- * @brief 
- * 尝试建立一个映射关系。映射指定的物理页地址至虚拟页地址，如果指定的虚拟页地址已被占用
+ * @brief 尝试建立一个映射关系。映射指定的物理页地址至虚拟页地址，如果指定的虚拟页地址已被占用
  * 则尝试寻找新的可用地址（改地址总是大于指定的地址）。
- * 
+ *
  * @param vpn 虚拟页地址
- * @param ppn 物理页地址
+ * @param pa 物理页地址
  * @param dattr PDE 的属性
  * @param tattr PTE 的属性
  * @return 虚拟页地址，如不成功，则为 NULL
  */
-void* vmm_map_page(void* vpn, void* ppn, pt_attr dattr, pt_attr tattr);
+void*
+vmm_map_page(void* va, void* pa, pt_attr tattr);
+
+/**
+ * @brief 建立一个映射关系，映射指定的物理页地址至虚拟页地址。如果指定的虚拟页地址已被占用，
+ * 则覆盖。
+ *
+ * @param va 虚拟页地址
+ * @param pa 物理页地址
+ * @param dattr PDE 的属性
+ * @param tattr PTE 的属性
+ * @return 虚拟页地址
+ */
+void*
+vmm_fmap_page(void* va, void* pa, pt_attr tattr);
 
 /**
  * @brief 尝试为一个虚拟页地址创建一个可用的物理页映射
- * 
- * @param vpn 虚拟页地址
+ *
+ * @param va 虚拟页地址
  * @return 物理页地址，如不成功，则为 NULL
  */
-void* vmm_alloc_page(void* vpn, pt_attr dattr, pt_attr tattr);
+void*
+vmm_alloc_page(void* va, pt_attr tattr);
 
 /**
  * @brief 删除一个映射
- * 
- * @param vpn 
+ *
+ * @param vpn
  */
-void vmm_unmap_page(void* vpn);
-
-ptd_t* get_pd();
-void set_pd(ptd_t* pd);
+void
+vmm_unmap_page(void* va);
 
 /**
  * @brief 将虚拟地址翻译为其对应的物理映射
- * 
+ *
  * @param va 虚拟地址
  * @return void* 物理地址，如映射不存在，则为NULL
  */
-void* vmm_v2p(void* va);
+void*
+vmm_v2p(void* va);
+
+/**
+ * @brief 查找一个映射
+ *
+ * @param va 虚拟地址
+ * @return v_mapping 映射相关属性
+ */
+v_mapping
+vmm_lookup(void* va);
 
 #endif /* __LUNAIX_VMM_H */
diff --git a/lunaix-os/kernel/k_init.c b/lunaix-os/kernel/k_init.c
index 642a421..07091ca 100644
--- a/lunaix-os/kernel/k_init.c
+++ b/lunaix-os/kernel/k_init.c
@@ -87,7 +87,7 @@ setup_memory(multiboot_memory_map_t* map, size_t map_size) {
         vmm_map_page(
             (void*)(VGA_BUFFER_VADDR + (i << PG_SIZE_BITS)), 
             (void*)(VGA_BUFFER_PADDR + (i << PG_SIZE_BITS)), 
-            PG_PREM_RW, PG_PREM_RW
+            PG_PREM_RW
         );
     }
     
@@ -101,7 +101,7 @@ void
 setup_kernel_runtime() {
     // 为内核创建一个专属栈空间。
     for (size_t i = 0; i < (K_STACK_SIZE >> PG_SIZE_BITS); i++) {
-        vmm_alloc_page((void*)(K_STACK_START + (i << PG_SIZE_BITS)), PG_PREM_RW, PG_PREM_RW);
+        vmm_alloc_page((void*)(K_STACK_START + (i << PG_SIZE_BITS)), PG_PREM_RW);
     }
     printf("[MM] Allocated %d pages for stack start at %p\n", K_STACK_SIZE>>PG_SIZE_BITS, K_STACK_START);
 }
\ No newline at end of file
diff --git a/lunaix-os/kernel/mm/dmm.c b/lunaix-os/kernel/mm/dmm.c
new file mode 100644
index 0000000..04a0249
--- /dev/null
+++ b/lunaix-os/kernel/mm/dmm.c
@@ -0,0 +1,19 @@
+#include <lunaix/mm/dmm.h>
+#include <lunaix/mm/vmm.h>
+
+// This is a temporary design. 
+//  We can do better when we are ready for multitasking
+void
+lxsbrk(void* current, void* next) {
+    // TODO: sbrk 
+}
+
+void
+lxmalloc(size_t size) {
+    // TODO: Malloc 
+}
+
+void
+lxfree(size_t size) {
+    // TODO: Free 
+}
\ No newline at end of file
diff --git a/lunaix-os/kernel/mm/vmm.c b/lunaix-os/kernel/mm/vmm.c
index a076a97..2f8d726 100644
--- a/lunaix-os/kernel/mm/vmm.c
+++ b/lunaix-os/kernel/mm/vmm.c
@@ -2,32 +2,8 @@
 #include <lunaix/mm/page.h>
 #include <lunaix/mm/pmm.h>
 #include <lunaix/mm/vmm.h>
-
-// TODO: Move these nasty inline asm stuff into hal
-//      These should be arch dependent
-ptd_t*
-get_pd()
-{
-    ptd_t* pd;
-#ifdef __ARCH_IA32
-    __asm__("movl %%cr3, %0\n"
-            "andl $0xfffff000, %0"
-            : "=r"(pd));
-#endif
-    return (ptd_t*)P2V(pd);
-}
-
-void
-set_pd(ptd_t* pd)
-{
-#ifdef __ARCH_IA32
-    __asm__("movl %0, %%eax\n"
-            "andl $0xfffff000, %%eax\n"
-            "movl %%eax, %%cr3\n"
-            :
-            : "r"(pd));
-#endif
-}
+#include <lunaix/assert.h>
+#include <hal/cpu.h>
 
 void
 vmm_init()
@@ -39,72 +15,112 @@ ptd_t*
 vmm_init_pd()
 {
     ptd_t* dir = pmm_alloc_page();
-    for (size_t i = 0; i < 1024; i++) {
+    for (size_t i = 0; i < PG_MAX_ENTRIES; i++) {
         dir[i] = 0;
     }
 
-    // 自己映射自己，方便我们在软件层面进行查表地址转换
-    dir[1023] = PDE(T_SELF_REF_PERM, dir);
+    // 递归映射，方便我们在软件层面进行查表地址转换
+    dir[PG_MAX_ENTRIES - 1] = NEW_L1_ENTRY(T_SELF_REF_PERM, dir);
 
     return dir;
 }
 
+int
+__vmm_map_internal(uint32_t l1_inx, uint32_t l2_inx, uintptr_t pa, pt_attr attr) {
+    ptd_t* l1pt = (ptd_t*)L1_BASE_VADDR;
+    pt_t* l2pt = (pt_t*)L2_VADDR(l1_inx);
+
+    // See if attr make sense
+    assert(attr <= 128);
+
+    if(!l1pt[l1_inx]) {
+        uint8_t* new_l1pt_pa = pmm_alloc_page();
+
+        // 物理内存已满！
+        if (!new_l1pt_pa) {
+            return 0;
+        }
+
+        l1pt[l1_inx] = NEW_L1_ENTRY(attr, new_l1pt_pa);
+        memset((void*)L2_VADDR(l1_inx), 0, PM_PAGE_SIZE);
+    }
+
+    l2pt[l2_inx] = NEW_L2_ENTRY(attr, pa);
+
+    return 1;
+}
+
 void*
-vmm_map_page(void* va, void* pa, pt_attr dattr, pt_attr tattr)
+vmm_map_page(void* va, void* pa, pt_attr tattr)
 {
     // 显然，对空指针进行映射没有意义。
     if (!pa || !va) {
         return NULL;
     }
 
-    uintptr_t pd_offset = PD_INDEX(va);
-    uintptr_t pt_offset = PT_INDEX(va);
-    ptd_t* ptd = (ptd_t*)PTD_BASE_VADDR;
+    assert(((uintptr_t)va & 0xFFFU) == 0)
+    assert(((uintptr_t)pa & 0xFFFU) == 0)
+
+    uint32_t l1_index = L1_INDEX(va);
+    uint32_t l2_index = L2_INDEX(va);
+    ptd_t* l1pt = (ptd_t*)L1_BASE_VADDR;
 
     // 在页表与页目录中找到一个可用的空位进行映射（位于va或其附近）
-    ptd_t pde = ptd[pd_offset];
-    pt_t* pt = (pt_t*)PT_VADDR(pd_offset);
-    while (pde && pd_offset < 1024) {
-        if (pt_offset == 1024) {
-            pd_offset++;
-            pt_offset = 0;
-            pde = ptd[pd_offset];
-            pt = (pt_t*)PT_VADDR(pd_offset);
+    ptd_t l1pte = l1pt[l1_index];
+    pt_t* l2pt = (pt_t*)L2_VADDR(l1_index);
+    while (l1pte && l1_index < PG_MAX_ENTRIES) {
+        if (l2_index == PG_MAX_ENTRIES) {
+            l1_index++;
+            l2_index = 0;
+            l1pte = l1pt[l1_index];
+            l2pt = (pt_t*)L2_VADDR(l1_index);
         }
-        // 页表有空位，只需要开辟一个新的 PTE
-        if (pt && !pt[pt_offset]) {
-            pt[pt_offset] = PTE(tattr, pa);
-            return (void*)V_ADDR(pd_offset, pt_offset, PG_OFFSET(va));
+        // 页表有空位，只需要开辟一个新的 PTE (Level 2)
+        if (l2pt && !l2pt[l2_index]) {
+            l2pt[l2_index] = NEW_L2_ENTRY(tattr, pa);
+            return (void*)V_ADDR(l1_index, l2_index, PG_OFFSET(va));
         }
-        pt_offset++;
+        l2_index++;
     }
 
     // 页目录与所有页表已满！
-    if (pd_offset > 1024) {
+    if (l1_index > PG_MAX_ENTRIES) {
         return NULL;
     }
 
-    // 页目录有空位，需要开辟一个新的 PDE
-    uint8_t* new_pt_pa = pmm_alloc_page();
+    if (!__vmm_map_internal(l1_index, l2_index, pa, tattr)) {
+        return NULL;
+    }
+
+    return (void*)V_ADDR(l1_index, l2_index, PG_OFFSET(va));
+}
 
-    // 物理内存已满！
-    if (!new_pt_pa) {
+void*
+vmm_fmap_page(void* va, void* pa, pt_attr tattr) {
+    if (!pa || !va) {
         return NULL;
     }
 
-    ptd[pd_offset] = PDE(dattr, new_pt_pa);
+    assert(((uintptr_t)va & 0xFFFU) == 0)
+    assert(((uintptr_t)pa & 0xFFFU) == 0)
+
+    uint32_t l1_index = L1_INDEX(va);
+    uint32_t l2_index = L2_INDEX(va);
+
+    if (!__vmm_map_internal(l1_index, l2_index, pa, tattr)) {
+        return NULL;
+    }
 
-    memset((void*)PT_VADDR(pd_offset), 0, PM_PAGE_SIZE);
-    pt[pt_offset] = PTE(tattr, pa);
+    cpu_invplg(va);
 
-    return (void*)V_ADDR(pd_offset, pt_offset, PG_OFFSET(va));
+    return (void*)V_ADDR(l1_index, l2_index, PG_OFFSET(va));
 }
 
 void*
-vmm_alloc_page(void* vpn, pt_attr dattr, pt_attr tattr)
+vmm_alloc_page(void* vpn, pt_attr tattr)
 {
     void* pp = pmm_alloc_page();
-    void* result = vmm_map_page(vpn, pp, dattr, tattr);
+    void* result = vmm_map_page(vpn, pp, tattr);
     if (!result) {
         pmm_free_page(pp);
     }
@@ -112,43 +128,53 @@ vmm_alloc_page(void* vpn, pt_attr dattr, pt_attr tattr)
 }
 
 void
-vmm_unmap_page(void* vpn)
+vmm_unmap_page(void* va)
 {
-    uintptr_t pd_offset = PD_INDEX(vpn);
-    uintptr_t pt_offset = PT_INDEX(vpn);
-    ptd_t* self_pde = (ptd_t*)PTD_BASE_VADDR;
-
-    ptd_t pde = self_pde[pd_offset];
-
-    if (pde) {
-        pt_t* pt = (pt_t*)PT_VADDR(pd_offset);
-        uint32_t pte = pt[pt_offset];
-        if (IS_CACHED(pte) && pmm_free_page((void*)pte)) {
-// 刷新TLB
-#ifdef __ARCH_IA32
-            __asm__("invlpg (%0)" ::"r"((uintptr_t)vpn) : "memory");
-#endif
+    assert(((uintptr_t)va & 0xFFFU) == 0)
+
+    uint32_t l1_index = L1_INDEX(va);
+    uint32_t l2_index = L2_INDEX(va);
+    ptd_t* l1pt = (ptd_t*)L1_BASE_VADDR;
+
+    ptd_t l1pte = l1pt[l1_index];
+
+    if (l1pte) {
+        pt_t* l2pt = (pt_t*)L2_VADDR(l1_index);
+        uint32_t l2pte = l2pt[l2_index];
+        if (IS_CACHED(l2pte) && pmm_free_page((void*)l2pte)) {
+            cpu_invplg(va);
         }
-        pt[pt_offset] = 0;
+        l2pt[l2_index] = 0;
     }
 }
 
-void*
-vmm_v2p(void* va)
+v_mapping
+vmm_lookup(void* va)
 {
-    uintptr_t pd_offset = PD_INDEX(va);
-    uintptr_t pt_offset = PT_INDEX(va);
-    uintptr_t po = PG_OFFSET(va);
-    ptd_t* self_pde = (ptd_t*)PTD_BASE_VADDR;
-
-    ptd_t pde = self_pde[pd_offset];
-    if (pde) {
-        pt_t pte = ((pt_t*)PT_VADDR(pd_offset))[pt_offset];
-        if (pte) {
-            uintptr_t ppn = pte >> 12;
-            return (void*)P_ADDR(ppn, po);
+    assert(((uintptr_t)va & 0xFFFU) == 0)
+
+    uint32_t l1_index = L1_INDEX(va);
+    uint32_t l2_index = L2_INDEX(va);
+    uint32_t po = PG_OFFSET(va);
+    ptd_t* l1pt = (ptd_t*)L1_BASE_VADDR;
+
+    ptd_t l1pte = l1pt[l1_index];
+
+    v_mapping mapping = { .flags = 0, .pa = 0, .pn = 0 };
+    if (l1pte) {
+        pt_t l2pte = ((pt_t*)L2_VADDR(l1_index))[l2_index];
+        if (l2pte) {
+            uintptr_t ppn = l2pte >> PG_SIZE_BITS;
+            mapping.flags = PG_ENTRY_FLAGS(l2pte);
+            mapping.pa = PG_ENTRY_ADDR(l2pte);
+            mapping.pn = mapping.pa >> PG_SIZE_BITS;
         }
     }
 
-    return NULL;
+    return mapping;
+}
+
+void*
+vmm_v2p(void* va) {
+    return vmm_lookup(va).pa;
 }
\ No newline at end of file
-- 
2.27.0