1 files changed, 44 insertions, 20 deletions

diff --git a/arch/x86_64/src/memory/main.cpp b/arch/x86_64/src/memory/main.cpp
index a6f91d9..2746a71 100644
--- a/arch/x86_64/src/memory/main.cpp
+++ b/arch/x86_64/src/memory/main.cpp
@@ -1,33 +1,57 @@
 #include "arch/memory/main.hpp"
 
 #include "arch/exception_handling/assert.hpp"
+#include "arch/kernel/cpu/control_register.hpp"
+#include "arch/kernel/cpu/msr.hpp"
 #include "arch/memory/allocator/area_frame_allocator.hpp"
-#include "arch/memory/cpu/control_register.hpp"
-#include "arch/memory/cpu/msr.hpp"
-#include "arch/memory/heap/heap_allocator.hpp"
+#include "arch/memory/allocator/concept.hpp"
+#include "arch/memory/heap/global_heap_allocator.hpp"
 #include "arch/memory/paging/active_page_table.hpp"
 #include "arch/memory/paging/kernel_mapper.hpp"
 
+#include <optional>
+
 namespace teachos::arch::memory
 {
   namespace
   {
-    auto remap_heap(allocator::area_frame_allocator allocator, paging::active_page_table & active_table) -> void
+    static std::optional<allocator::area_frame_allocator> frame_allocator;
+
+    auto create_frame_allocator(multiboot::memory_information const & memory_information)
+        -> allocator::area_frame_allocator &
     {
-      auto const start_page = paging::virtual_page::containing_address(memory::heap::HEAP_START);
-      auto const end_page =
-          ++(paging::virtual_page::containing_address(memory::heap::HEAP_START + memory::heap::HEAP_SIZE - 1));
-      paging::page_container::iterator const begin{start_page};
-      paging::page_container::iterator const end{end_page};
-      paging::page_container const pages{begin, end};
-
-      for (auto const & page : pages)
-      {
-        active_table.map_page_to_next_free_frame(allocator, page, paging::entry::WRITABLE);
-      }
+      frame_allocator.emplace(memory_information);
+      return frame_allocator.value();
+    }
+
+    auto get_frame_allocator() -> allocator::area_frame_allocator &
+    {
+      exception_handling::assert(frame_allocator.has_value(),
+                                 "[Initialization] Frame allocator has not been created yet");
+      return frame_allocator.value();
     }
   }  // namespace
 
+  auto remap_heap(std::size_t heap_start, std::size_t heap_size, paging::entry::bitset additional_flags = {}) -> void
+  {
+    decltype(auto) allocator = get_frame_allocator();
+    decltype(auto) active_table = paging::active_page_table::create_or_get();
+    auto const start_page = paging::virtual_page::containing_address(heap_start);
+    auto const end_page = ++(paging::virtual_page::containing_address(heap_start + heap_size - 1));
+
+    paging::page_container::iterator const begin{start_page};
+    paging::page_container::iterator const end{end_page};
+    paging::page_container const pages{begin, end};
+
+    constexpr auto base_flags = paging::entry::WRITABLE;
+    auto const flags = base_flags | additional_flags;
+
+    for (auto const & page : pages)
+    {
+      active_table.map_page_to_next_free_frame(allocator, page, flags);
+    }
+  }
+
   auto initialize_memory_management() -> void
   {
     static bool has_been_called = false;
@@ -36,17 +60,17 @@ namespace teachos::arch::memory
     has_been_called = true;
 
     auto const memory_information = multiboot::read_multiboot2();
-    allocator::area_frame_allocator allocator(memory_information);
+    decltype(auto) allocator = create_frame_allocator(memory_information);
 
-    cpu::set_cr0_bit(memory::cpu::cr0_flags::WRITE_PROTECT);
-    cpu::set_efer_bit(memory::cpu::efer_flags::NXE);
+    kernel::cpu::set_cr0_bit(kernel::cpu::cr0_flags::WRITE_PROTECT);
+    kernel::cpu::set_efer_bit(kernel::cpu::efer_flags::NXE);
 
     paging::kernel_mapper kernel(allocator, memory_information);
-    auto & active_table = kernel.remap_kernel();
+    kernel.remap_kernel();
     video::vga::text::write("Kernel remapping successful", video::vga::text::common_attributes::green_on_black);
     video::vga::text::newline();
 
-    remap_heap(allocator, active_table);
+    remap_heap(heap::KERNEL_HEAP_START, heap::KERNEL_HEAP_SIZE);
     video::vga::text::write("Heap remapping successful", video::vga::text::common_attributes::green_on_black);
     video::vga::text::newline();
   }