From 8e351de4defb5d6ebd1ce51d56115218c6d746e3 Mon Sep 17 00:00:00 2001
From: Stephen Miller <millerste004@gmail.com>
Date: Thu, 17 Apr 2025 21:31:33 -0500
Subject: [PATCH] AllocateDirectMemory fixes

search_start and search_end were ignored in certain cases, this fixes that issue.
I've also basically rewritten the function in the process, since the lack of documentation made it difficult to make the proper adjustments.
---
 src/core/memory.cpp | 35 +++++++++++++++--------------------
 1 file changed, 15 insertions(+), 20 deletions(-)

diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index 8b108a654..e71112443 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -139,35 +139,30 @@ PAddr MemoryManager::Allocate(PAddr search_start, PAddr search_end, size_t size,
     alignment = alignment > 0 ? alignment : 16_KB;
 
     auto dmem_area = FindDmemArea(search_start);
+    auto mapping_start = search_start > dmem_area->second.base ? Common::AlignUp(search_start, alignment) : Common::AlignUp(dmem_area->second.base, alignment);
+    auto mapping_end = Common::AlignUp(mapping_start + size, alignment);
 
-    const auto is_suitable = [&] {
-        if (dmem_area == dmem_map.end()) {
-            return false;
-        }
-        const auto aligned_base = Common::AlignUp(dmem_area->second.base, alignment);
-        const auto alignment_size = aligned_base - dmem_area->second.base;
-        const auto remaining_size =
-            dmem_area->second.size >= alignment_size ? dmem_area->second.size - alignment_size : 0;
-        return dmem_area->second.is_free && remaining_size >= size;
-    };
-    while (dmem_area != dmem_map.end() && !is_suitable() &&
-           dmem_area->second.GetEnd() <= search_end) {
-        ++dmem_area;
+    // Find the first free, large enough dmem area in the range.
+    while ((!dmem_area->second.is_free || dmem_area->second.GetEnd() < mapping_end) && dmem_area != dmem_map.end()) {
+        // The current dmem_area isn't suitable, move to the next one.
+        dmem_area++;
+
+        // Update local variables based on the new dmem_area
+        mapping_start = search_start > dmem_area->second.base ? Common::AlignUp(search_start, alignment) : Common::AlignUp(dmem_area->second.base, alignment);
+        mapping_end = Common::AlignUp(mapping_start + size, alignment);
     }
-    if (!is_suitable()) {
+
+    if (dmem_area == dmem_map.end()) {
+        // There are no suitable mappings in this range
         LOG_ERROR(Kernel_Vmm, "Unable to find free direct memory area: size = {:#x}", size);
         return -1;
     }
 
-    // Align free position
-    PAddr free_addr = dmem_area->second.base;
-    free_addr = Common::AlignUp(free_addr, alignment);
-
     // Add the allocated region to the list and commit its pages.
-    auto& area = CarveDmemArea(free_addr, size)->second;
+    auto& area = CarveDmemArea(mapping_start, size)->second;
     area.memory_type = memory_type;
     area.is_free = false;
-    return free_addr;
+    return mapping_start;
 }
 
 void MemoryManager::Free(PAddr phys_addr, size_t size) {