1 files changed, 24 insertions, 2 deletions
diff --git a/arch/x86_64/src/memory/heap/linked_list_allocator.cpp b/arch/x86_64/src/memory/heap/linked_list_allocator.cpp
index 1c5c443..d922dc8 100644
--- a/arch/x86_64/src/memory/heap/linked_list_allocator.cpp
+++ b/arch/x86_64/src/memory/heap/linked_list_allocator.cpp
@@ -34,7 +34,11 @@ namespace teachos::arch::memory::heap
 
     while (current != nullptr)
     {
-      if (current->size >= size + min_allocatable_size())
+      if (current->size == size)
+      {
+        return remove_free_memory_block(previous, current, size);
+      }
+      else if (current->size >= size + min_allocatable_size())
       {
         auto memory_address = split_free_memory_block(previous, current, size);
         mutex.unlock();
@@ -77,6 +81,24 @@ namespace teachos::arch::memory::heap
     mutex.unlock();
   }
 
+  auto linked_list_allocator::remove_free_memory_block(memory_block * previous_block, memory_block * current_block,
+                                                       std::size_t size) -> void *
+  {
+    auto const start_address = reinterpret_cast<std::size_t>(current_block);
+    // If we want to allocate into the first block that is before any other free block, then there exists no previous
+    // free block (nullptr). Therefore we have to overwrite the first block instead of overwriting its next value.
+    if (previous_block == nullptr)
+    {
+      first = nullptr;
+    }
+    else
+    {
+      previous_block->next = current_block->next;
+    }
+    clear_memory_block_header(current_block);
+    return reinterpret_cast<void *>(start_address);
+  }
+
   auto linked_list_allocator::split_free_memory_block(memory_block * previous_block, memory_block * current_block,
                                                       std::size_t size) -> void *
   {
@@ -92,7 +114,7 @@ namespace teachos::arch::memory::heap
     }
     else
     {
-      previous_block = new_block;
+      previous_block->next = new_block;
     }
     clear_memory_block_header(current_block);
     return reinterpret_cast<void *>(start_address);