dev:Refactor Memory Management and Process Management

3 files changed, 795 insertions, 198 deletions

diff --git a/src/kernel/mm/pmm.cc b/src/kernel/mm/pmm.cc
index 9b68fcd..09bd65b 100644
--- a/src/kernel/mm/pmm.cc
+++ b/src/kernel/mm/pmm.cc
@@ -1,71 +1,356 @@
+#include "common/types.h"
+#include "kernel/riscv.h"
 #include "kernel/pmm.h"
+#include "kernel/mm.h"
+// #include "kernel/spinlock.h"
 #include "drivers/uart.h"
+#include "lib/string.h"
 
 extern "C" char kernel_end[];
+extern uint64 g_dtb_addr;
+
+// namespace PMM
+// {
+//     struct Run
+//     {
+//         struct Run *next;
+//     };
+
+//     struct
+//     {
+//         struct Run *freelist;
+//     } kmem;
+
+//     inline uint64_t PGROUNDUP(uint64_t a)
+//     {
+//         return (a + PGSIZE - 1) & ~(PGSIZE - 1);
+//     }
+
+//     void init()
+//     {
+//         uint64_t start = PGROUNDUP((uint64_t)kernel_end);
+
+//         Drivers::uart_puts("PMM: Initializing from");
+
+//         for (uint64_t p = start; p + PGSIZE <= PHYSTOP; p += PGSIZE)
+//         {
+//             free_page((void *)p);
+//         }
+//     }
+
+//     void free_page(void* pa)
+//     {
+//         struct Run *r;
+//         if (((uint64_t)pa % PGSIZE) != 0 || (uint64_t)pa < (uint64_t)kernel_end || (uint64_t)pa >= PHYSTOP)
+//         {
+//             Drivers::uart_puts("PMM: Error - Invalid free address\n");
+//             return;
+//         }
+//         r = (struct Run*)pa;
+
+//         r->next = kmem.freelist;
+//         kmem.freelist = r;
+        
+//     }
+
+//     void* alloc_page()
+//     {
+//         struct Run* r;
+
+//         r = kmem.freelist;
+//         if (r)
+//         {
+//             kmem.freelist = r->next;
+
+//             uint64_t *p = (uint64_t *)r;
+//             for (int i = 0; i < PGSIZE / 8; i++)
+//             {
+//                 p[i] = 0;
+//             }
+//         }
+//         else
+//         {
+//             Drivers::uart_puts("PMM: Out of memory!\n");
+//         }
+//         return (void *)r;
+//     }
+// }
 
 namespace PMM
 {
-    struct Run
-    {
-        struct Run *next;
-    };
+    // Global metadata array pointer
+    struct Page *mem_map = nullptr;
+    // Total number of physical memory pages
+    uint64 total_pages = 0;
+    // Free list array (free_area[k] stores the head of the free block list of order=k)
+    struct Page *free_area[MAX_ORDER];
+    // Protective lock
+    Spinlock pmm_lock;
 
-    struct
+    struct Page *pa_to_page(uint64 pa)
     {
-        struct Run *freelist;
-    } kmem;
+        if (pa < KERNBASE || pa >= PHYSTOP)
+            return nullptr;
+        uint64 pfn = (pa - KERNBASE) >> PGSHIFT;
+        if (pfn >= total_pages)
+            return nullptr;
+        return &mem_map[pfn];
+    }
 
-    inline uint64_t PGROUNDUP(uint64_t a)
+    uint64 page_to_pa(struct Page *page)
     {
-        return (a + PGSIZE - 1) & ~(PGSIZE - 1);
+        uint64 pfn = page - mem_map;
+        return KERNBASE + (pfn << PGSHIFT);
     }
 
-    void init()
+    void list_add(struct Page *page, int order)
     {
-        uint64_t start = PGROUNDUP((uint64_t)kernel_end);
+        page->order = order;
+        page->flags = 0; // Mark as available
+        page->next = free_area[order];
+        page->prev = nullptr;
 
-        Drivers::uart_puts("PMM: Initializing from");
+        if (free_area[order])
+        {
+            free_area[order]->prev = page;
+        }
+        free_area[order] = page;
+    }
 
-        for (uint64_t p = start; p + PGSIZE <= PHYSTOP; p += PGSIZE)
+    void list_del(struct Page* page, int order)
+    {
+        if (page->prev)
+        {
+            page->prev->next = page->next;
+        }
+        else
+        {
+            free_area[order] = page->next;
+        }
+        if (page->next)
         {
-            free_page((void *)p);
+            page->next->prev = page->prev;
         }
+        page->next = nullptr;
+        page->prev = nullptr;
     }
 
