kernel/heap: improve large alignment handling

2 files changed, 68 insertions, 29 deletions

diff --git a/kernel/include/kernel/memory/block_list_allocator.hpp b/kernel/include/kernel/memory/block_list_allocator.hpp
index 977ce89..5e81c44 100644
--- a/kernel/include/kernel/memory/block_list_allocator.hpp
+++ b/kernel/include/kernel/memory/block_list_allocator.hpp
@@ -42,12 +42,21 @@ namespace kernel::memory
   private:
     struct block_header final
     {
-      std::size_t size;
+      std::size_t usable_size;
       bool free;
       block_header * next;
       block_header * prev;
     };
 
+    //! The size of the metadata required for each allocated block.
+    //!
+    //! Each allocated block carries a block header, like any unallocated one, but in addition also has a back-pointer
+    //! to the block header to support padding due to alignment.
+    constexpr auto static allocated_metadata_size = sizeof(block_header) + sizeof(block_header *);
+
+    //! The minimum number of bytes for an allocation.
+    constexpr auto static minimum_allocation_size = 16uz;
+
     //! Try to expand the heap to accommodate the given size.
     //!
     //! @param delta The size to expand the heap by.
diff --git a/kernel/src/memory/block_list_allocator.cpp b/kernel/src/memory/block_list_allocator.cpp
index be83236..1bd17ff 100644
--- a/kernel/src/memory/block_list_allocator.cpp
+++ b/kernel/src/memory/block_list_allocator.cpp
@@ -18,10 +18,10 @@ namespace kernel::memory
 
   namespace
   {
-    [[nodiscard]] constexpr auto align_up(std::uintptr_t value, std::size_t alignment) noexcept -> std::uintptr_t
+    [[nodiscard]] constexpr auto align_up(std::byte * pointer, std::align_val_t alignment) noexcept -> std::byte *
     {
-      auto const remainder = value % alignment;
-      return remainder == 0 ? value : value + alignment - remainder;
+      auto const remainder = std::bit_cast<std::uintptr_t>(pointer) % static_cast<std::size_t>(alignment);
+      return remainder == 0 ? pointer : pointer + static_cast<std::size_t>(alignment) - remainder;
     }
   }  // namespace
 
