[Mammoth] Write a custom buddy allocator to replace libc malloc.

lib/libc/CMakeLists.txt

@@ -1,6 +1,5 @@
 
 add_library(c STATIC
-  src/malloc.cpp
   src/string.cpp
   )
 

lib/libc/include/stdlib.h

@@ -1,5 +0,0 @@
-#pragma once
-
-#include "stddef.h"
-
-void* malloc(size_t size);

lib/libc/src/malloc.cpp

@@ -1,46 +0,0 @@
-#include <zcall.h>
-#include <zglobal.h>
-
-#include "stdlib.h"
-
-namespace {
-class NaiveAllocator {
- public:
-  constexpr static uint64_t kSize = 0x10000;
-  NaiveAllocator() {}
-  bool is_init() { return next_addr_ != 0; }
-  void Init() {
-    uint64_t vmmo_cap;
-    uint64_t err = ZMemoryObjectCreate(kSize, &vmmo_cap);
-    if (err != 0) {
-      (void)ZProcessExit(err);
-    }
-    err = ZAddressSpaceMap(gSelfVmasCap, 0, vmmo_cap, 0, &next_addr_);
-    max_addr_ = next_addr_ + kSize;
-  }
-
-  void* Allocate(size_t size) {
-    uint64_t addr = next_addr_;
-    next_addr_ += size;
-    if (next_addr_ >= max_addr_) {
-      (void)ZProcessExit(0xBEEF);
-      return 0;
-    }
-    return reinterpret_cast<void*>(addr);
-  }
-
- private:
-  uint64_t next_addr_ = 0;
-  uint64_t max_addr_ = 0;
-};
-
-NaiveAllocator gAlloc;
-
-}  // namespace
-
-void* malloc(size_t size) {
-  if (!gAlloc.is_init()) {
-    gAlloc.Init();
-  }
-  return gAlloc.Allocate(size);
-}

lib/mammoth/util/new.cpp

@@ -1,8 +1,166 @@
+#include <glacier/container/intrusive_list.h>
+#include <glacier/string/str_format.h>
 #include <stdint.h>
-#include <stdlib.h>
+#include <zcall.h>
+#include <zglobal.h>
 
-[[nodiscard]] void* operator new(uint64_t size) { return malloc(size); }
+#include "util/debug.h"
-[[nodiscard]] void* operator new[](uint64_t size) { return malloc(size); }
+
+namespace {
+
+class PageAllocator {
+ public:
+  static uint64_t AllocatePagePair() {
+    uint64_t mem_cap;
+    check(ZMemoryObjectCreate(0x2000, &mem_cap));
+
+    uint64_t vaddr;
+    check(ZAddressSpaceMap(gSelfVmasCap, 0, mem_cap, /* align= */ 0x2000,
+                           &vaddr));
+
+    // The address space mapping will keep this memory capability alive.
+    check(ZCapRelease(mem_cap));
+
+    return vaddr;
+  }
+};
+
+struct BuddySlot {
+  BuddySlot* next;
+  BuddySlot* prev;
+  uint64_t size;
+};
+
+uint64_t NeededSize(uint64_t size) {
+  uint64_t needed = size + sizeof(BuddySlot);
+  // Start at 32 because sizeof(BuddySlot) is already 24;
+  uint64_t pow2 = 32;
+
+  while (pow2 < needed) {
+    pow2 <<= 1;
+  }
+
+  return pow2;
+}
+
+class BuddyAllocator {
+ public:
+  BuddyAllocator() {}
+
+  void Init() {
+    free_front_ = nullptr;
+    AddPage();
+    check(ZMutexCreate(&mutex_cap_));
+  }
+
+  void* Allocate(uint64_t size) {
+    check(ZMutexLock(mutex_cap_));
+    if (size > (0x2000 - sizeof(BuddySlot))) {
+      crash("Can't allocate greater than one page", glcr::UNIMPLEMENTED);
+    }
+    if (free_front_ == nullptr) {
+      AddPage();
+    }
+    BuddySlot* slot = free_front_;
+    uint64_t needed = NeededSize(size);
+    BuddySlot* best_fit = nullptr;
+
+    while (slot != nullptr) {
+      bool fits = slot->size >= needed;
+      bool better = best_fit == nullptr || slot->size < best_fit->size;
+      if (fits && better) {
+        best_fit = slot;
+      }
+      slot = slot->next;
+    }
+
+    if (best_fit == nullptr) {
+      AddPage();
+      best_fit = free_front_;
+    }
+
+    while (best_fit->size > needed) {
+      best_fit = Split(best_fit);
+    }
+
+    Remove(best_fit);
+    // TODO: We may need to align the pointer here.
+    void* ptr = reinterpret_cast<uint8_t*>(best_fit) + sizeof(BuddySlot);
+    check(ZMutexRelease(mutex_cap_));
+    return ptr;
+  }
+
+ private:
+  BuddySlot* free_front_ = nullptr;
+  z_cap_t mutex_cap_ = 0;
+
+  void AddPage() {
+    uint64_t vaddr = PageAllocator::AllocatePagePair();
+    BuddySlot* slot = reinterpret_cast<BuddySlot*>(vaddr);
+    slot->prev = nullptr;
+    slot->next = free_front_;
+    free_front_ = slot;
+    slot->size = 0x2000;
+  }
+
+  BuddySlot* Split(BuddySlot* slot) {
+    if (slot->size <= 32) {
+      crash("Splitting smallest buddy chunk", glcr::INTERNAL);
+    }
+
+    slot->size /= 2;
+    BuddySlot* new_slot = reinterpret_cast<BuddySlot*>(
+        reinterpret_cast<uint64_t>(slot) ^ slot->size);
+    new_slot->size = slot->size;
+    new_slot->next = slot->next;
+    new_slot->prev = slot;
+    if (slot->next) {
+      slot->next->prev = new_slot;
+    }
+    slot->next = new_slot;
+    return slot;
+  }
+
+  void Remove(BuddySlot* slot) {
+    if (slot->prev) {
+      slot->prev->next = slot->next;
+    }
+    if (slot->next) {
+      slot->next->prev = slot->prev;
+    }
+
+    if (free_front_ == slot) {
+      free_front_ = slot->next;
+    }
+
+    slot->next = nullptr;
+    slot->prev = nullptr;
+  }
+};
+
+BuddyAllocator gAllocator;
+bool gHasInit = false;
+
+// FIXME: Race condition.
+void* Allocate(uint64_t size) {
+  if (!gHasInit) {
+    // Call Init since we don't call global constructors yet.
+    gAllocator.Init();
+    gHasInit = true;
+  }
+  void* ptr = gAllocator.Allocate(size);
+
+  char buffer[64];
+  // glcr::FixedStringBuilder builder(buffer, 64);
+  // glcr::StrFormatIntoBuffer(builder, "Allocated {x}", (uint64_t)ptr);
+  // dbgln(builder.operator glcr::StringView());
+  return ptr;
+}
+
+}  // namespace
+
+[[nodiscard]] void* operator new(uint64_t size) { return Allocate(size); }
+[[nodiscard]] void* operator new[](uint64_t size) { return Allocate(size); }
 
 void operator delete(void*, uint64_t) {}
 void operator delete[](void*) {}

sys/yellowstone/yellowstone_server.cpp

@@ -6,7 +6,6 @@
 #include <mammoth/util/debug.h>
 #include <mammoth/util/init.h>
 #include <mammoth/util/memory_region.h>
-#include <stdlib.h>
 #include <zcall.h>
 
 #include "hw/gpt.h"