x86_64/memory: simplify region allocator

2 files changed, 56 insertions, 30 deletions

diff --git a/arch/x86_64/include/x86_64/memory/region_allocator.hpp b/arch/x86_64/include/x86_64/memory/region_allocator.hpp
index 84b7a94..bf657a0 100644
--- a/arch/x86_64/include/x86_64/memory/region_allocator.hpp
+++ b/arch/x86_64/include/x86_64/memory/region_allocator.hpp
@@ -61,9 +61,12 @@ namespace teachos::memory::x86_64
     //! @note This implementation will never actually release any frames.
     auto release(frame frame) -> void override;
 
+    auto next_free_frame() noexcept -> std::optional<frame>;
+
   private:
     //! Find the next memory area and write it into current_area.
-    auto choose_next_area() -> void;
+    auto choose_next_region() -> void;
+    auto find_next_frame() -> std::optional<frame>;
 
     frame m_next_frame;                      //!< The next available frame.
     std::optional<region> m_current_region;  //!< The memory region currently used for allocation
diff --git a/arch/x86_64/src/memory/region_allocator.cpp b/arch/x86_64/src/memory/region_allocator.cpp
index 0f65b3a..efa1f4a 100644
--- a/arch/x86_64/src/memory/region_allocator.cpp
+++ b/arch/x86_64/src/memory/region_allocator.cpp
@@ -34,59 +34,82 @@ namespace teachos::memory::x86_64
       , m_multiboot_start(frame::containing(mem_info.mbi_range.first))
       , m_multiboot_end(frame::containing(mem_info.mbi_range.second))
   {
-    choose_next_area();
+    choose_next_region();
   }
 
-  auto region_allocator::choose_next_area() -> void
+  auto region_allocator::choose_next_region() -> void
   {
     m_current_region.reset();
-    auto next_area_with_free_frames =
-        m_memory_map | std::views::filter(&multiboot2::memory_map::region::available) |
-        std::views::filter([next = m_next_frame](auto const & region) -> bool { return last_frame(region) >= next; });
 
-    auto lowest_region_with_free_frames = std::ranges::min_element(
-        next_area_with_free_frames, [](auto lhs, auto rhs) -> bool { return lhs.base < rhs.base; });
+    auto remaining_regions =
+        m_memory_map |  //
+        std::views::filter(&multiboot2::memory_map::region::available) |
+        std::views::filter([this](auto const & region) -> bool { return last_frame(region) >= m_next_frame; });
 
-    if (lowest_region_with_free_frames != next_area_with_free_frames.end())
+    auto lowest_region =
+        std::ranges::min_element(remaining_regions, [](auto lhs, auto rhs) -> bool { return lhs.base < rhs.base; });
+
+    if (lowest_region == remaining_regions.end())
     {
-      m_current_region = *lowest_region_with_free_frames;
-      auto start_frame = frame::containing(physical_address{m_current_region->base});
-      if (m_next_frame < start_frame)
-      {
-        m_next_frame = start_frame;
-      }
+      return;
     }
-  }
 
-  auto region_allocator::allocate() noexcept -> std::optional<frame>
-  {
-    if (!m_current_region)
+    m_current_region = *lowest_region;
+    if (auto start_of_region = frame::containing(physical_address{m_current_region->base});
+        start_of_region < m_next_frame)
     {
-      return {};
+      m_next_frame = start_of_region;
     }
+  }
 
-    auto const end_address = physical_address{m_current_region->base + m_current_region->size_in_B - 1};
-    auto end_frame = frame::containing(end_address);
-
-    if (m_next_frame > end_frame)
+  auto region_allocator::find_next_frame() -> std::optional<frame>
+  {
+    if (!m_current_region || m_next_frame > last_frame(*m_current_region))
     {
-      choose_next_area();
+      choose_next_region();
+      if (!m_current_region)
+      {
+        return std::nullopt;
+      }
     }
-    else if (falls_within(m_next_frame, m_kernel_start, m_kernel_end))
+
+    if (falls_within(m_next_frame, m_kernel_start, m_kernel_end))
     {
       m_next_frame = m_kernel_end + 1;
     }
-    else if (falls_within(m_next_frame, m_multiboot_start, m_multiboot_end))
+
+    if (falls_within(m_next_frame, m_multiboot_start, m_multiboot_end))
     {
       m_next_frame = m_multiboot_end + 1;
     }
-    else
+
+    if (falls_within(m_next_frame, m_kernel_start, m_kernel_end))
     {
-      return m_next_frame++;
+      m_next_frame = m_kernel_end + 1;
     }
 
-    return allocate();
+    if (m_next_frame > last_frame(*m_current_region))
+    {
+      choose_next_region();
+    }
+
+    return m_current_region.transform([this](auto) -> auto { return m_next_frame; });
+  }
+
+  auto region_allocator::allocate() noexcept -> std::optional<frame>
+  {
+    return find_next_frame().transform([this](auto frame) -> auto {
+      auto result = frame;
+      ++m_next_frame;
+      return result;
+    });
   }
 
   auto region_allocator::release(frame) -> void {}
+
+  auto region_allocator::next_free_frame() noexcept -> std::optional<frame>
+  {
+    return find_next_frame();
+  }
+
 }  // namespace teachos::memory::x86_64