@@ -37,31 +37,59 @@ namespace kernel::memory
   {
     kstd::lock_guard guard{m_lock};
 
-    auto const alignment_value = static_cast<std::size_t>(alignment);
-    auto const search_size = size + alignment_value + sizeof(block_header *);
-
     for (auto attempt = 0uz; attempt < 2uz; ++attempt)
     {
       auto current = m_block_list;
       while (current != nullptr)
       {
-        if (current->free && current->size >= search_size)
+        if (!current->free)
+        {
+          current = current->next;
+          continue;
+        }
+
+        auto const raw_block = reinterpret_cast<std::byte *>(current);
+        auto const unaligned_payload = raw_block + allocated_metadata_size;
+        auto const aligned_payload = align_up(unaligned_payload, alignment);
+        auto const required_padding = static_cast<std::size_t>(aligned_payload - unaligned_payload);
+        auto const total_required_size = required_padding + allocated_metadata_size + size;
+
+        if (current->usable_size >= total_required_size)
         {
-          auto const payload_start = std::bit_cast<std::uintptr_t>(current) + sizeof(block_header);
-          auto const aligned_pointer = align_up(payload_start + sizeof(block_header *), alignment_value);
-          auto const padding = aligned_pointer - payload_start;
+          auto const payload_header = aligned_payload - sizeof(block_header *) - sizeof(block_header);
+          auto const front_padding = static_cast<std::size_t>(payload_header - raw_block);
+
+          auto payload_block = current;
 
-          split(current, size, padding);
-          current->free = false;
+          if (front_padding >= allocated_metadata_size + minimum_allocation_size)
+          {
+            payload_block = reinterpret_cast<block_header *>(payload_header);
+            std::construct_at(payload_block, current->usable_size - front_padding, true, current->next);
 
-          auto back_pointer = std::bit_cast<block_header **>(aligned_pointer - sizeof(block_header *));
-          *back_pointer = current;
+            if (payload_block->next)
+            {
+              payload_block->next->prev = payload_block;
+            }
 
-          return std::bit_cast<void *>(aligned_pointer);
+            current->usable_size = front_padding - allocated_metadata_size;
+            current->next = payload_block;
+            payload_block->prev = current;
+          }
+
+          auto const payload_size =
+              aligned_payload - reinterpret_cast<std::byte *>(payload_block) - allocated_metadata_size + size;
+          split(payload_block, payload_size, 0uz);
+
+          payload_block->free = false;
+
+          auto back_pointer = reinterpret_cast<block_header **>(aligned_payload - sizeof(block_header *));
+          *back_pointer = payload_block;
+
+          return reinterpret_cast<void *>(aligned_payload);
         }
-        current = current->next;
       }
 
+      auto const search_size = size + static_cast<std::size_t>(alignment);
       if (attempt == 0uz && !expand(search_size))
       {
         return nullptr;
@@ -80,8 +108,8 @@ namespace kernel::memory
 
     kstd::lock_guard guard{m_lock};
 
-    auto const aligned_pointer = std::bit_cast<std::uintptr_t>(ptr);
-    auto back_pointer = std::bit_cast<block_header **>(aligned_pointer - sizeof(block_header *));
+    auto const aligned_payload = reinterpret_cast<std::byte *>(ptr);
+    auto back_pointer = reinterpret_cast<block_header **>(aligned_payload - sizeof(block_header *));
     auto block = *back_pointer;
 
     block->free = true;
@@ -90,7 +118,7 @@ namespace kernel::memory
 
   auto block_list_allocator::expand(std::size_t size) noexcept -> bool
   {
-    auto const total_required_size = size + sizeof(block_header);
+    auto const total_required_size = size + allocated_metadata_size;
     auto const frames_needed = (total_required_size + kapi::memory::frame::size - 1) / kapi::memory::frame::size;
 
     auto const flags = kapi::memory::page_mapper::flags::writable | kapi::memory::page_mapper::flags::supervisor_only |
@@ -110,7 +138,8 @@ namespace kernel::memory
     }
 
     auto block = static_cast<block_header *>(m_frontier);
-    std::construct_at(block, frames_needed * kapi::memory::frame::size - sizeof(block_header), true, nullptr, nullptr);
+    std::construct_at(block, frames_needed * kapi::memory::frame::size - allocated_metadata_size, true, nullptr,
+                      nullptr);
 
     m_frontier += frames_needed * kapi::memory::frame::size;
 
@@ -138,7 +167,7 @@ namespace kernel::memory
     if (block->next && block->next->free)
     {
       auto follower = block->next;
-      block->size += follower->size + sizeof(block_header);
+      block->usable_size += follower->usable_size + allocated_metadata_size;
       block->next = follower->next;
       if (block->next)
       {
@@ -150,7 +179,7 @@ namespace kernel::memory
     if (block->prev && block->prev->free)
     {
       auto leader = block->prev;
-      leader->size += block->size + sizeof(block_header);
+      leader->usable_size += block->usable_size + allocated_metadata_size;
       leader->next = block->next;
       if (block->next)
       {
@@ -162,20 +191,21 @@ namespace kernel::memory
 
   auto block_list_allocator::split(block_header * block, std::size_t size, std::size_t padding) noexcept -> void
   {
-    auto const allocation_size = size + padding;
+    auto const new_block_size = size + padding;
 
-    if (block->size > allocation_size + sizeof(block_header) + 16uz)
+    if (block->usable_size > new_block_size + allocated_metadata_size + minimum_allocation_size)
     {
-      auto new_block =
-          std::bit_cast<block_header *>(std::bit_cast<std::uintptr_t>(block) + sizeof(block_header) + allocation_size);
-      std::construct_at(new_block, block->size - allocation_size - sizeof(block_header), true, block->next, block);
+      auto raw_block = reinterpret_cast<std::byte *>(block);
+      auto new_block = reinterpret_cast<block_header *>(raw_block + allocated_metadata_size + new_block_size);
+      std::construct_at(new_block, block->usable_size - new_block_size - allocated_metadata_size, true, block->next,
+                        block);
 
       if (new_block->next)
       {
         new_block->next->prev = new_block;
       }
 
-      block->size = allocation_size;
+      block->usable_size = new_block_size;
       block->next = new_block;
     }
   }