From d2bc8f2bf82964d2ca8a178a50f4a79abcbd8b76 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= <loic@gregoire.dev>
Date: Fri, 3 Apr 2026 21:16:53 -0400
Subject: [PATCH] feat(mm/phys): implement alloc and free

---
 Kconfig               |   2 +-
 src/mm/amd64/impl.zig |   2 +-
 src/mm/phys.zig       | 176 +++++++++++++++++++++++++++++++++++++++---
 src/mm/pmap.zig       |   2 +-
 src/mm/radix_pmap.zig |   2 +-
 src/mm/root.zig       |  16 +++-
 6 files changed, 185 insertions(+), 15 deletions(-)

diff --git a/Kconfig b/Kconfig
index 196b31e..3590454 100644
--- a/Kconfig
+++ b/Kconfig
@@ -1,4 +1,4 @@
-mainmenu "Bach Configuration"
+mainmenu "Zag Configuration"
 
 
 choice
diff --git a/src/mm/amd64/impl.zig b/src/mm/amd64/impl.zig
index bf46a75..ee2fe03 100644
--- a/src/mm/amd64/impl.zig
+++ b/src/mm/amd64/impl.zig
@@ -98,7 +98,7 @@ pub fn init_kernel() void {
     const higher_half_start = phys_to_virt(0);
     const higher_half_size: usize = 256 * (1 << 39);
 
-    mi.kernel_pmap.impl.root_pa = mi.phys.alloc_page();
+    mi.kernel_pmap.impl.root_pa = mi.phys.alloc();
     const table_ptr: [*]u64 = @ptrFromInt(mm.p2v(mi.kernel_pmap.impl.root_pa));
     @memset(table_ptr[0..512], 0);
 
diff --git a/src/mm/phys.zig b/src/mm/phys.zig
index 6fdc04d..cbbdc7c 100644
--- a/src/mm/phys.zig
+++ b/src/mm/phys.zig
@@ -1,17 +1,95 @@
+//! Physical memory allocator:
+//! - Uses a per-CPU magazined allocator with a global Treiber stack as the fallback.
+//! Each CPU has 4 batches of 64 pages each for fast allocation and deallocation without contention.
+//!  When a CPU's active batch is exhausted, it tries to get a new batch from its local depot.
+//!  If the depot is empty, it pops a batch from the global stack.
+//! When freeing pages, they are added to the active batch until it reaches capacity,
+//!  at which point the full batch is moved to the local depot.
+//!  If the depot is full when trying to add a new batch, one batch is pushed back to the global stack to make room.
 const b = @import("base");
 const std = @import("std");
 const rtl = @import("rtl");
 const pl = b.pl;
 const mm = b.mm;
 const mi = mm.private;
+const ke = b.ke;
 
+// === Early state ===
+var bootstrapped = false;
 var early_alloc_entry_idx: usize = 0;
 var memory_map: *pl.BootInfo.MemMap = undefined;
 var early_allocs: usize = 0;
 
-pub var bootstrapped = false;
+// === Memory allocator state ===
+/// Head for the global treiber stack.
+const GlobalHead = packed struct(u64) {
+    pfn: mm.Pfn,
+    // A 32-bit ABA tag is more than enough for this,
+    // on 32-bit machines which do not support cmpxchg8b, we could use 12 bits instead.
+    tag: u32,
+};
 
-pub fn early_alloc() usize {
+/// Per-CPU state.
+const PerCpu = struct {
+    depot: [4]*mm.Page,
+    depot_count: u8,
+    active_batch: ?*mm.Page,
+    active_count: u8,
+};
+
+const percpu = ke.CpuLocal(PerCpu, undefined);
+
+const null_pfn: mm.Pfn = std.math.maxInt(mm.Pfn);
+
+// TODO: make this configurable (or dynamic).
+const batch_size: u8 = 64;
+
+var batch_pool: std.atomic.Value(GlobalHead) = .init(.{
+    .pfn = null_pfn,
+    .tag = 0,
+});
+
+/// Push a batch of free pages to the global pool.
+fn push_to_pool(batch: *mm.Page) void {
+    var old_head = batch_pool.load(.acquire);
+
+    while (true) {
+        batch.batch_next = old_head.pfn;
+
+        const new_head = GlobalHead{
+            .pfn = mm.struct_page_to_pfn(batch),
+            .tag = old_head.tag + 1,
+        };
+
+        old_head = batch_pool.cmpxchgWeak(old_head, new_head, .acq_rel, .acquire) orelse break;
+    }
+}
+
+/// Pop a batch of free pages from the global pool. Returns null if the pool is empty.
+fn pop_from_pool() ?*mm.Page {
+    var old_head = batch_pool.load(.acquire);
+
+    while (true) {
+        const batch_pfn = old_head.pfn;
+
+        if (batch_pfn == null_pfn) {
+            return null;
+        }
+
+        const batch = mm.pfn_to_struct_page(batch_pfn);
+
+        const new_head = GlobalHead{
+            .pfn = batch.batch_next,
+            .tag = old_head.tag + 1,
+        };
+
+        old_head = batch_pool.cmpxchgWeak(old_head, new_head, .acq_rel, .acquire) orelse return batch;
+    }
+
+    return null;
+}
+
+fn early_alloc() usize {
     while (early_alloc_entry_idx < memory_map.entry_count) {
         const entry = &memory_map.entries[early_alloc_entry_idx];
 
@@ -32,11 +110,92 @@ pub fn early_alloc() usize {
     @panic("mm/phys: Out of memory for early allocation");
 }
 
-pub fn alloc_page() usize {
+/// Allocate a page of physical memory.
+pub fn alloc() b.PAddr {
     if (!bootstrapped) {
         return early_alloc();
     }
-    @panic("not implemented");
+
+    const ipl = ke.ipl.raise(.High);
+    const cpu = percpu.local();
+
+    defer ke.ipl.lower(ipl);
+
+    if (cpu.active_count == 0) {
+        // Our batch is empty, get another one.
+        if (cpu.depot_count > 0) {
+            // Get a new batch from the local depot.
+            cpu.depot_count -= 1;
+            cpu.active_batch = cpu.depot[cpu.depot_count];
+            cpu.active_count = batch_size;
+        } else {
+            // Get one from the global pool.
+            const new_batch = pop_from_pool();
+
+            if (new_batch) |new| {
+                cpu.active_batch = new;
+                cpu.active_count = new.batch_count;
+            } else {
+                // TODO: call an IPI here for the other CPUs to fill the global pool before giving up.
+                @panic("OOM");
+            }
+        }
+    }
+
+    const batch = cpu.active_batch.?;
+
+    if (cpu.active_count > 1) {
+        // Pop the next page from the batch and set it as the new head.
+        const next_page = mm.pfn_to_struct_page(batch.next_pfn);
+        cpu.active_batch = next_page;
+    } else {
+        cpu.active_batch = null;
+    }
+
+    cpu.active_count -= 1;
+
+    return mm.pfn_to_page(mm.struct_page_to_pfn(batch));
+}
+
+/// Free a page of physical memory.
+pub fn free(addr: b.PAddr) void {
+    const ipl = ke.ipl.raise(.High);
+    defer ke.ipl.lower(ipl);
+
+    const cpu = percpu.local();
+
+    const page = mm.pfn_to_struct_page(mm.page_to_pfn(addr));
+
+    if (cpu.active_batch) |curr| {
+        // Add the page to the current batch.
+        page.next_pfn = mm.struct_page_to_pfn(curr);
+    } else {
+        page.next_pfn = null_pfn;
+    }
+
+    cpu.active_batch = page;
+    cpu.active_count += 1;
+
+    if (cpu.active_count == batch_size) {
+        // The batch is full, move it to the depot.
+        page.batch_count = @intCast(cpu.active_count);
+
+        if (cpu.depot_count == cpu.depot.len) {
+            // The depot is full, push one batch back to the global pool to make room.
+            push_to_pool(cpu.depot[0]);
+
+            cpu.depot[0] = cpu.depot[1];
+            cpu.depot[1] = cpu.depot[2];
+            cpu.depot[2] = cpu.depot[3];
+            cpu.depot_count -= 1;
+        }
+
+        cpu.depot[cpu.depot_count] = cpu.active_batch.?;
+        cpu.depot_count += 1;
+
+        cpu.active_batch = null;
+        cpu.active_count = 0;
+    }
 }
 
 pub fn init(boot_info: *pl.BootInfo) linksection(b.init) void {
@@ -91,8 +250,6 @@ pub fn init(boot_info: *pl.BootInfo) linksection(b.init) void {
 }
 
 pub fn init_pfndb() void {
-    const pfndb: [*]mm.Page = @ptrFromInt(mi.impl.pfndb_base);
-
     for (0..memory_map.entry_count) |i| {
         const entry = memory_map.entries[i];
 
@@ -100,13 +257,12 @@ pub fn init_pfndb() void {
             continue;
         }
 
-        const start_pfn: usize = mm.page_to_pfn(entry.base);
         const npages = entry.size / mm.page_size;
 
         for (0..npages) |j| {
-            const pfn = start_pfn + j;
-            const page = &pfndb[pfn];
-            page.batch_count = 0;
+            free(entry.base + (j * mm.page_size));
         }
     }
+
+    bootstrapped = true;
 }
diff --git a/src/mm/pmap.zig b/src/mm/pmap.zig
index fe05b02..980dddd 100644
--- a/src/mm/pmap.zig
+++ b/src/mm/pmap.zig
@@ -55,7 +55,7 @@ const AllocatingSource = struct {
     flags: mm.MapFlags,
 
     pub fn next(self: *const AllocatingSource, _: b.VAddr) MapItem {
-        const pa = mi.phys.alloc_page();
+        const pa = mi.phys.alloc();
         return .{
             .pa = pa,
             .len = mm.page_size,
diff --git a/src/mm/radix_pmap.zig b/src/mm/radix_pmap.zig
index 16467ac..95db64a 100644
--- a/src/mm/radix_pmap.zig
+++ b/src/mm/radix_pmap.zig
@@ -133,7 +133,7 @@ pub fn RadixPmap(comptime Impl: type) type {
                             return error.PageNotMapped;
                         }
 
-                        const new_table_pa = mi.phys.alloc_page();
+                        const new_table_pa = mi.phys.alloc();
                         const table_ptr: [*]Impl.Pte = @ptrFromInt(mm.p2v(new_table_pa));
                         @memset(table_ptr[0..entries_per_table], std.mem.zeroes(Impl.Pte));
 
diff --git a/src/mm/root.zig b/src/mm/root.zig
index 7fe944d..821a377 100644
--- a/src/mm/root.zig
+++ b/src/mm/root.zig
@@ -28,9 +28,23 @@ pub const MapFlags = packed struct {
     cache_disable: bool = false,
 };
 
-pub fn page_to_pfn(addr: usize) Pfn {
+pub inline fn page_to_pfn(addr: usize) Pfn {
     return @intCast(addr >> 12);
 }
 
+pub inline fn pfn_to_page(pfn: Pfn) usize {
+    return @intCast(pfn << 12);
+}
+
+pub inline fn struct_page_to_pfn(page: *Page) Pfn {
+    return @intCast((@intFromPtr(page) - private.impl.pfndb_base) / @sizeOf(Page));
+}
+
+pub inline fn pfn_to_struct_page(pfn: Pfn) *Page {
+    return @ptrFromInt(private.impl.pfndb_base + (@as(usize, @intCast(pfn)) * @sizeOf(Page)));
+}
+
+pub var pfndb: [*]Page = @ptrFromInt(private.impl.pfndb_base);
+
 pub const p2v = private.impl.phys_to_virt;
 pub const v2p = private.impl.virt_to_phys;