-    void free_page(void* pa)
+    void init()
     {
-        struct Run *r;
-        if (((uint64_t)pa % PGSIZE) != 0 || (uint64_t)pa < (uint64_t)kernel_end || (uint64_t)pa >= PHYSTOP)
+        pmm_lock.init("pmm");
+
+        uint64 kern_end_pa = PGROUNDUP((uint64)kernel_end);
+        total_pages = (PHYSTOP - KERNBASE) / PGSIZE;
+
+        Drivers::uart_puts("PMM: Total RAM size: 128MB\n");
+        // Place the mem_map array
+        // mem_map is located immediately after kernel_end
+        mem_map = (struct Page*)kern_end_pa;
+
+        uint64 mem_map_size = total_pages * sizeof(struct Page);
+        uint64 free_mem_start = PGROUNDUP(kern_end_pa + mem_map_size);
+
+        Drivers::uart_puts("PMM: MemMap placed at: ");
+        Drivers::print_hex((uint64)mem_map);
+        Drivers::uart_putc('\n');
+
+        Drivers::uart_puts("PMM: Free memory starts at: ");
+        Drivers::print_hex(free_mem_start);
+        Drivers::uart_puts("\n");
+
+        if (free_mem_start >= PHYSTOP)
         {
-            Drivers::uart_puts("PMM: Error - Invalid free address\n");
-            return;
+            Drivers::uart_puts("PMM: Error - Kernel too large, no memory left!\n");
+            while (1)
+                ;
+        }
+
+        // Initialize all page states
+        for (uint64 i = 0; i < total_pages; i++)
+        {
+            mem_map[i].flags = PG_reserved;
+            mem_map[i].order = 0;
+            mem_map[i].refcnt = 0;
+            mem_map[i].next = nullptr;
+            mem_map[i].prev = nullptr;
         }
-        r = (struct Run*)pa;
 
-        r->next = kmem.freelist;
-        kmem.freelist = r;
+        for (int i = 0; i < MAX_ORDER; i++)
+        {
+            free_area[i] = nullptr;
+        }
+
+        // Release actual free memory to the allocator
+        // Range: [free_mem_start, PHYSTOP)
+        // Note: need to avoid the page where the DTB is located
+        uint64 dtb_pa_start = PGROUNDDOWN(g_dtb_addr);
+        uint64 dtb_pa_end = PGROUNDUP(g_dtb_addr + 0x10000);
+
+        uint64 p = free_mem_start;
+        while (p < PHYSTOP)
+        {
+            if (p >= dtb_pa_start && p < dtb_pa_end)
+            {
+                p += PGSIZE;
+                continue;
+            }
+
+            // Mark the page as Used, then call free_pages to release it
+            // This allows free_pages to handle the merging logic automatically
+            struct Page *page = pa_to_page(p);
+            if (page)
+            {
+                page->flags = PG_used;
+                free_pages((void *)p, 0); // Release page by page, let Buddy merge automatically
+            }
+            p += PGSIZE;
+
+            
+        }
         
     }
 
