From bee2de4b3109221ed8376d74e811dddf072f8862 Mon Sep 17 00:00:00 2001
From: Idan Horowitz <idan.horowitz@gmail.com>
Date: Wed, 8 Sep 2021 02:40:44 +0300
Subject: [PATCH] Kernel: Use an IntrusiveRedBlackTree for storing the cr3
 mappings

This ensures we don't allocate when intializing the PageDirectory.
---
 Kernel/Memory/PageDirectory.cpp | 8 ++++----
 Kernel/Memory/PageDirectory.h   | 4 ++++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/Kernel/Memory/PageDirectory.cpp b/Kernel/Memory/PageDirectory.cpp
index d6693f82fa..133bfc9165 100644
--- a/Kernel/Memory/PageDirectory.cpp
+++ b/Kernel/Memory/PageDirectory.cpp
@@ -17,9 +17,9 @@ extern u8 end_of_kernel_image[];
 
 namespace Kernel::Memory {
 
-static Singleton<HashMap<FlatPtr, PageDirectory*>> s_cr3_map;
+static Singleton<IntrusiveRedBlackTree<FlatPtr, PageDirectory, RawPtr<PageDirectory>, &PageDirectory::m_tree_node>> s_cr3_map;
 
-static HashMap<FlatPtr, PageDirectory*>& cr3_map()
+static IntrusiveRedBlackTree<FlatPtr, PageDirectory, RawPtr<PageDirectory>, &PageDirectory::m_tree_node>& cr3_map()
 {
     VERIFY_INTERRUPTS_DISABLED();
     return *s_cr3_map;
@@ -28,7 +28,7 @@ static HashMap<FlatPtr, PageDirectory*>& cr3_map()
 RefPtr<PageDirectory> PageDirectory::find_by_cr3(FlatPtr cr3)
 {
     SpinlockLocker lock(s_mm_lock);
-    return cr3_map().get(cr3).value_or({});
+    return cr3_map().find(cr3);
 }
 
 UNMAP_AFTER_INIT NonnullRefPtr<PageDirectory> PageDirectory::must_create_kernel_page_directory()
@@ -132,7 +132,7 @@ KResultOr<NonnullRefPtr<PageDirectory>> PageDirectory::try_create_for_userspace(
     auto* new_pd = MM.quickmap_pd(*directory, 0);
     memcpy(new_pd, &buffer, sizeof(PageDirectoryEntry));
 
-    cr3_map().set(directory->cr3(), directory.ptr());
+    cr3_map().insert(directory->cr3(), directory);
     return directory;
 }
 
diff --git a/Kernel/Memory/PageDirectory.h b/Kernel/Memory/PageDirectory.h
index 348f43857b..d1755f4f1c 100644
--- a/Kernel/Memory/PageDirectory.h
+++ b/Kernel/Memory/PageDirectory.h
@@ -7,6 +7,7 @@
 #pragma once
 
 #include <AK/HashMap.h>
+#include <AK/IntrusiveRedBlackTree.h>
 #include <AK/RefCounted.h>
 #include <AK/RefPtr.h>
 #include <Kernel/Forward.h>
@@ -46,6 +47,9 @@ public:
 
     RecursiveSpinlock& get_lock() { return m_lock; }
 
+    // This has to be public to let the global singleton access the member pointer
+    IntrusiveRedBlackTreeNode<FlatPtr, PageDirectory, RawPtr<PageDirectory>> m_tree_node;
+
 private:
     PageDirectory();