2 files changed, 68 insertions, 4 deletions

diff --git a/kernel/kernel/memory/bitmap_allocator.cpp b/kernel/kernel/memory/bitmap_allocator.cpp
index 240e2afe..a4e6a8ca 100644
--- a/kernel/kernel/memory/bitmap_allocator.cpp
+++ b/kernel/kernel/memory/bitmap_allocator.cpp
@@ -34,14 +34,42 @@ namespace kernel::memory
     auto free_count = 0uz;
     auto first_free = 0uz;
 
-    for (auto i = 0uz; i < m_frame_count; ++i)
+    for (auto i = m_last_index; i < m_frame_count; ++i)
     {
-      auto const current = (m_last_index + i) % m_frame_count;
-      if (!test(current))
+      if (!test(i))
       {
         if (free_count == 0)
         {
-          first_free = current;
+          first_free = i;
+        }
+
+        ++free_count;
+
+        if (free_count == count)
+        {
+          for (auto j = 0uz; j < count; ++j)
+          {
+            set(first_free + j);
+          }
+          m_last_index = first_free + count;
+          return std::pair{kapi::memory::frame{first_free}, count};
+        }
+      }
+      else
+      {
+        free_count = 0;
+      }
+    }
+
+    free_count = 0;
+
+    for (auto i = 0uz; i < m_last_index; ++i)
+    {
+      if (!test(i))
+      {
+        if (free_count == 0)
+        {
+          first_free = i;
         }
 
         ++free_count;
diff --git a/kernel/kernel/memory/bitmap_allocator.tests.cpp b/kernel/kernel/memory/bitmap_allocator.tests.cpp
index 05d11a33..7aaf43cd 100644
--- a/kernel/kernel/memory/bitmap_allocator.tests.cpp
+++ b/kernel/kernel/memory/bitmap_allocator.tests.cpp
@@ -10,6 +10,7 @@
 #include <limits>
 #include <ranges>
 #include <span>
+#include <utility>
 #include <vector>
 
 constexpr auto all_bits_set = std::numeric_limits<std::uint64_t>::max();
@@ -285,4 +286,39 @@ SCENARIO("Bitmap allocator frame allocation", "[memory][bitmap_allocator]")
       }
     }
   }
+
+  GIVEN("A small storage region")
+  {
+    auto storage = std::vector(1, 0uz);
+
+    AND_GIVEN("an allocator with 3 free frames straddling the wraparound boundary")
+    {
+      auto allocator = kernel::memory::bitmap_frame_allocator{storage, 10};
+      allocator.release_many({kapi::memory::frame{0}, 8});
+      CHECK(allocator.allocate_many(8) == std::pair{kapi::memory::frame{0}, 8});
+
+      allocator.release(kapi::memory::frame{0});
+      allocator.release_many({kapi::memory::frame{8}, 2});
+
+      WHEN("allocating 3 frames")
+      {
+        auto result = allocator.allocate_many(3);
+
+        THEN("the allocation fails")
+        {
+          REQUIRE_FALSE(result.has_value());
+        }
+      }
+
+      WHEN("allocating 2 frames")
+      {
+        auto result = allocator.allocate_many(2);
+
+        THEN("the allocation succeeds")
+        {
+          REQUIRE(result.has_value());
+        }
+      }
+    }
+  }
 }
\ No newline at end of file