From e8845b501243dbe73865ebac0eec3a9618568afb Mon Sep 17 00:00:00 2001
From: Alexander Myltsev <a.myltsev@gmail.com>
Date: Sat, 18 Jan 2025 02:38:44 +0400
Subject: [PATCH] Free user VM.

---
 cpu/memlayout.h |  3 +++
 kernel/mem.c    |  4 ++--
 kernel/mem.h    |  3 ++-
 kernel/vm.c     | 18 ++++++++++++++++++
 proc.c          | 10 ++++++----
 5 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/cpu/memlayout.h b/cpu/memlayout.h
index 414a0be..fd0dcbd 100644
--- a/cpu/memlayout.h
+++ b/cpu/memlayout.h
@@ -7,6 +7,9 @@
 #define PGROUNDUP(sz)  (((sz)+PGSIZE-1) & ~(PGSIZE-1))
 #define PGROUNDDOWN(a) (((a)) & ~((uintptr_t)(PGSIZE-1)))
 
+#define NPDENTRIES      1024    // # directory entries per page directory
+#define NPTENTRIES      1024    // # PTEs per page table
+
 #define PTXSHIFT        12      // offset of PTX in a linear address
 #define PDXSHIFT        22      // offset of PDX in a linear address
 
diff --git a/kernel/mem.c b/kernel/mem.c
index c59f685..5fd7f82 100644
--- a/kernel/mem.c
+++ b/kernel/mem.c
@@ -41,7 +41,7 @@ memset(void *dst, unsigned c, uint64_t n)
 // call to kalloc().  (The exception is when
 // initializing the allocator; see kinit above.)
 void
-kfree(char *v)
+kfree(void *v)
 {
   struct run *r;
 
@@ -51,7 +51,7 @@ kfree(char *v)
   // Fill with junk to catch dangling refs.
   memset(v, 1, PGSIZE);
 
-  r = (struct run*)v;
+  r = v;
   r->next = kmem.freelist;
   kmem.freelist = r;
 }
diff --git a/kernel/mem.h b/kernel/mem.h
index 5d66fcd..be1f5e6 100644
--- a/kernel/mem.h
+++ b/kernel/mem.h
@@ -9,9 +9,10 @@ void* memset(void *dst, unsigned c, uint64_t n);
 
 void freerange(void *vstart, void *vend);
 void* kalloc(void);
-void kfree(char*);
+void kfree(void*);
 
 pde_t *setupkvm();
 void kvmalloc();
 void switchkvm();
 int allocuvm(pde_t *pgdir, uintptr_t base, uintptr_t top);
+void freevm(pde_t *pgdir);
diff --git a/kernel/vm.c b/kernel/vm.c
index 99ba35b..7ece1ec 100644
--- a/kernel/vm.c
+++ b/kernel/vm.c
@@ -92,3 +92,21 @@ int allocuvm(pde_t *pgdir, uintptr_t base, uintptr_t top) {
     }
     return 0;
 }
+
+static void freept(pte_t *pt) {
+    for (int i = 0; i < NPTENTRIES; i++) {
+        if (pt[i] & PTE_P) {
+            kfree((char*)P2V(PTE_ADDR(pt[i])));
+        }
+    }
+    kfree((char*)pt);
+}
+
+void freevm(pde_t *pgdir) {
+    for (int i = 0; i < NPDENTRIES / 2; i++) {
+        if (pgdir[i] & PTE_P) {
+            freept((pte_t*)P2V(PTE_ADDR(pgdir[i])));
+        }
+    }
+    kfree(pgdir);
+}
diff --git a/proc.c b/proc.c
index a89f124..851471b 100644
--- a/proc.c
+++ b/proc.c
@@ -44,11 +44,12 @@ void run_elf(const char* name) {
     }
     if (!vm.user_task) {
         vm.user_task = kalloc();
-        vm.user_task->pgdir = setupkvm();
-        allocuvm(vm.user_task->pgdir, USER_BASE, USER_BASE + statbuf.size);
-        allocuvm(vm.user_task->pgdir, USER_STACK_BASE - 2 * PGSIZE, USER_STACK_BASE);
-        switchuvm(&vm.user_task->tss, vm.user_task->stack.bottom, vm.user_task->pgdir);
     }
+    vm.user_task->pgdir = setupkvm();
+    allocuvm(vm.user_task->pgdir, USER_BASE, USER_BASE + statbuf.size);
+    allocuvm(vm.user_task->pgdir, USER_STACK_BASE - 2 * PGSIZE, USER_STACK_BASE);
+    switchuvm(&vm.user_task->tss, vm.user_task->stack.bottom, vm.user_task->pgdir);
+
     if (read_file(&statbuf, (void*)USER_BASE, 100 << 20) <= 0) {
         printk(name);
         printk(": file not found\n");
@@ -80,6 +81,7 @@ void run_elf(const char* name) {
 _Noreturn void killproc() {
     void* task_stack;
     switchkvm();
+    freevm(vm.user_task->pgdir);
     swtch(&task_stack, vm.kernel_thread);
     __builtin_unreachable();
 }