From c83f6c03b508b46d1ebdefdcf812d7b9e97bf9fa Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Mon, 9 Jan 2023 18:08:50 +0100
Subject: [PATCH] Bitmap: Add a 'find_and_toggle' method

Just like find(), but toggles the value when finding it.
Avoids doing this manually in MemoryManager and KernelVM.
---
 kernel/src/memory/KernelVM.cpp      | 3 +--
 kernel/src/memory/MemoryManager.cpp | 6 ++++--
 luna/include/luna/Bitmap.h          | 4 +++-
 luna/src/Bitmap.cpp                 | 7 +++++++
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/kernel/src/memory/KernelVM.cpp b/kernel/src/memory/KernelVM.cpp
index f09eeaa7..48c4b80a 100644
--- a/kernel/src/memory/KernelVM.cpp
+++ b/kernel/src/memory/KernelVM.cpp
@@ -30,12 +30,11 @@ namespace KernelVM
     {
         auto kernelvm_bitmap = g_kernelvm_bitmap.lock();
 
-        const auto maybe_index = kernelvm_bitmap->find(false);
+        const auto maybe_index = kernelvm_bitmap->find_and_toggle(false);
         if (!maybe_index.has_value()) return err(ENOMEM);
 
         const usize index = maybe_index.value();
 
-        kernelvm_bitmap->set(index, true);
         g_used_vm += ARCH_PAGE_SIZE;
 
         return KERNEL_VM_RANGE_START + (index * ARCH_PAGE_SIZE);
diff --git a/kernel/src/memory/MemoryManager.cpp b/kernel/src/memory/MemoryManager.cpp
index 62fd0acc..35fdcc4c 100644
--- a/kernel/src/memory/MemoryManager.cpp
+++ b/kernel/src/memory/MemoryManager.cpp
@@ -146,13 +146,15 @@ namespace MemoryManager
     {
         auto frame_bitmap = g_frame_bitmap.lock();
 
-        const auto maybe_index = frame_bitmap->find(false, start_index);
+        const auto maybe_index = frame_bitmap->find_and_toggle(false, start_index);
         if (!maybe_index.has_value()) return err(ENOMEM);
 
         const usize index = maybe_index.value();
 
         start_index = index + 1;
-        do_lock_frame(index, *frame_bitmap);
+
+        used_mem += ARCH_PAGE_SIZE;
+        free_mem -= ARCH_PAGE_SIZE;
 
         return index * ARCH_PAGE_SIZE;
     }
diff --git a/luna/include/luna/Bitmap.h b/luna/include/luna/Bitmap.h
index 46d6f757..1ef3e239 100644
--- a/luna/include/luna/Bitmap.h
+++ b/luna/include/luna/Bitmap.h
@@ -1,5 +1,5 @@
 #pragma once
-#include <luna/Option.h>
+#include <luna/Result.h>
 #include <luna/Types.h>
 
 class Bitmap
@@ -37,6 +37,8 @@ class Bitmap
 
     Option<usize> find(bool value, usize begin = 0) const;
 
+    Option<usize> find_and_toggle(bool value, usize begin = 0);
+
     void clear(bool value);
     void clear_region(usize start, usize bits, bool value);
 
diff --git a/luna/src/Bitmap.cpp b/luna/src/Bitmap.cpp
index 90530da0..9f4dae55 100644
--- a/luna/src/Bitmap.cpp
+++ b/luna/src/Bitmap.cpp
@@ -125,3 +125,10 @@ Option<usize> Bitmap::find(bool value, usize begin) const
 
     return {};
 }
+
+Option<usize> Bitmap::find_and_toggle(bool value, usize begin)
+{
+    usize index = TRY(find(value, begin));
+    set(index, !value);
+    return index;
+}