kernel/memory: implement basic free-list heap

9 files changed, 449 insertions, 5 deletions

diff --git a/kernel/CMakeLists.txt b/kernel/CMakeLists.txt
index ae2325a..566ab9e 100644
--- a/kernel/CMakeLists.txt
+++ b/kernel/CMakeLists.txt
@@ -11,6 +11,9 @@ add_executable("kernel"
     # Kernel Implementation
     "src/main.cpp"
     "src/memory/bitmap_allocator.cpp"
+    "src/memory/free_list_allocator.cpp"
+    "src/memory/operators.cpp"
+    "src/memory.cpp"
 )
 
 target_include_directories("kernel" PRIVATE
diff --git a/kernel/include/kernel/memory.hpp b/kernel/include/kernel/memory.hpp
new file mode 100644
index 0000000..f09c519
--- /dev/null
+++ b/kernel/include/kernel/memory.hpp
@@ -0,0 +1,19 @@
+#ifndef TEACHOS_KERNEL_MEMORY_HPP
+#define TEACHOS_KERNEL_MEMORY_HPP
+
+#include "kapi/memory.hpp"
+
+#include "kernel/memory/heap_allocator.hpp"  // IWYU pragma: export
+
+namespace kernel::memory
+{
+
+  //! Get a handle to the global heap allocator.
+  auto get_heap_allocator() -> heap_allocator &;
+
+  //! Initialize the kernel heap.
+  auto init_heap(kapi::memory::linear_address base) -> void;
+
+}  // namespace kernel::memory
+
+#endif
\ No newline at end of file
diff --git a/kernel/include/kernel/memory/free_list_allocator.hpp b/kernel/include/kernel/memory/free_list_allocator.hpp
new file mode 100644
index 0000000..53a7b61
--- /dev/null
+++ b/kernel/include/kernel/memory/free_list_allocator.hpp
@@ -0,0 +1,80 @@
+#ifndef TEACHOS_KERNEL_MEMORY_FREE_LIST_ALLOCATOR_HPP
+#define TEACHOS_KERNEL_MEMORY_FREE_LIST_ALLOCATOR_HPP
+
+#include "kapi/memory.hpp"
+
+#include "kernel/memory/heap_allocator.hpp"
+
+#include <kstd/mutex>
+
+#include <cstddef>
+#include <new>
+
+namespace kernel::memory
+{
+
+  //! A simple, dynamically sized heap allocator.
+  //!
+  //! This allocator is designed to perform first-fit allocation using an intrusive doubly-linked list. The heap is
+  //! expanded as needed.
+  struct free_list_allocator final : heap_allocator
+  {
+    //! Construct a new free list allocator with the given base address.
+    explicit free_list_allocator(kapi::memory::linear_address base) noexcept;
+
+    free_list_allocator(free_list_allocator const &) = delete;
+    free_list_allocator(free_list_allocator &&) = delete;
+    auto operator=(free_list_allocator const &) -> free_list_allocator & = delete;
+    auto operator=(free_list_allocator &&) -> free_list_allocator & = delete;
+
+    //! Allocate a block of memory with the given alignment.
+    //!
+    //! @param size The size of the block to allocate
+    //! @param alignment The desired alignment of the allocated block
+    //! @return A pointer to the beginning of the block on success, @p nullptr otherwise.
+    [[nodiscard]] auto allocate(std::size_t size, std::align_val_t alignment) noexcept -> void * override;
+
+    //! Deallocate a block of memory previously allocated.
+    //!
+    //! @param block The pointer to the block to free.
+    auto deallocate(void * block) noexcept -> void override;
+
+  private:
+    struct block_header final
+    {
+      std::size_t size;
+      bool free;
+      block_header * next;
+      block_header * prev;
+    };
+
+    //! Try to expand the heap to accommodate the given size.
+    //!
+    //! @param delta The size to expand the heap by.
+    //! @return @p true if the heap was expanded, @p false otherwise.
+    auto expand(std::size_t delta) noexcept -> bool;
+
+    //! Split a given free block to accommodate and allocation.
+    //!
+    //! The newly split free block is inserted into the free list.
+    //!
+    //! @param block The block to split.
+    //! @param size The size of the allocation.
+    //! @param padding The amount of padding to apply.
+    auto split(block_header * block, std::size_t size, std::size_t padding) noexcept -> void;
+
+    //! Try to coalesce a given block with it's preceding and/or following block.
+    //!
+    //! @param block The block to coalesce.
+    auto coalesce(block_header * block) noexcept -> void;
+
+    kapi::memory::linear_address m_base;
+    kapi::memory::linear_address m_frontier;
+    block_header * m_free_list;
+
+    kstd::mutex m_lock;
+  };
+
+}  // namespace kernel::memory
+
+#endif
\ No newline at end of file
diff --git a/kernel/include/kernel/memory/heap_allocator.hpp b/kernel/include/kernel/memory/heap_allocator.hpp
new file mode 100644
index 0000000..bc771e3
--- /dev/null
+++ b/kernel/include/kernel/memory/heap_allocator.hpp
@@ -0,0 +1,32 @@
+#ifndef TEACHOS_KERNEL_MEMORY_HEAP_ALLOCATOR_HPP
+#define TEACHOS_KERNEL_MEMORY_HEAP_ALLOCATOR_HPP
+
+#include <kstd/mutex>
+
+#include <cstddef>
+#include <new>
+
+namespace kernel::memory
+{
+
+  //! The interface for all kernel heap allocators.
+  struct heap_allocator
+  {
+    virtual ~heap_allocator() = default;
+
+    //! Allocate a block of memory with the given alignment.
+    //!
+    //! @param size The size of the block to allocate
+    //! @param alignment The desired alignment of the allocated block
+    //! @return A pointer to the beginning of the block on success, @p nullptr otherwise.
+    [[nodiscard]] virtual auto allocate(std::size_t size, std::align_val_t alignment) noexcept -> void * = 0;
+
+    //! Deallocate a block of memory previously allocated.
+    //!
+    //! @param block The pointer to the block to free.
+    virtual auto deallocate(void * block) noexcept -> void = 0;
+  };
+
+}  // namespace kernel::memory
+
+#endif
\ No newline at end of file
diff --git a/kernel/kapi/memory.cpp b/kernel/kapi/memory.cpp
index d657527..767bcd7 100644
--- a/kernel/kapi/memory.cpp
+++ b/kernel/kapi/memory.cpp
@@ -101,9 +101,9 @@ namespace kapi::memory
     return get_frame_allocator().allocate_many(count);
   }
 
