From eb04242a59aa45c3862cc453cf841ee02ce9186e Mon Sep 17 00:00:00 2001
From: Drew Galbraith <dgalbraith33@gmail.com>
Date: Wed, 7 Jun 2023 10:01:22 -0700
Subject: [PATCH] Add a debug mode for the KernelHeap.

---
 zion/interrupt/interrupt.cpp |  4 ++++
 zion/memory/kernel_heap.cpp  | 33 +++++++++++++++++++++++++++++++++
 zion/memory/kernel_heap.h    | 19 +++++++++++++++++++
 zion/zion.cpp                |  3 +++
 4 files changed, 59 insertions(+)

diff --git a/zion/interrupt/interrupt.cpp b/zion/interrupt/interrupt.cpp
index 4f8d0a4..da41838 100644
--- a/zion/interrupt/interrupt.cpp
+++ b/zion/interrupt/interrupt.cpp
@@ -4,6 +4,7 @@
 
 #include "common/port.h"
 #include "debug/debug.h"
+#include "memory/kernel_heap.h"
 #include "scheduler/scheduler.h"
 
 #define IDT_INTERRUPT_GATE 0x8E
@@ -132,6 +133,9 @@ extern "C" void isr_timer();
 extern "C" void interrupt_timer(InterruptFrame*) {
   cnt++;
   if (cnt % 20 == 0) {
+    if (cnt == 20) {
+      KernelHeap::DumpDistribution();
+    }
     dbgln("timer: %us", cnt * 50 / 1000);
   }
   outb(PIC1_COMMAND, PIC_EOI);
diff --git a/zion/memory/kernel_heap.cpp b/zion/memory/kernel_heap.cpp
index f9babc8..910a112 100644
--- a/zion/memory/kernel_heap.cpp
+++ b/zion/memory/kernel_heap.cpp
@@ -3,6 +3,8 @@
 #include "debug/debug.h"
 #include "memory/paging_util.h"
 
+#define K_HEAP_DEBUG 0
+
 namespace {
 
 static KernelHeap* gKernelHeap = nullptr;
@@ -24,12 +26,43 @@ void* KernelHeap::Allocate(uint64_t size) {
   if (next_addr_ + size >= upper_bound_) {
     panic("Kernel Heap Overrun");
   }
+#if K_HEAP_DEBUG
+  RecordSize(size);
+#endif
   EnsureResident(next_addr_, size);
   uint64_t address = next_addr_;
   next_addr_ += size;
   return reinterpret_cast<void*>(address);
 }
 
+void KernelHeap::DumpDistribution() {
+#if K_HEAP_DEBUG
+  uint64_t* distributions = gKernelHeap->distributions;
+  dbgln("<=4B:   %u", distributions[0]);
+  dbgln("<=8B:   %u", distributions[1]);
+  dbgln("<=16B:  %u", distributions[2]);
+  dbgln("<=32B:  %u", distributions[3]);
+  dbgln("<=64B:  %u", distributions[4]);
+  dbgln("<=128B: %u", distributions[5]);
+  dbgln("<=256B: %u", distributions[6]);
+  dbgln("<=512B: %u", distributions[7]);
+  dbgln("<=1KiB: %u", distributions[8]);
+  dbgln("<=2KiB: %u", distributions[9]);
+  dbgln("<=4KiB: %u", distributions[10]);
+  dbgln("> 4KiB: %u", distributions[11]);
+#endif
+}
+
+void KernelHeap::RecordSize(uint64_t size) {
+  size >>= 3;
+  uint64_t index = 0;
+  while (size && index < 11) {
+    size >>= 1;
+    index++;
+  }
+  distributions[index]++;
+}
+
 void* operator new(uint64_t size) { return GetKernelHeap().Allocate(size); }
 void* operator new[](uint64_t size) { return GetKernelHeap().Allocate(size); }
 
diff --git a/zion/memory/kernel_heap.h b/zion/memory/kernel_heap.h
index 456b387..ae8f772 100644
--- a/zion/memory/kernel_heap.h
+++ b/zion/memory/kernel_heap.h
@@ -8,7 +8,26 @@ class KernelHeap {
 
   void* Allocate(uint64_t size);
 
+  static void DumpDistribution();
+
  private:
   uint64_t next_addr_;
   uint64_t upper_bound_;
+
+  // Distribution collection for the purpose of investigating a slab allocator.
+  // 0: 0-4B
+  // 1: 4B-8B
+  // 2: 8B-16B
+  // 3: 16B-32B
+  // 4: 32B-64B
+  // 5: 64B-128B
+  // 6: 128B-256B
+  // 7: 256B-512B
+  // 8: 512B-1KiB
+  // 9: 1KiB-2KiB
+  // 10: 2KiB-4KiB
+  // 11: 4KiB+
+  uint64_t distributions[12];
+
+  void RecordSize(uint64_t size);
 };
diff --git a/zion/zion.cpp b/zion/zion.cpp
index 5cc2a0e..abdc4a8 100644
--- a/zion/zion.cpp
+++ b/zion/zion.cpp
@@ -40,6 +40,9 @@ extern "C" void zion() {
   dbgln("[boot] Loading sys init program.");
   LoadInitProgram();
 
+  dbgln("[boot] Allocs during boot:");
+  heap.DumpDistribution();
+
   dbgln("[boot] Init finished, yielding.");
   gScheduler->Enable();
   gScheduler->Yield();