-    void* alloc_page()
+    void *alloc_pages(int order)
     {
-        struct Run* r;
+        if (order < 0 || order >= MAX_ORDER)
+            return nullptr;
+        
+        pmm_lock.acquire();
 
-        r = kmem.freelist;
-        if (r)
+        int current_order = order;
+        while (current_order < MAX_ORDER)
         {
-            kmem.freelist = r->next;
+            if (free_area[current_order])
+            {
+                // Find free block
+                struct Page *page = free_area[current_order];
+                list_del(page, current_order);
+
+                // If the block is too large, split it.
+                while (current_order > order)
+                {
+                    current_order--;
+                    // page is the left part (Buddy A)
+                    // buddy is the right part (Buddy B)
+                    // Pointer addition: Page* + n is equivalent to address offset n * sizeof(Page)
+                    // Since mem_map is linearly mapped, mem_map[i + (1<<k)] exactly corresponds to physical address + (1<<k)*PGSIZE
+                    struct Page *buddy = page + (1 << current_order);
+
+                    // Add the buddy to a lower-level linked list
+                    list_add(buddy, current_order);
+                }
+
+                // Mark as Occupied
+                page->flags = PG_used;
+                page->order = order; // Record the orders that have been assigned
+                page->refcnt = 1;
 
-            uint64_t *p = (uint64_t *)r;
-            for (int i = 0; i < PGSIZE / 8; i++)
+                pmm_lock.release();
+                return (void *)page_to_pa(page);
+            }
+            current_order++;
+        }
+        pmm_lock.release();
+        Drivers::uart_puts("PMM: OOM alloc_pages\n");
+        return nullptr;
+    }
+
+    void free_pages(void *pa, int order)
+    {
+        if (!pa)
+            return;
+        struct Page *page = pa_to_page((uint64)pa);
+        if (!page)
+            return;
+        pmm_lock.acquire();
+
+        if (!(page->flags & PG_used))
+        {
+            Drivers::uart_puts("PMM: Double free warning\n");
+            pmm_lock.release();
+            return;
+        }
+
+        // Mark as basic status
+        page->flags = 0;
+        page->refcnt = 0;
+
+        // try merge
+        uint64 page_idx = page - mem_map;
+
+        while (order < MAX_ORDER - 1)
+        {
+            // calculate Buddy index: index ^ (1 << order)
+            uint64 buddy_idx = page_idx ^ (1 << order);
+            struct Page *buddy = &mem_map[buddy_idx];
+
+            // Check if the Buddy can be merged:
+            // 1. Must be within the memory range
+            // 2. Must be free (flags == 0)
+            // 3. Must be at the current order (to prevent merging already split small blocks)
+            if (buddy_idx >= total_pages || (buddy->flags & PG_used) || (buddy->flags & PG_reserved) || buddy->order != order)
             {
-                p[i] = 0;
+                break;
             }
+
+            // Can be merged!
+            // Remove Buddy from the free list
+            list_del(buddy, order);
+
+            // Update the current block to the merged block
+            // The starting address of the merged block is the smaller one between index and buddy_index
+            if (buddy_idx < page_idx)
+            {
+                page = buddy;
+                page_idx = buddy_idx;
+            }
+
+            order++;
         }
-        else
+
+        // Insert into the final order linked list
+        list_add(page, order);
+
+        pmm_lock.release();
+    }
+
+    void *alloc_page()
+    {
+        void *p = alloc_pages(0);
+        if (p)
+        {
+            memset(p, 0, PGSIZE);
+        }
+        return p;
+    }
+
+    void free_page(void *pa)
+    {
+        free_pages(pa, 0);
+    }
+
+    void print_free_lists()
+    {
+        // Debug print the number of free blocks at each level
+        Drivers::uart_puts("PMM Stats: [ ");
+        for (int i = 0; i < MAX_ORDER; i++)
         {
-            Drivers::uart_puts("PMM: Out of memory!\n");
+            int count = 0;
+            struct Page *p = free_area[i];
+            while (p)
+            {
+                count++;
+                p = p->next;
+            }
+            if (count > 0)
+            {
+                Drivers::uart_puts("Order");
+                // Drivers::uart_put_int(i);
+                Drivers::uart_puts(":");
+                // Drivers::uart_put_int(count);
+                Drivers::uart_puts(" ");
+            }
         }
-        return (void *)r;
+        Drivers::uart_puts("]\n");
     }
