From b15c037b16dba4a734e249181717b4a08bf453c4 Mon Sep 17 00:00:00 2001
From: nub31 <oliverlhs@outlook.com>
Date: Sat, 6 Sep 2025 18:13:33 +0200
Subject: [PATCH] improve vmm api

---
 src/arch/x86_64/mem/pmm.h |  5 +++--
 src/arch/x86_64/mem/vmm.c | 43 ++++++++++++++++++++++++++++-----------
 src/arch/x86_64/mem/vmm.h | 17 +++++++++++++---
 3 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/src/arch/x86_64/mem/pmm.h b/src/arch/x86_64/mem/pmm.h
index ceecc1b..647ff42 100644
--- a/src/arch/x86_64/mem/pmm.h
+++ b/src/arch/x86_64/mem/pmm.h
@@ -12,8 +12,9 @@
 
 void pmm_init(multiboot_info_t* info);
 
-// Allocate a 2mb physical page
+// Low level function to allocate a 2mb physical page and return the physical address
+// Return value 0 indicates out of memory
 u64 pmm_alloc();
 
-// Free the 2mb physical page at the specified address
+// Low level function to free a 2mb physical page
 void pmm_free(u64 physical_address);
\ No newline at end of file
diff --git a/src/arch/x86_64/mem/vmm.c b/src/arch/x86_64/mem/vmm.c
index 3194591..685e07c 100644
--- a/src/arch/x86_64/mem/vmm.c
+++ b/src/arch/x86_64/mem/vmm.c
@@ -15,7 +15,7 @@
 
 extern u64 pml4[];
 
-static u64 create_pte(u64 physical_address)
+static u64 create_pte(u64 physical_address, u32 flags)
 {
     if (physical_address & (PAGE_SIZE - 1))
     {
@@ -23,10 +23,10 @@ static u64 create_pte(u64 physical_address)
         panic("Failed to create PTE");
     }
 
-    return (physical_address & PTE_MASK) | PTE_PRESENT | PTE_WRITABLE | PTE_PS;
+    return (physical_address & PTE_MASK) | (flags | PTE_PRESENT);
 }
 
-void vmm_map(u64 physical_address, u64 virtual_address)
+void vmm_map(u64 physical_address, u64 virtual_address, u32 flags)
 {
     u64 pml4_idx = PML4_INDEX(virtual_address);
     u64 pdpt_idx = PDPT_INDEX(virtual_address);
@@ -58,7 +58,7 @@ void vmm_map(u64 physical_address, u64 virtual_address)
         panic("Failed to map virtual to physical page");
     }
 
-    pd_physical_address[pd_idx] = create_pte(physical_address);
+    pd_physical_address[pd_idx] = create_pte(physical_address, flags);
 }
 
 u64 vmm_unmap(u64 virtual_address)
@@ -97,20 +97,39 @@ u64 vmm_unmap(u64 virtual_address)
     return physical_address;
 }
 
-void* vmm_alloc(u64 virtual_address)
+u64 vmm_alloc_address(u64 page_count)
 {
-    u64 physical_address = pmm_alloc();
-    if (!physical_address)
+    panic("not implemented");
+}
+
+u64 vmm_free_address(u64 virtual_address, u64 page_count)
+{
+    panic("not implemented");
+}
+
+void* vmm_alloc(u64 page_count)
+{
+    u64 virtual_address = vmm_alloc_address(page_count);
+
+    for (size_t i = 0; i < page_count; i++)
     {
-        panic("Out of physical memory");
+        u64 physical_address = pmm_alloc();
+        if (!physical_address)
+        {
+            panic("Out of physical memory");
+        }
+
+        vmm_map(physical_address, virtual_address + (i * PAGE_SIZE), PTE_WRITABLE | PTE_PS);
     }
 
-    vmm_map(physical_address, virtual_address);
     return (void*)virtual_address;
 }
 
-void vmm_free(u64 virtual_address)
+void vmm_free(u64 virtual_address, u64 page_count)
 {
-    u64 physical_address = vmm_unmap(virtual_address);
-    pmm_free(physical_address);
+    for (size_t i = 0; i < page_count; i++)
+    {
+        u64 physical_address = vmm_unmap(virtual_address + (i * PAGE_SIZE));
+        pmm_free(physical_address);
+    }
 }
\ No newline at end of file
diff --git a/src/arch/x86_64/mem/vmm.h b/src/arch/x86_64/mem/vmm.h
index 8afa56a..5a5a968 100644
--- a/src/arch/x86_64/mem/vmm.h
+++ b/src/arch/x86_64/mem/vmm.h
@@ -2,8 +2,19 @@
 
 #include "std.h"
 
-void vmm_map(u64 physical_address, u64 virtual_address);
+// Low level function to map a virtual address to a physical address
+void vmm_map(u64 physical_address, u64 virtual_address, u32 flags);
+// Low level function to unmap a virtual address from a physical address
 u64 vmm_unmap(u64 virtual_address);
 
-void* vmm_alloc(u64 virtual_address);
-void vmm_free(u64 virtual_address);
\ No newline at end of file
+// Allocates a free page aligned block of virtual addresses
+u64 vmm_alloc_address(u64 page_count);
+// Frees a block of virtual addresses previously allocated via `vmm_alloc_address`
+// Only use this function for pages mapped via `vmm_alloc_address`
+u64 vmm_free_address(u64 virtual_address, u64 page_count);
+
+// Allocates and maps `page_count` continuous pages and returns the virtual address of the first page
+void* vmm_alloc(u64 page_count);
+// Frees the pages allocated via `vmm_alloc` at the specified virtual address
+// Only use this function for pages mapped via `vmm_alloc`
+void vmm_free(u64 virtual_address, u64 page_count);
\ No newline at end of file