feat: Enhance Kata state management with new states and deferred cleanup for page tables

10 files changed, 124 insertions, 34 deletions

diff --git a/mirai/drivers/serial/serial.zig b/mirai/drivers/serial/serial.zig
index 9d5770f..07d3af7 100644
--- a/mirai/drivers/serial/serial.zig
+++ b/mirai/drivers/serial/serial.zig
@@ -64,9 +64,20 @@ pub fn printf(comptime fmt: []const u8, args: anytype) void {
             print_arg(arg);
             arg_index += 1;
             i += 2;
-        } else if (fmt[i] == '{' and i + 2 < fmt.len and fmt[i + 1] == 'x' and fmt[i + 2] == '}') {
+        } else if (fmt[i] == '{' and i + 2 < fmt.len and fmt[i + 2] == '}') {
+            const spec = fmt[i + 1];
             const arg = args[arg_index];
-            print_arg_hex(arg);
+            if (spec == 'x') {
+                print_arg_hex(arg);
+            } else if (spec == 'd') {
+                print_arg_decimal(arg);
+            } else if (spec == 's') {
+                print_arg_string(arg);
+            } else {
+                write('{');
+                write(spec);
+                write('}');
+            }
             arg_index += 1;
             i += 3;
         } else {
@@ -89,6 +100,44 @@ fn print_arg(arg: anytype) void {
     }
 }
 