-}
\ No newline at end of file
+} // namespace PMM
diff --git a/src/kernel/mm/slab.cc b/src/kernel/mm/slab.cc
new file mode 100644
index 0000000..370074a
--- /dev/null
+++ b/src/kernel/mm/slab.cc
@@ -0,0 +1,246 @@
+#include "kernel/slab.h"
+#include "kernel/pmm.h"
+#include "kernel/spinlock.h"
+#include "kernel/riscv.h"
+#include "drivers/uart.h"
+#include "lib/string.h"
+
+#define PG_slab 0x04
+
+// Supported object size (must be a multiple of 8 and <= 2048)
+const uint64 kmalloc_sizes[] = {
+    16, 32, 64, 128, 256, 512, 1024, 2048};
+const int NUM_CACHES = 8;
+
+struct SlabHeader
+{
+    struct SlabHeader *next; // Next slab in the linked list
+    uint16 free_idx;         // Index of the head of the free list (0xFFFF means none)
+    uint16 used_count;       // Number of allocated objects
+    // Followed immediately by object data...
+    // The free list of objects is embedded in the first 2 bytes of the object itself (stores the index of the next object)
+};
+
+struct KmemCache
+{
+    Spinlock lock;
+    uint64 obj_size;
+    uint64 obj_count_per_slab;
+    struct SlabHeader *slab_list; // Slab list with free objects
+    const char *name;
+};
+
+struct KmemCache caches[NUM_CACHES];
+
+// Virtual address -> struct Page*
+static struct Page *virt_to_page(void *ptr)
+{
+    uint64 pa = (uint64)ptr;
+    return PMM::pa_to_page(pa);
+}
+
+// Slab Implementation
+
+namespace Slab
+{
+
+    void init()
+    {
+        for (int i = 0; i < NUM_CACHES; i++)
+        {
+            caches[i].lock.init("slab_lock");
+            caches[i].obj_size = kmalloc_sizes[i];
+            caches[i].name = "kmalloc";
+
+            // Calculate how many objects can fit on a page
+            // Excluding the overhead of the SlabHeader
+            uint64 available = PGSIZE - sizeof(struct SlabHeader);
+            caches[i].obj_count_per_slab = available / caches[i].obj_size;
+            caches[i].slab_list = nullptr;
+        }
+        Drivers::uart_puts("[Slab] Initialized kmalloc caches.\n");
+    }
+
+    // Allocate a new slab page for the cache
+    // Must hold cache->lock
+    static struct SlabHeader *grow_cache(struct KmemCache *c)
+    {
+        void *p = PMM::alloc_page();
+        if (!p)
+            return nullptr;
+
+        // Mark this page as a Slab page
+        struct Page *page = virt_to_page(p);
+        if (page)
+        {
+            page->flags |= PG_slab;
+            int cache_idx = c - caches;
+            page->order = (uint8)cache_idx;
+        }
+
+        struct SlabHeader *slab = (struct SlabHeader *)p;
+        slab->next = nullptr;
+        slab->used_count = 0;
+        slab->free_idx = 0;
+
+        // Initialize the free list of objects
+        // The starting position of each object stores the index of the next object
+        // Starting position of the object data area
+        char *base = (char *)(slab + 1);
+
+        for (uint i = 0; i < c->obj_count_per_slab - 1; i++)
+        {
+            // The content of the i-th object is i+1
+            uint16 *ptr = (uint16 *)(base + i * c->obj_size);
+            *ptr = i + 1;
+        }
+        // The last object points to 0xFFFF
+        uint16 *ptr = (uint16 *)(base + (c->obj_count_per_slab - 1) * c->obj_size);
+        *ptr = 0xFFFF;
+
+        return slab;
+    }
+
+    void *kmalloc(uint64 size)
+    {
+        if (size == 0)
+            return nullptr;
+
+        // Large objects go directly through PMM
+        if (size > 2048)
+        {
+            // calculate order
+            int order = 0;
+            uint64 req = size;
+            uint64 pages = (size + PGSIZE - 1) / PGSIZE;
+            (void)req;
+
+            // calculate order: 2^order >= pages
+            while ((1UL << order) < pages)
+            {
+                order++;
+            }
+            if (order > MAX_ORDER)
+                return nullptr;
+
+            void *p = PMM::alloc_pages(order);
+            // Here, page->flags defaults to PG_used, not PG_slab
+            // kfree needs to be able to recognize it
+            return p;
+        }
+
+        int idx = 0;
+        while (idx < NUM_CACHES && kmalloc_sizes[idx] < size)
+        {
+            idx++;
+        }
+
+        struct KmemCache *c = &caches[idx];
+        c->lock.acquire();
+
+        if (!c->slab_list)
+        {
+            c->slab_list = grow_cache(c);
+            if (!c->slab_list)
+            {
+                c->lock.release();
+                return nullptr;
+            }
+        }
+
+        struct SlabHeader *slab = c->slab_list;
+
+        int obj_idx = slab->free_idx;
+        char *base = (char *)(slab + 1);
+        void *ptr = base + obj_idx * c->obj_size;
+
+        uint16 *link = (uint16 *)ptr;
+        slab->free_idx = *link;
+        slab->used_count++;
+
+        if (slab->free_idx == 0xFFFF)
+        {
+            if (slab->next)
+            {
+                c->slab_list = slab->next;
+                slab->next = nullptr;
+            }
+            else
+            {
+                c->slab_list = nullptr;
+            }
+        }
+
+        c->lock.release();
+
+        memset(ptr, 0, size);
+        return ptr;
+    }
+
+    void kfree(void *ptr)
+    {
+        if (!ptr)
+            return;
+
+        struct Page *page = virt_to_page(ptr);
+        if (!page)
+            return;
+
+        if (!(page->flags & PG_slab))
+        {
+            PMM::free_pages(ptr, page->order);
+            return;
+        }
+
+        int cache_idx = page->order;
+        if (cache_idx >= NUM_CACHES)
+        {
+            Drivers::uart_puts("Slab: Corrupted page metadata\n");
+            return;
+        }
+
+        struct KmemCache *c = &caches[cache_idx];
+        c->lock.acquire();
+
+        struct SlabHeader *slab = (struct SlabHeader *)PGROUNDDOWN((uint64)ptr);
+
+        char *base = (char *)(slab + 1);
+        int obj_idx = ((uint64)ptr - (uint64)base) / c->obj_size;
+
+        uint16 *link = (uint16 *)ptr;
+        *link = slab->free_idx;
+        slab->free_idx = (uint16)obj_idx;
+
+        slab->used_count--;
+
+        if (slab->used_count == 0)
+        {
+            struct SlabHeader **pp = &c->slab_list;
+            // bool found = false;
+            while (*pp)
+            {
+                if (*pp == slab)
+                {
+                    *pp = slab->next;
+                    // found = true;
+                    break;
+                }
+                pp = &(*pp)->next;
+            }
+
+            page->flags &= ~PG_slab;
+            PMM::free_page(slab);
+        }
+        else
+        {
+
+            if (slab->used_count == c->obj_count_per_slab - 1)
+            {
+                slab->next = c->slab_list;
+                c->slab_list = slab;
+            }
+        }
+
+        c->lock.release();
+    }
+}
\ No newline at end of file
diff --git a/src/kernel/mm/vm.cc b/src/kernel/mm/vm.cc
index 650fcf1..c3c5f01 100644
--- a/src/kernel/mm/vm.cc
+++ b/src/kernel/mm/vm.cc
@@ -1,75 +1,88 @@
 #include "kernel/mm.h"
 #include "kernel/pmm.h"
 #include "kernel/riscv.h"