-  auto map(page page, frame frame) -> std::byte *
+  auto map(page page, frame frame, page_mapper::flags flags) -> std::byte *
   {
-    return active_page_mapper->map(page, frame, page_mapper::flags::empty);
+    return active_page_mapper->map(page, frame, flags);
   }
 
   auto unmap(page page) -> void
@@ -122,16 +122,15 @@ namespace kapi::memory
       system::panic("[OS:MEM] Not enough memory for bitmap allocator!");
     }
 
-    auto const base_address = pmm_metadata_base.raw();
     auto const flags = page_mapper::flags::writable | page_mapper::flags::supervisor_only;
 
     std::ranges::for_each(std::views::iota(0uz, bitmap_pages), [&](auto index) {
-      auto page = page::containing(linear_address{base_address + index * page::size});
+      auto page = page::containing(pmm_metadata_base + index * page::size);
       auto frame = memory::frame(bitmap_frames->first.number() + index);
       active_page_mapper->map(page, frame, flags);
     });
 
-    auto bitmap_ptr = std::bit_cast<std::uint64_t *>(base_address);
+    auto bitmap_ptr = std::bit_cast<std::uint64_t *>(pmm_metadata_base.raw());
     auto bitmap = std::span{bitmap_ptr, (bitmap_bytes + sizeof(std::uint64_t) - 1uz) / sizeof(std::uint64_t)};
 
     allocator.emplace(bitmap, frame_count);
diff --git a/kernel/src/main.cpp b/kernel/src/main.cpp
index d69903d..01bcbb0 100644
--- a/kernel/src/main.cpp
+++ b/kernel/src/main.cpp
@@ -2,6 +2,8 @@
 #include "kapi/memory.hpp"
 #include "kapi/system.hpp"
 
+#include "kernel/memory.hpp"
+
 #include <kstd/print>
 
 auto main() -> int
@@ -10,6 +12,7 @@ auto main() -> int
   kstd::println("[OS] IO subsystem initialized.");
 
   kapi::memory::init();
+  kernel::memory::init_heap(kapi::memory::heap_base);
   kstd::println("[OS] Memory subsystem initialized.");
   kapi::system::memory_initialized();
 