+fn print_arg_decimal(arg: anytype) void {
+    const T = @TypeOf(arg);
+    if (@typeInfo(T) == .int or @typeInfo(T) == .comptime_int) {
+        print_decimal(arg);
+    } else {
+        print("?");
+    }
+}
+
+fn print_arg_string(arg: anytype) void {
+    const T = @TypeOf(arg);
+    const info = @typeInfo(T);
+    if (T == []const u8) {
+        print(arg);
+    } else if (info == .pointer) {
+        const child = info.pointer.child;
+        const child_info = @typeInfo(child);
+        if (child == u8) {
+            // [*]const u8 - null-terminated string pointer
+            var p = arg;
+            while (p[0] != 0) : (p += 1) {
+                write(p[0]);
+            }
+        } else if (child_info == .array and child_info.array.child == u8) {
+            // *const [N]u8 or *const [N:0]u8 - pointer to string literal
+            const slice = arg.*;
+            for (slice) |c| {
+                if (c == 0) break;
+                write(c);
+            }
+        } else {
+            print("?ptr");
+        }
+    } else {
+        print("?type");
+    }
+}
+
 fn print_arg_hex(arg: anytype) void {
     const T = @TypeOf(arg);
     if (T == u8) {
diff --git a/mirai/invocations/kata/wait.zig b/mirai/invocations/kata/wait.zig
index 1d6afa5..e46b08b 100644
--- a/mirai/invocations/kata/wait.zig
+++ b/mirai/invocations/kata/wait.zig
@@ -17,7 +17,7 @@ pub fn invoke(ctx: *handler.InvocationContext) void {
 
     const current = sensei.get_current_kata() orelse return result.set_error(ctx);
 
-    current.state = .Waiting;
+    current.state = .Stalled;
     current.waiting_for = target_id;
 
     current.context.rax = 0;
diff --git a/mirai/invocations/kata/yield.zig b/mirai/invocations/kata/yield.zig
index 0c2715a..bc4028a 100644
--- a/mirai/invocations/kata/yield.zig
+++ b/mirai/invocations/kata/yield.zig
@@ -8,7 +8,7 @@ const sensei = @import("../../kata/sensei/sensei.zig");
 pub fn invoke(ctx: *handler.InvocationContext) void {
     const kata = sensei.get_current_kata() orelse return result.set_error(ctx);
 
-    kata.state = kata_mod.State.Ready;
+    kata.state = kata_mod.State.Alive;
 
     if (!sensei.is_in_queue(kata)) {
         sensei.enqueue_kata(kata);
diff --git a/mirai/kata/memory.zig b/mirai/kata/memory.zig
index ac16d76..5639bf5 100644
--- a/mirai/kata/memory.zig
+++ b/mirai/kata/memory.zig
@@ -13,6 +13,39 @@ const PAGE_SIZE = memory_const.PAGE_SIZE;
 const KERNEL_VMALLOC_START: u64 = 0xFFFFFF8000000000;
 var next_vmalloc_addr: u64 = KERNEL_VMALLOC_START;
 
+// Deferred page table destruction queue
+const MAX_DEFERRED: usize = 16;
+var deferred_page_tables: [MAX_DEFERRED]u64 = [_]u64{0} ** MAX_DEFERRED;
+var deferred_count: usize = 0;
+
+/// Queue a page table for deferred destruction
+fn queue_deferred_destroy(page_table: u64) void {
+    if (deferred_count < MAX_DEFERRED) {
+        deferred_page_tables[deferred_count] = page_table;
+        deferred_count += 1;
+    }
+    // If queue full, leak the page table (shouldn't happen in practice)
+}
+
+/// Process deferred page table destructions.
+/// Safe to call when CR3 points to a page table we want to keep.
+pub fn process_deferred_cleanup(exclude_pt: u64) void {
+    const current_cr3 = asm_memory.read_page_table_base();
+    var write_idx: usize = 0;
+
+    for (0..deferred_count) |i| {
+        const pt = deferred_page_tables[i];
+        if (pt != 0 and pt != current_cr3 and pt != exclude_pt) {
+            paging.destroy_page_table(pt);
+        } else if (pt != 0) {
+            // Keep in queue
+            deferred_page_tables[write_idx] = pt;
+            write_idx += 1;
+        }
+    }
+    deferred_count = write_idx;
+}
+
 pub const VirtualBuffer = struct {
     data: []u8,
     virt_base: u64,
@@ -220,12 +253,13 @@ pub fn load_segment(
 
 pub fn cleanup(kata: *types.Kata) void {
     if (kata.page_table != 0) {
-        // CRITICAL: If we're currently using this kata's page table,
-        // we CANNOT destroy it - we'd be freeing pages while CR3 points to them.
-        // Skip destruction (memory leak) to prevent crash.
         const current_cr3 = asm_memory.read_page_table_base();
         if (current_cr3 != kata.page_table) {
+            // Safe to destroy immediately
             paging.destroy_page_table(kata.page_table);
+        } else {
+            // Queue for deferred destruction
+            queue_deferred_destroy(kata.page_table);
         }
         kata.page_table = 0;
     }
diff --git a/mirai/kata/pool.zig b/mirai/kata/pool.zig
index 3fff3e0..030909a 100644
--- a/mirai/kata/pool.zig
+++ b/mirai/kata/pool.zig
@@ -5,7 +5,9 @@ const attachment_const = @import("../common/constants/attachment.zig");
 const heap = @import("../memory/heap.zig");
 const kata_const = @import("../common/constants/kata.zig");
 const kata_limits = @import("../common/limits/kata.zig");
+const memory = @import("memory.zig");
 const types = @import("types.zig");
+const waker = @import("sensei/waker.zig");
 
 pub var pool: [kata_limits.MAX_KATAS]types.Kata = undefined;
 pub var used: [kata_limits.MAX_KATAS]bool = [_]bool{false} ** kata_limits.MAX_KATAS;
@@ -28,7 +30,7 @@ pub fn create() !*types.Kata {
 
             kata.* = create_empty();
             kata.id = kata_id;
-            kata.state = .Ready;
+            kata.state = .Alive;
 
             kata.current_location[0] = '/';
 
@@ -73,9 +75,6 @@ pub fn get(id: u32) ?*types.Kata {
 }
 
 pub fn dissolve(kata_id: u32) void {
-    const waker = @import("sensei/waker.zig");
-    const memory = @import("memory.zig");
-
     for (&pool, 0..) |*kata, i| {
         if (used[i] and kata.id == kata_id) {
             for (&kata.attachments) |*slot| {
diff --git a/mirai/kata/sensei/queue.zig b/mirai/kata/sensei/queue.zig
index d39d98d..943d70c 100644
--- a/mirai/kata/sensei/queue.zig
+++ b/mirai/kata/sensei/queue.zig
@@ -9,7 +9,7 @@ pub fn get_head() ?*types.Kata {
 }
 
 pub fn enqueue(kata: *types.Kata) void {
-    kata.state = .Ready;
+    kata.state = .Alive;
 
     if (head == null) {
         head = kata;
diff --git a/mirai/kata/sensei/sensei.zig b/mirai/kata/sensei/sensei.zig
index 2e801b3..7985ad9 100644
--- a/mirai/kata/sensei/sensei.zig
+++ b/mirai/kata/sensei/sensei.zig
@@ -41,20 +41,20 @@ pub fn schedule() void {
 
     if (next == current and next != null) {
         queue.dequeue(next.?);
-        next.?.state = .Running;
+        next.?.state = .Flowing;
         return;
     }
 
     if (current) |curr| {
-        if (curr.state == .Running) {
-            curr.state = .Ready;
+        if (curr.state == .Flowing) {
+            curr.state = .Alive;
             queue.enqueue(curr);
         }
     }
 
     if (next) |n| {
         queue.dequeue(n);
-        n.state = .Running;
+        n.state = .Flowing;
         n.last_run = tick_count;
         current = n;
         shift.to_kata(n);
@@ -104,11 +104,11 @@ fn pick_next() ?*types.Kata {
     var fallback: ?*types.Kata = null;
 
     while (curr) |kata| {
-        if (kata.state == .Ready and kata.vruntime < local_min) {
+        if (kata.state == .Alive and kata.vruntime < local_min) {
             local_min = kata.vruntime;
             chosen = kata;
         }
-        if (kata.state == .Running) {
+        if (kata.state == .Flowing) {
             fallback = kata;
         }
         curr = kata.next;
@@ -120,7 +120,7 @@ fn pick_next() ?*types.Kata {
     }
 
     if (current) |c| {
-        if (c.state == .Running) return c;
+        if (c.state == .Flowing) return c;
     }
 
     if (fallback) |k| return k;
diff --git a/mirai/kata/sensei/waker.zig b/mirai/kata/sensei/waker.zig
index fc7f15a..c4334e3 100644
--- a/mirai/kata/sensei/waker.zig
+++ b/mirai/kata/sensei/waker.zig
@@ -11,11 +11,11 @@ pub fn wake_all_waiting() void {
         if (!pool.used[i]) continue;
 
         const kata = &pool.pool[i];
-        if (kata.state != .Waiting) continue;
+        if (kata.state != .Stalled) continue;
 
         const target = pool.get(kata.waiting_for);
         if (target == null or target.?.state == .Dissolved) {
-            kata.state = .Ready;
+            kata.state = .Alive;
             queue.enqueue(kata);
             kata.waiting_for = 0;
         }
@@ -27,8 +27,8 @@ pub fn wake_waiting(target_id: u32) void {
         if (!pool.used[i]) continue;
 
         const kata = &pool.pool[i];
-        if (kata.state == .Waiting and kata.waiting_for == target_id) {
-            kata.state = .Ready;
+        if (kata.state == .Stalled and kata.waiting_for == target_id) {
+            kata.state = .Alive;
             queue.enqueue(kata);
             kata.waiting_for = 0;
         }
@@ -42,8 +42,8 @@ pub fn wake_blocked() void {
         if (!pool.used[i]) continue;
 
         const kata = &pool.pool[i];
-        if (kata.state == .Blocked) {
-            kata.state = .Ready;
+        if (kata.state == .Frozen) {
+            kata.state = .Alive;
             queue.enqueue(kata);
             return;
         }
@@ -55,8 +55,8 @@ pub fn wake_one_blocked() void {
         if (!pool.used[i]) continue;
 
         const kata = &pool.pool[i];
-        if (kata.state == .Blocked) {
-            kata.state = .Ready;
+        if (kata.state == .Frozen) {
+            kata.state = .Alive;
             queue.enqueue(kata);
             return;
         }
@@ -66,8 +66,8 @@ pub fn wake_one_blocked() void {
 pub fn wake(kata_id: u32) void {
     const kata = pool.get(kata_id) orelse return;
 
-    if (kata.state == .Blocked) {
-        kata.state = .Ready;
+    if (kata.state == .Frozen) {
+        kata.state = .Alive;
         queue.enqueue(kata);
     }
 }
diff --git a/mirai/kata/shift.zig b/mirai/kata/shift.zig
index 3e45379..3ad64be 100644
--- a/mirai/kata/shift.zig
+++ b/mirai/kata/shift.zig
@@ -2,12 +2,17 @@
 
 const context = @import("../asm/context.zig");
 const gdt = @import("../boot/gdt/gdt.zig");
+const kata_memory = @import("memory.zig");
 const tss = @import("../boot/tss/tss.zig");
 const types = @import("types.zig");
 
 var current_context: ?*types.Context = null;
 
 pub fn to_kata(kata: *types.Kata) void {
+    // Process any deferred page table cleanups before switching
+    // Exclude the target kata's page table from cleanup
+    kata_memory.process_deferred_cleanup(kata.page_table);
+
     tss.set_kernel_stack(kata.stack_top);
     current_context = &kata.context;
 
diff --git a/mirai/kata/types.zig b/mirai/kata/types.zig
index efbbe8d..85f653f 100644
--- a/mirai/kata/types.zig
+++ b/mirai/kata/types.zig
@@ -5,11 +5,14 @@ const kata_const = @import("../common/constants/kata.zig");
 const kata_limits = @import("../common/limits/kata.zig");
 
 pub const State = enum {
-    Ready,
-    Running,
-    Waiting,
-    Blocked,
-    Dissolved,
+    Born, // Just created, being initialized
+    Alive, // Ready to run, waiting for scheduler
+    Flowing, // Currently executing
+    Stalled, // Waiting for child/event
+    Frozen, // Blocked on I/O
+    Dying, // Exit called, cleanup starting
+    Zombie, // Exited, awaiting Shinigami
+    Dissolved, // Gone, slot reusable
 };
 
 pub const Context = packed struct {