-#include "drivers/uart.h"
 #include "common/types.h"
-#include "kernel/config.h"
 #include "lib/string.h"
+#include "drivers/uart.h"
 
+// COW Flag (PTE bit 8)
+#define PTE_COW (1L << 8)
 
-extern char text_start[];
-extern char text_end[];
-extern char rodata_start[];
-extern char rodata_end[];
-extern char data_start[];
-extern char data_end[];
-extern char bss_start[];
-extern char bss_end[];
-extern char kernel_end[];
-
-extern char uservec[];
-extern char userret[];
-extern char trampoline[];
-
-VM::pagetable_t kernel_pagetable;
+extern "C"
+{
+    extern char kernel_end[];
+    extern char trampoline[];
+}
 
 namespace VM
 {
-    // pagetable_t kernel_pagetable;
 
-    pte_t *walk(pagetable_t pagetable, uint64 va, int alloc)
+    uint64 *kernel_pagetable;
+
+    static void page_ref_inc(uint64 pa)
+    {
+        struct Page *p = PMM::pa_to_page(pa);
+        if (p)
+        {
+            __atomic_add_fetch(&p->refcnt, 1, __ATOMIC_RELAXED);
+        }
+    }
+
+    static bool page_ref_dec(uint64 pa)
+    {
+        struct Page *p = PMM::pa_to_page(pa);
+        if (p)
+        {
+            if (__atomic_sub_fetch(&p->refcnt, 1, __ATOMIC_RELEASE) == 0)
+            {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    static uint64 *walk(uint64 *pagetable, uint64 va, int alloc)
     {
         if (va >= MAXVA)
-            return 0;
-        
+            return nullptr;
+
         for (int level = 2; level > 0; level--)
         {
-            pte_t *pte = &pagetable[PX(level, va)];
-
+            uint64 *pte = &pagetable[PX(level, va)];
             if (*pte & PTE_V)
             {
-                pagetable = (pagetable_t)PTE2PA(*pte);
+                pagetable = (uint64 *)PTE2PA(*pte);
             }
             else
             {
-                if (!alloc || (pagetable = (pagetable_t)PMM::alloc_page()) == 0)
-                    return 0;
+                if (!alloc || (pagetable = (uint64 *)PMM::alloc_page()) == nullptr)
+                    return nullptr;
+
                 memset(pagetable, 0, PGSIZE);
-                // pagetable = (pagetable_t)PMM::alloc_page();
                 *pte = PA2PTE(pagetable) | PTE_V;
             }
         }
         return &pagetable[PX(0, va)];
     }
 
-    int mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
+    int mappages(uint64 *pagetable, uint64 va, uint64 size, uint64 pa, int perm)
     {
         uint64 a, last;
-        pte_t *pte;
+        uint64 *pte;
+
+        if (size == 0)
+            return -1;
 
         a = PGROUNDDOWN(va);
         last = PGROUNDDOWN(va + size - 1);
 
-        for(;;)
+        for (;;)
         {
-            if((pte = walk(pagetable, a, 1)) == 0)
+            if ((pte = walk(pagetable, a, 1)) == nullptr)
                 return -1;
 
             if (*pte & PTE_V)
             {
-                Drivers::uart_puts("Remap panic: ");
-                Drivers::print_hex(a);
-                Drivers::uart_puts("\n");
                 return -1;
             }
 
@@ -83,228 +96,281 @@ namespace VM
         return 0;
     }
 
-    void kvmmap(uint64 va, uint64 pa, uint64 sz, int perm)
+    uint64 *uvmcreate()
     {
-        if (mappages(kernel_pagetable, va, sz, pa, perm) != 0)
+        uint64 *pagetable = (uint64 *)PMM::alloc_page();
+        if (pagetable)
         {
-            Drivers::uart_puts("kvmmap: failed to map\n");
+            memset(pagetable, 0, PGSIZE);
+
+            if (mappages(pagetable, TRAMPOLINE, PGSIZE, (uint64)trampoline, PTE_R | PTE_X) < 0)
+            {
+                PMM::free_page(pagetable);
+                return nullptr;
+            }
         }
+        return pagetable;
     }
 
-    void kvminit()
+    void uvminit(uint64 *pagetable, uchar *src, uint64 sz)
     {
-        kernel_pagetable = (pagetable_t)PMM::alloc_page();
-        memset(kernel_pagetable, 0, PGSIZE);
-
-        kvmmap(g_uart_base, g_uart_base, PGSIZE, PTE_R | PTE_W);
-
-        kvmmap(g_plic_base, g_plic_base, 0x400000, PTE_R | PTE_W);
-
-        // kvmmap(g_clint_base, g_clint_base, 0x10000, PTE_R | PTE_W);
-
-        kvmmap((uint64)text_start, (uint64)text_start,
-               (uint64)text_end - (uint64)text_start, PTE_R | PTE_X);
-
-        kvmmap((uint64)rodata_start, (uint64)rodata_start,
-               (uint64)rodata_end - (uint64)rodata_start, PTE_R);
+        char *mem;
 
-        kvmmap((uint64)data_start, (uint64)data_start,
-               (uint64)kernel_end - (uint64)data_start, PTE_R | PTE_W);
+        if (sz >= PGSIZE)
+        {
+            Drivers::uart_puts("uvminit: initcode too big!\n");
+            return;
+        }
 
-        kvmmap((uint64)kernel_end, (uint64)kernel_end,
-               (uint64)PHYSTOP - (uint64)kernel_end, PTE_R | PTE_W);
+        mem = (char *)PMM::alloc_page();
+        memset(mem, 0, PGSIZE);
 
-        kvmmap(TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
-    }
-    void kvminithart()
-    {
-        w_satp(MAKE_SATP(kernel_pagetable));
-        sfence_vma();
+        mappages(pagetable, 0, PGSIZE, (uint64)mem, PTE_W | PTE_R | PTE_X | PTE_U);
+        memmove(mem, src, sz);
     }
 
-
-    pagetable_t uvmcreate()
+    uint64 uvmalloc(uint64 *pagetable, uint64 oldsz, uint64 newsz, int xperm)
     {
-        pagetable_t pagetable = (pagetable_t)PMM::alloc_page();
-        if (pagetable == 0)
-            return 0;
-        memset(pagetable, 0, PGSIZE);
+        char *mem;
+        uint64 a;
+
+        if (newsz < oldsz)
+            return oldsz;
 
-        if (mappages(pagetable, TRAMPOLINE, PGSIZE, (uint64)trampoline, PTE_R | PTE_X) < 0)
+        oldsz = PGROUNDUP(oldsz);
+        for (a = oldsz; a < newsz; a += PGSIZE)
         {
-            return 0;
+            mem = (char *)PMM::alloc_page();
+            if (mem == nullptr)
+            {
+                uvmdealloc(pagetable, a, oldsz);
+                return 0;
+            }
+            memset(mem, 0, PGSIZE);
+            if (mappages(pagetable, a, PGSIZE, (uint64)mem, PTE_R | PTE_U | xperm) != 0)
+            {
+                PMM::free_page(mem);
+                uvmdealloc(pagetable, a, oldsz);
+                return 0;
+            }
         }
-
-        return pagetable;
+        return newsz;
     }
 
-    void uvminit(pagetable_t pagetable, uchar *src, uint sz)
+    uint64 uvmdealloc(uint64 *pagetable, uint64 oldsz, uint64 newsz)
     {
-        char *mem;
+        if (newsz >= oldsz)
+            return oldsz;
 
-        if (sz >= PGSIZE)
+        if (PGROUNDUP(newsz) < PGROUNDUP(oldsz))
         {
-            Drivers::uart_puts("uvminit: initcode too big\n");
-            return;
+            int npages = (PGROUNDUP(oldsz) - PGROUNDUP(newsz)) / PGSIZE;
+            uvmunmap(pagetable, PGROUNDUP(newsz), npages, 1);
         }
 
-        mem = (char*)PMM::alloc_page();
-        memset(mem, 0, PGSIZE);
-
-        mappages(pagetable, 0, PGSIZE, (uint64)mem, PTE_W|PTE_R|PTE_X|PTE_U);
-
-        memmove(mem, src, sz);
-
-        uint64 trap_pa = (uint64)uservec;
-        uint64 trap_page_start = PGROUNDDOWN(trap_pa);
-
-        mappages(pagetable, trap_page_start, PGSIZE, trap_page_start, PTE_R | PTE_X);
-
-        mappages(pagetable, trap_page_start + PGSIZE, PGSIZE, trap_page_start + PGSIZE, PTE_R | PTE_X);
+        return newsz;
     }
 
-    void uvmunmap(pagetable_t pagetable, uint64 va, uint64 npages, int do_free)
+    void uvmunmap(uint64 *pagetable, uint64 va, uint64 npages, int do_free)
     {
         uint64 a;
-        pte_t *pte;
+        uint64 *pte;
 
         if ((va % PGSIZE) != 0)
-        {
-            Drivers::uart_puts("uvmunmap: not aligned\n");
             return;
-        }
 
-        for(a = va; a < va + npages * PGSIZE; a += PGSIZE)
+        for (a = va; a < va + npages * PGSIZE; a += PGSIZE)
         {
-            if ((pte = walk(pagetable, a, 0)) == 0)
+            if ((pte = walk(pagetable, a, 0)) == nullptr)
                 continue;
 
-            if((*pte & PTE_V) == 0)
+            if ((*pte & PTE_V) == 0)
                 continue;
 
-            if (PTE_FLAGS(*pte) == PTE_V)
+            if ((PTE_FLAGS(*pte) & (PTE_R | PTE_W | PTE_X)) == 0)
             {
-                Drivers::uart_puts("uvmunmap: not a leaf\n");
+                continue;
             }
 
             if (do_free)
             {
                 uint64 pa = PTE2PA(*pte);
-                PMM::free_page((void *)pa);
+                if (page_ref_dec(pa))
+                {
+                    PMM::free_page((void *)pa);
+                }
             }
             *pte = 0;
         }
     }
 
-    void uvmfree(pagetable_t pagetable, uint64 sz)
+    static void freewalk(uint64 *pagetable)
     {
-        if(sz > 0)
-            uvmunmap(pagetable, 0, PGROUNDDOWN(sz)/PGSIZE, 1);
+        for (int i = 0; i < 512; i++)
+        {
+            uint64 pte = pagetable[i];
+            if ((pte & PTE_V) && (pte & (PTE_R | PTE_W | PTE_X)) == 0)
+            {
+                uint64 *child = (uint64 *)PTE2PA(pte);
+                freewalk(child);
+                pagetable[i] = 0;
+            }
+        }
+        PMM::free_page(pagetable);
+    }
 
-        PMM::free_page((void *)pagetable);
+    void uvmfree(uint64 *pagetable, uint64 sz)
+    {
+        if (sz > 0)
+            uvmunmap(pagetable, 0, PGROUNDUP(sz) / PGSIZE, 1);
+        freewalk(pagetable);
     }
 
-    int uvmcopy(pagetable_t oldpt, pagetable_t newpt, uint64 sz)
+    int uvmcopy(uint64 *oldpt, uint64 *newpt, uint64 sz)
     {
-        pte_t *pte;
+        uint64 *pte;
         uint64 pa, i;
         uint flags;
-        char *mem;
 
-        for(i = 0; i < sz; i += PGSIZE)
+        for (i = 0; i < sz; i += PGSIZE)
         {
-            if((pte = walk(oldpt, i, 0)) == 0)
+            if ((pte = walk(oldpt, i, 0)) == nullptr)
                 continue;
             if ((*pte & PTE_V) == 0)
                 continue;
-            
+
             pa = PTE2PA(*pte);
             flags = PTE_FLAGS(*pte);
 
-            if ((mem = (char *)PMM::alloc_page()) == 0)
-                goto err;
-
-            memmove(mem, (char *)pa, PGSIZE);
+            if (flags & PTE_W)
+            {
+                flags &= ~PTE_W;
+                flags |= PTE_COW;
+                *pte = PA2PTE(pa) | flags;
+            }
 
-            if (mappages(newpt, i, PGSIZE, (uint64)mem, flags) != 0)
+            if (mappages(newpt, i, PGSIZE, pa, flags) != 0)
             {
-                PMM::free_page(mem);
-                goto err;
+                return -1;
             }
+
+            page_ref_inc(pa);
         }
+
+        sfence_vma();
         return 0;
-    err:
-        uvmunmap(newpt, 0, i / PGSIZE, 1);
-        return -1;
     }
 
-    uint64 uvmalloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz)
+    int copyout(uint64 *pagetable, uint64 dstva, char *src, uint64 len)
     {
-        char *mem;
-        uint64 a;
-
-        if (newsz < oldsz)
-            return oldsz;
+        uint64 n, va0, pa0;
+        uint64 *pte;
 
-        oldsz = PGROUNDUP(oldsz);
-        for (a = oldsz; a < newsz; a += PGSIZE)
+        while (len > 0)
         {
-            mem = (char *)PMM::alloc_page();
-            if (mem == 0)
-            {
-                uvmdealloc(pagetable, a, oldsz);
-                return 0;
-            }
-            memset(mem, 0, PGSIZE);
-            if (mappages(pagetable, a, PGSIZE, (uint64)mem, PTE_W | PTE_X | PTE_R | PTE_U) != 0)
-            {
-                PMM::free_page(mem);
-                uvmdealloc(pagetable, a, oldsz);
-                return 0;
-            }
-        }
-        return newsz;
-    }
+            va0 = PGROUNDDOWN(dstva);
+            pte = walk(pagetable, va0, 0);
+            if (pte == nullptr || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
+                return -1;
 
-    uint64 uvmdealloc(pagetable_t pagetable, uint64 oldsz, uint64 newsz)
-    {
-        if (newsz >= oldsz)
-            return oldsz;
+            pa0 = PTE2PA(*pte);
+            n = PGSIZE - (dstva - va0);
+            if (n > len)
+                n = len;
 
-        if (PGROUNDUP(newsz) < PGROUNDUP(oldsz))
-        {
-            int npages = (PGROUNDUP(oldsz) - PGROUNDUP(newsz)) / PGSIZE;
-            uvmunmap(pagetable, PGROUNDUP(newsz), npages, 1);
-        }
+            memmove((void *)(pa0 + (dstva - va0)), src, n);
 
-        return newsz;
+            len -= n;
+            src += n;
+            dstva = va0 + PGSIZE;
+        }
+        return 0;
     }
 
-    int copyin(pagetable_t pagetable, char *dst, uint64 srcva, uint64 len)
+    int copyin(uint64 *pagetable, char *dst, uint64 srcva, uint64 len)
     {
         uint64 n, va0, pa0;
-        pte_t *pte;
+        uint64 *pte;
 
         while (len > 0)
         {
             va0 = PGROUNDDOWN(srcva);
             pte = walk(pagetable, va0, 0);
-            if(pte == 0 || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
-            {
+            if (pte == nullptr || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
                 return -1;
-            }
+
             pa0 = PTE2PA(*pte);
             n = PGSIZE - (srcva - va0);
-            if(n > len)
+            if (n > len)
                 n = len;
-            
-            memmove(dst, (void*)(pa0 + (srcva - va0)), n);
+
+            memmove(dst, (void *)(pa0 + (srcva - va0)), n);
 
             len -= n;
             dst += n;
             srcva = va0 + PGSIZE;
         }
         return 0;
-        
+    }
+
+    int handle_cow_fault(uint64 *pagetable, uint64 va)
+    {
+        if (va >= MAXVA)
+            return -1;
+
+        uint64 *pte = walk(pagetable, va, 0);
+        if (pte == nullptr || (*pte & PTE_V) == 0 || (*pte & PTE_U) == 0)
+            return -1;
+
+        if ((*pte & PTE_COW) == 0)
+            return -1;
+
+        uint64 pa = PTE2PA(*pte);
+        uint64 flags = PTE_FLAGS(*pte);
+
+        char *mem = (char *)PMM::alloc_page();
+        if (mem == nullptr)
+            return -1;
+
+        memmove(mem, (char *)pa, PGSIZE);
+
+        uint64 new_flags = (flags | PTE_W) & ~PTE_COW;
+        *pte = PA2PTE((uint64)mem) | new_flags | PTE_V;
+
+        if (page_ref_dec(pa))
+        {
+            PMM::free_page((void *)pa);
+        }
+
+        sfence_vma();
+        return 0;
+    }
+
+    static void kvmmap(uint64 va, uint64 pa, uint64 sz, int perm)
+    {
+        if (mappages(kernel_pagetable, va, sz, pa, perm) != 0)
+        {
+            Drivers::uart_puts("kvmmap: failed\n");
+            while (1)
+                ;
+        }
+    }
+
+    void kvminit()
+    {
+        kernel_pagetable = uvmcreate();
+
+        kvmmap(0x10000000, 0x10000000, PGSIZE, PTE_R | PTE_W);
+        kvmmap(0x0c000000, 0x0c000000, 0x400000, PTE_R | PTE_W);
+        kvmmap(0x02000000, 0x02000000, 0x10000, PTE_R | PTE_W); // CLINT
+        kvmmap(KERNBASE, KERNBASE, PHYSTOP - KERNBASE, PTE_R | PTE_W | PTE_X);
+        // kvmmap(TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
+    }
+
+    void kvminithart()
+    {
+        w_satp(MAKE_SATP(kernel_pagetable));
+        sfence_vma();
+        Drivers::uart_puts("[Boot] MMU ENABLED! Using virtual addresses.\n");
     }
 }
\ No newline at end of file