diff --git a/kernel/src/memory.cpp b/kernel/src/memory.cpp
new file mode 100644
index 0000000..14bedf1
--- /dev/null
+++ b/kernel/src/memory.cpp
@@ -0,0 +1,55 @@
+#include "kernel/memory.hpp"
+
+#include "kapi/memory.hpp"
+#include "kapi/system.hpp"
+
+#include "kernel/memory/free_list_allocator.hpp"
+#include "kernel/memory/heap_allocator.hpp"
+
+#include <atomic>
+#include <cstddef>
+#include <new>
+#include <optional>
+
+namespace kernel::memory
+{
+
+  namespace
+  {
+    struct null_allocator : heap_allocator
+    {
+      null_allocator static instance;
+
+      [[nodiscard]] auto allocate(std::size_t, std::align_val_t) noexcept -> void * override
+      {
+        return nullptr;
+      }
+
+      auto deallocate(void *) noexcept -> void override {}
+    };
+
+    constinit null_allocator null_allocator::instance{};
+
+    auto constinit active_heap_allocator = std::optional<heap_allocator *>{&null_allocator::instance};
+    auto constinit basic_allocator = std::optional<free_list_allocator>{};
+  }  // namespace
+
+  auto get_heap_allocator() -> heap_allocator &
+  {
+    return *active_heap_allocator.value();
+  }
+
+  auto init_heap(kapi::memory::linear_address base) -> void
+  {
+    auto static constinit is_initialized = std::atomic_flag{};
+
+    if (is_initialized.test_and_set())
+    {
+      kapi::system::panic("[OS] The heap has already been initialized.");
+    }
+
+    auto & instance = basic_allocator.emplace(base);
+    active_heap_allocator = &instance;
+  }
+
+}  // namespace kernel::memory
\ No newline at end of file
diff --git a/kernel/src/memory/free_list_allocator.cpp b/kernel/src/memory/free_list_allocator.cpp
new file mode 100644
index 0000000..255e643
--- /dev/null
+++ b/kernel/src/memory/free_list_allocator.cpp
@@ -0,0 +1,187 @@
+#include "kernel/memory/free_list_allocator.hpp"
+
+#include "kapi/memory.hpp"
+#include "kapi/system.hpp"
+
+#include "kernel/memory/heap_allocator.hpp"
+
+#include <kstd/mutex>
+
+#include <bit>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <new>
+
+namespace kernel::memory
+{
+
+  namespace
+  {
+    [[nodiscard]] constexpr auto align_up(std::uintptr_t value, std::size_t alignment) noexcept -> std::uintptr_t
+    {
+      auto const remainder = value % alignment;
+      return remainder == 0 ? value : value + alignment - remainder;
+    }
+  }  // namespace
+
+  free_list_allocator::free_list_allocator(kapi::memory::linear_address base) noexcept
+      : heap_allocator{}
+      , m_base{base}
+      , m_frontier{base}
+      , m_free_list{}
+      , m_lock{}
+  {}
+
+  auto free_list_allocator::allocate(std::size_t size, std::align_val_t alignment) noexcept -> void *
+  {
+    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_free_list;
+      while (current != nullptr)
+      {
+        if (current->free && current->size >= search_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;
+
+          split(current, size, padding);
+          current->free = false;
+
+          auto back_pointer = std::bit_cast<block_header **>(aligned_pointer - sizeof(block_header *));
+          *back_pointer = current;
+
+          return std::bit_cast<void *>(aligned_pointer);
+        }
+        current = current->next;
+      }
+
+      if (attempt == 0uz && !expand(search_size))
+      {
+        return nullptr;
+      }
+    }
+
+    return nullptr;
+  }
+
+  auto free_list_allocator::deallocate(void * ptr) noexcept -> void
+  {
+    if (!ptr)
+    {
+      return;
+    }
+
+    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 block = *back_pointer;
+
+    block->free = true;
+    coalesce(block);
+  }
+
+  auto free_list_allocator::expand(std::size_t size) noexcept -> bool
+  {
+    auto const total_required_size = size + sizeof(block_header);
+    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;
+
+    for (auto i = 0uz; i < frames_needed; ++i)
+    {
+      auto frame = kapi::memory::allocate_frame();
+      if (!frame)
+      {
+        kapi::system::panic("[OS:Heap] OOM when expanding heap.");
+        return false;
+      }
+
+      auto page = kapi::memory::page::containing(m_frontier + i * kapi::memory::page::size);
+      kapi::memory::map(page, *frame, flags);
+    }
+
+    auto block = std::bit_cast<block_header *>(m_frontier.raw());
+    // std::construct_at(block, frames_needed * kapi::memory::frame::size - sizeof(block_header), true, nullptr,
+    // nullptr);
+    block->size = frames_needed * kapi::memory::frame::size - sizeof(block_header);
+    block->free = true;
+    block->next = nullptr;
+    block->prev = nullptr;
+
+    m_frontier += frames_needed * kapi::memory::frame::size;
+
+    if (!m_free_list)
+    {
+      m_free_list = block;
+    }
+    else
+    {
+      auto current = m_free_list;
+      while (current->next)
+      {
+        current = current->next;
+      }
+      current->next = block;
+      block->prev = current;
+    }
+
+    coalesce(block);
+    return true;
+  }
+
+  auto free_list_allocator::coalesce(block_header * block) noexcept -> void
+  {
+    if (block->next && block->next->free)
+    {
+      auto follower = block->next;
+      block->size += follower->size + sizeof(block_header);
+      block->next = follower->next;
+      if (block->next)
+      {
+        block->next->prev = block;
+      }
+      std::destroy_at(follower);
+    }
+
+    if (block->prev && block->prev->free)
+    {
+      auto leader = block->prev;
+      leader->size += block->size + sizeof(block_header);
+      leader->next = block->next;
+      if (block->next)
+      {
+        block->next->prev = leader;
+      }
+      std::destroy_at(block);
+    }
+  }
+
+  auto free_list_allocator::split(block_header * block, std::size_t size, std::size_t padding) noexcept -> void
+  {
+    auto const allocation_size = size + padding;
+
+    if (block->size > allocation_size + sizeof(block_header) + 16uz)
+    {
+      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);
+
+      if (new_block->next)
+      {
+        new_block->next->prev = new_block;
+      }
+
+      block->size = allocation_size;
+      block->next = new_block;
+    }
+  }
+
+}  // namespace kernel::memory
\ No newline at end of file
diff --git a/kernel/src/memory/operators.cpp b/kernel/src/memory/operators.cpp
new file mode 100644
index 0000000..1422133
--- /dev/null
+++ b/kernel/src/memory/operators.cpp
@@ -0,0 +1,66 @@
+#include "kapi/system.hpp"
+
+#include "kernel/memory.hpp"
+
+#include <cstddef>
+#include <new>
+
+[[nodiscard]] auto operator new(std::size_t size) -> void *
+{
+  auto & allocator = kernel::memory::get_heap_allocator();
+  auto pointer = allocator.allocate(size, std::align_val_t{__STDCPP_DEFAULT_NEW_ALIGNMENT__});
+
+  if (pointer == nullptr)
+  {
+    kapi::system::panic("[OS:Heap] Out of memory!");
+  }
+
+  return pointer;
+}
+
+[[nodiscard]] auto operator new[](std::size_t size) -> void *
+{
+  return ::operator new(size);
+}
+
+[[nodiscard]] auto operator new(std::size_t size, std::align_val_t alignment) -> void *
+{
+  auto & allocator = kernel::memory::get_heap_allocator();
+  auto pointer = allocator.allocate(size, alignment);
+
+  if (pointer == nullptr)
+  {
+    kapi::system::panic("[OS:Heap] Out of memory!");
+  }
+
+  return pointer;
+}
+
+auto operator delete(void * pointer) noexcept -> void
+{
+  if (pointer != nullptr)
+  {
+    auto & allocator = kernel::memory::get_heap_allocator();
+    allocator.deallocate(pointer);
+  }
+}
+
+auto operator delete(void * pointer, [[maybe_unused]] std::size_t size) noexcept -> void
+{
+  ::operator delete(pointer);
+}
+
+auto operator delete[](void * pointer) noexcept -> void
+{
+  ::operator delete(pointer);
+}
+
+auto operator delete[](void * pointer, [[maybe_unused]] std::size_t size) noexcept -> void
+{
+  ::operator delete(pointer);
+}
+
+auto operator delete(void * pointer, [[maybe_unused]] std::align_val_t alignment) noexcept -> void
+{
+  ::operator delete(pointer);
+}
\ No newline at end of file