1 files changed, 93 insertions, 18 deletions

diff --git a/sys/vm/vm_physmem.c b/sys/vm/vm_physmem.c
index c7fcedb..0bd2d54 100644
--- a/sys/vm/vm_physmem.c
+++ b/sys/vm/vm_physmem.c
@@ -32,15 +32,22 @@
 #include <sys/limine.h>
 #include <sys/syslog.h>
 #include <sys/spinlock.h>
+#include <sys/panic.h>
 #include <vm/physmem.h>
 #include <vm/vm.h>
 #include <string.h>
 
-size_t highest_frame_idx = 0;
-size_t bitmap_size = 0;
-size_t bitmap_free_start = 0;
+#define BYTES_PER_MIB 8388608
 
-uint8_t *bitmap;
+static size_t pages_free = 0;
+static size_t pages_used = 0;
+static size_t pages_total = 0;
+static size_t highest_frame_idx = 0;
+static size_t bitmap_size = 0;
+static size_t bitmap_free_start = 0;
+static ssize_t last_idx = 0;
+
+static uint8_t *bitmap;
 static struct limine_memmap_response *resp = NULL;
 static struct spinlock lock = {0};
 
@@ -59,9 +66,11 @@ physmem_populate_bitmap(void)
 
     for (size_t i = 0; i < resp->entry_count; ++i) {
         ent = resp->entries[i];
+        pages_total += ent->length / DEFAULT_PAGESIZE;
 
         if (ent->type != LIMINE_MEMMAP_USABLE) {
             /* This memory is not usable */
+            pages_used += ent->length / DEFAULT_PAGESIZE;
             continue;
         }
 
@@ -72,6 +81,8 @@ physmem_populate_bitmap(void)
         for (size_t j = 0; j < ent->length; j += DEFAULT_PAGESIZE) {
             clrbit(bitmap, (ent->base + j) / DEFAULT_PAGESIZE);
         }
+
+        pages_free += ent->length / DEFAULT_PAGESIZE;
     }
 }
 
@@ -137,29 +148,59 @@ physmem_init_bitmap(void)
  *
  * @count: Number of frames to allocate.
  */
-uintptr_t
-vm_alloc_frame(size_t count)
+static uintptr_t
+__vm_alloc_frame(size_t count)
 {
     size_t frames = 0;
+    ssize_t idx = -1;
     uintptr_t ret = 0;
 
-    spinlock_acquire(&lock);
-    for (size_t i = 0; i < highest_frame_idx; ++i) {
+    for (size_t i = last_idx; i < highest_frame_idx; ++i) {
         if (!testbit(bitmap, i)) {
-            /* We have a free page */
-            if (++frames != count) {
-                continue;
-            }
+            if (idx < 0)
+                idx = i;
+            if (++frames >= count)
+                break;
 
-            for (size_t j = i; j < i + count; ++j) {
-                setbit(bitmap, j);
-            }
-
-            ret = i * DEFAULT_PAGESIZE;
-            break;
+            continue;
         }
+
+        idx = -1;
+        frames = 0;
+    }
+
+    if (idx < 0 || frames != count) {
+        ret = 0;
+        goto done;
     }
 
+    for (size_t i = idx; i < idx + count; ++i) {
+        setbit(bitmap, i);
+    }
+    ret = idx * DEFAULT_PAGESIZE;
+    last_idx = idx;
+    memset(PHYS_TO_VIRT(ret), 0, count * DEFAULT_PAGESIZE);
+done:
+    return ret;
+}
+
+uintptr_t
+vm_alloc_frame(size_t count)
+{
+    uintptr_t ret;
+
+    spinlock_acquire(&lock);
+    if ((ret = __vm_alloc_frame(count)) == 0) {
+        last_idx = 0;
+        ret = __vm_alloc_frame(count);
+    }
+
+    if (ret == 0) {
+        panic("out of memory\n");
+    }
+
+    pages_used += count;
+    pages_free -= count;
     spinlock_release(&lock);
     return ret;
 }
@@ -169,13 +210,47 @@ vm_free_frame(uintptr_t base, size_t count)
 {
     size_t stop_at = base + (count * DEFAULT_PAGESIZE);
 
+    base = ALIGN_UP(base, DEFAULT_PAGESIZE);
+
     spinlock_acquire(&lock);
     for (uintptr_t p = base; p < stop_at; p += DEFAULT_PAGESIZE) {
         clrbit(bitmap, p / DEFAULT_PAGESIZE);
     }
+    pages_used -= count;
+    pages_free += count;
     spinlock_release(&lock);
 }
 
+/*
+ * Return the amount of memory in MiB that is
+ * currently allocated.
+ */
+uint32_t
+vm_mem_used(void)
+{
+    return (pages_used * DEFAULT_PAGESIZE) / BYTES_PER_MIB;
+}
+
+/*
+ * Return the amount of memory in MiB that is
+ * currently free.
+ */
+uint32_t
+vm_mem_free(void)
+{
+    return (pages_free * DEFAULT_PAGESIZE) / BYTES_PER_MIB;
+}
+
+/*
+ * Return the total amount of memory supported
+ * by the machine.
+ */
+size_t
+vm_mem_total(void)
+{
+    return (pages_total * DEFAULT_PAGESIZE) / BYTES_PER_MIB;
+}
+
 void
 vm_physmem_init(void)
 {