mm: pmap overhaul

src/mm/amd64/impl.zig

@@ -17,17 +17,14 @@ pub const hhdm_base = 0xffff800000000000;
 pub const kernel_heap_base = 0xffffe00000000000;
 pub const pfndb_base = 0xffffff0000000000;
 
-pub const PagingImpl = struct {
+pub const levels = [_]mi.PMapLevel{
-    const Self = @This();
+    .{ .shift = 12, .mask = 0x1ff, .leaf = true }, // 4K
+    .{ .shift = 21, .mask = 0x1ff, .leaf = true }, // 2M
+    .{ .shift = 30, .mask = 0x1ff, .leaf = true }, // 1G
+    .{ .shift = 39, .mask = 0x1ff, .leaf = false }, // root
+};
 
-    pub const levels = [_]mi.radix_pmap.Level{
+pub const Pte = packed struct(u64) {
-        .{ .shift = 12, .leaf = true }, // 4K
-        .{ .shift = 21, .leaf = true }, // 2M
-        .{ .shift = 30, .leaf = true }, // 1G
-        .{ .shift = 39, .leaf = false }, // root
-    };
-
-    pub const Pte = packed struct(u64) {
     present: bool,
     writable: bool,
     user: bool,
@@ -45,9 +42,35 @@ pub const PagingImpl = struct {
     pub fn address(self: Pte) r.PAddr {
         return @as(r.PAddr, self.addr) << 12;
     }
-    };
 
-    pub inline fn make_table_pte(pa: r.PAddr) Pte {
+    pub fn is_present(self: Pte) bool {
+        return self.present;
+    }
+
+    pub fn load(table: *Pte) Pte {
+        return @atomicLoad(Pte, table, .monotonic);
+    }
+
+    pub fn zero() Pte {
+        return Pte{
+            .present = false,
+            .writable = false,
+            .user = false,
+            .write_through = false,
+            .cache_disable = false,
+            .accessed = false,
+            .dirty = false,
+            .huge_page = false,
+            .global = false,
+            .available = 0,
+            .addr = 0,
+            .ignored_2 = 0,
+            .nx = false,
+        };
+    }
+};
+
+pub inline fn make_table_pte(pa: r.PAddr) Pte {
     return Pte{
         .present = true,
         .writable = true,
@@ -63,9 +86,9 @@ pub const PagingImpl = struct {
         .ignored_2 = 0,
         .nx = false,
     };
-    }
+}
 
-    pub inline fn make_leaf_pte(pa: r.PAddr, flags: mm.MapFlags, level: usize) Pte {
+pub inline fn make_leaf_pte(pa: r.PAddr, flags: mm.MapFlags, level: usize) Pte {
     return Pte{
         .present = true,
         .writable = flags.write,
@@ -81,28 +104,24 @@ pub const PagingImpl = struct {
         .ignored_2 = 0,
         .nx = !flags.execute,
     };
-    }
+}
 
-    pub inline fn activate(root_pa: r.PAddr) void {
+pub inline fn activate(root_pa: r.PAddr) void {
     amd64.write_cr(3, root_pa);
-    }
+}
 
-    pub fn is_leaf_level_enabled(level: usize) bool {
+pub fn is_leaf_level_enabled(level: usize) bool {
     if (!levels[level].leaf) return false;
     if (level == 2) return amd64.cpu_features.gib_pages;
     return true;
-    }
+}
-};
-
-pub const PMap = mi.radix_pmap.RadixPmap(PagingImpl);
 
 pub fn init_kernel() void {
-    const higher_half_start = phys_to_virt(0);
+    mi.kernel_space.pmap.root_pa = mi.phys.early_alloc();
-    const higher_half_size: usize = 256 * (1 << 39);
+    const table_ptr: [*]Pte = @ptrFromInt(mm.p2v(mi.kernel_space.pmap.root_pa));
+    @memset(table_ptr[0..512], Pte.zero());
 
-    mi.kernel_space.pmap.impl.root_pa = mi.phys.alloc();
+    for (256..512) |i| {
-    const table_ptr: [*]u64 = @ptrFromInt(mm.p2v(mi.kernel_space.pmap.impl.root_pa));
+        table_ptr[i] = make_table_pte(mi.phys.early_alloc());
-    @memset(table_ptr[0..512], 0);
+    }
-
-    mi.kernel_space.pmap.impl.preallocate(higher_half_start, higher_half_size, 2);
 }

src/mm/heap.zig

@@ -8,16 +8,33 @@ pub fn init() void {
 }
 
 pub fn alloc(size: usize, policy: mm.WaitPolicy) mm.Error!*anyopaque {
-    _ = policy;
-
     mi.kernel_space.lock.acquire();
 
-    const addr = try mi.kernel_space.arena.alloc(size, .{});
+    const addr = mi.kernel_space.arena.alloc(size, .{}) catch {
+        mi.kernel_space.lock.release();
+        return mm.Error.OutOfMemory;
+    };
 
-    mi.kernel_space.pmap.map_range_allocating(addr, size, .{ .read = true, .write = true }, .DontWaitForMemory);
+    const npages = size / mm.page_size;
+
+    for (0..npages) |i| {
+        const pte = mi.pmap.wire_pte(&mi.kernel_space, addr + (i * mm.page_size), policy) catch {
+            mi.kernel_space.lock.release();
+            return mm.Error.OutOfMemory;
+        };
 
         mi.kernel_space.lock.release();
 
+        const page = mi.phys.alloc_opts(.{ .policy = policy }) orelse {
+            return mm.Error.OutOfMemory;
+        };
+
+        mi.kernel_space.lock.acquire();
+
+        pte.* = mi.impl.make_leaf_pte(page, .{ .read = true, .write = true }, 0);
+    }
+
+    mi.kernel_space.lock.release();
     return @ptrFromInt(addr);
 }
 

src/mm/init.zig

@@ -33,30 +33,36 @@ fn map_kernel(boot_info: *pl.BootInfo) void {
         .read = true,
         .write = true,
         .global = true,
-    });
+    }, .Boot);
 
     mi.kernel_space.pmap.map_contiguous_range(text_start, kaddr.physical_base + (text_start - kaddr.virtual_base), text_size, .{
         .read = true,
         .execute = true,
         .global = true,
-    });
+    }, .Boot);
 
     mi.kernel_space.pmap.map_contiguous_range(rodata_start, kaddr.physical_base + (rodata_start - kaddr.virtual_base), rodata_size, .{
         .read = true,
         .global = true,
-    });
+    }, .Boot);
 
     mi.kernel_space.pmap.map_contiguous_range(data_start, kaddr.physical_base + (data_start - kaddr.virtual_base), data_size, .{
         .read = true,
         .write = true,
         .global = true,
-    });
+    }, .Boot);
 
-    mi.kernel_space.pmap.map_contiguous_range(mm.p2v(0), 0, mi.impl.hhdm_minimum_max_address, .{
+    mi.kernel_space.pmap.map_contiguous_range(
+        mm.p2v(0),
+        0,
+        mi.impl.hhdm_minimum_max_address,
+        .{
             .read = true,
             .write = true,
             .global = true,
-    });
+        },
+        .Boot,
+    );
 
     // Now go through every usable entry and map to the HHDM every part that isnt covered by [0, hhdm_minimum_max_address).
     // We can't blindly map until the maximum usable physical address because on some CPUs this might cause MCEs.
@@ -81,11 +87,17 @@ fn map_kernel(boot_info: *pl.BootInfo) void {
             entry_size -= adjust;
         }
 
-        mi.kernel_space.pmap.map_contiguous_range(mm.p2v(entry_start), entry_start, entry_size, .{
+        mi.kernel_space.pmap.map_contiguous_range(
+            mm.p2v(entry_start),
+            entry_start,
+            entry_size,
+            .{
                 .read = true,
                 .write = true,
                 .global = true,
-        });
+            },
+            .Boot,
+        );
     }
 }
 

src/mm/phys.zig

@@ -23,7 +23,7 @@ var free_pages: usize = 0;
 
 pub var usable_memory: std.atomic.Value(usize) = .init(0);
 
-fn early_alloc() usize {
+pub fn early_alloc() usize {
     while (early_alloc_entry_idx < memory_map.entry_count) {
         const entry = &memory_map.entries[early_alloc_entry_idx];
 
@@ -72,10 +72,6 @@ fn wait_for_pages(old_ipl: ke.Ipl) void {
 /// Allocate a page of physical memory.
 /// This may block if no memory is available.
 pub fn alloc() r.PAddr {
-    if (!bootstrapped) {
-        @branchHint(.unlikely);
-        return early_alloc();
-    }
     const ipl = list_lock.acquire();
     defer list_lock.release(ipl);
 
@@ -97,11 +93,6 @@ pub fn alloc() r.PAddr {
 }
 
 pub fn alloc_opts(opts: struct { policy: mm.WaitPolicy }) ?r.PAddr {
-    if (!bootstrapped) {
-        @branchHint(.unlikely);
-        return early_alloc();
-    }
-
     const ipl = list_lock.acquire();
     defer list_lock.release(ipl);
 
@@ -216,7 +207,7 @@ pub fn init(boot_info: *pl.BootInfo) linksection(r.init) void {
                 .write = true,
                 .global = true,
             },
-            .DontWaitForMemory,
+            .Boot,
         );
     }
 }

src/mm/pmap.zig

@@ -1,36 +1,47 @@
-//! Platform-agnostic page map interface.
+//! Generic radix-tree page map implementation.
 //!
-//! This file defines the `PMap` wrapper and the two built-in mapping sources
+//! `RadixPmap` is a comptime-parameterized page table walker that implements
-//! (`ContiguousSource`, `AllocatingSource`) used to drive `PMap.map_from`.
+//! virtual-to-physical address mapping over a multi-level radix tree (e.g.
+//! x86-64's four-level PT). The caller supplies an
+//! `Impl` type that describes the tree's level structure and provides
+//! platform-specific PTE constructors.
 //!
-//! `PMap` itself is a thin wrapper over `mi.impl.PMap`, which is expected to
+//! The `Impl` type passed to `RadixPmap` must provide:
-//! be a `RadixPmap` instantiation (or any type that satisfies the same
-//! interface). To add a new platform, provide a concrete `impl.PMap` that
-//! exposes:
 //!
-//!   - `map_from(va: VAddr, size: usize, source: anytype) void`
+//!   - `levels: []const Level`: ordered from the leaf up to the root,
-//!   - `activate() void`
+//!     each entry carrying the VA index shift and whether that level can
-//!
+//!     hold a leaf (large-page) mapping.
-//! Mapping sources are ordinary structs with a single method:
+//!   - `is_leaf_level_enabled(level: usize) bool`:runtime gate that
-//!
+//!     controls whether large pages are used at a given level (e.g. to
-//!   - `next(self: *const Source, va: VAddr) MapItem`
+//!     disable 1 GiB pages on hardware that lacks them).
-//!
+//!   - `Pte`: the PTE type. Must be zero-initializable and expose a
-//! `next` is called once per mapping step and returns the physical address,
+//!     `present: bool` field and an `address() PAddr` method.
-//! length, and flags for that step. The two sources here cover the common
+//!   - `make_table_pte(pa: PAddr) Pte`: construct a non-leaf PTE pointing
-//! cases (contiguous PA range and on-demand allocation); custom sources can
+//!     to a child page table at `pa`.
-//! be passed directly to `map_from` for other needs such as file-backed or
+//!   - `make_leaf_pte(pa: PAddr, flags: MapFlags, level: usize) Pte`:
-//! scatter-gather mappings.
+//!     construct a leaf PTE mapping `pa` with the given flags at the
+//!     chosen level.
+//!   - `activate(root_pa: PAddr) void`: install the root page table into
+//!     the CPU (e.g. write CR3 or satp).
+
 const r = @import("root");
+const rtl = @import("rtl");
+const std = @import("std");
 const mm = r.mm;
 const mi = mm.private;
 
 /// Describes a contiguous region of physical memory to be mapped at a given virtual address.
-pub const MapItem = struct {
+const MapItem = struct {
     pa: r.PAddr,
     len: usize,
     flags: mm.MapFlags,
 };
 
+const Allocator = enum {
+    Boot,
+    Normal,
+};
+
 /// A mapping source that maps a single contiguous virtual-to-physical region.
 /// All calls to `next` return the same physical base address and size, so this
 /// is intended to be consumed exactly once by `map_from`.
@@ -53,10 +64,10 @@ const ContiguousSource = struct {
 /// returns a single page, so this is intended to be consumed by `map_from` in a loop until the entire range is mapped.
 const AllocatingSource = struct {
     flags: mm.MapFlags,
-    policy: mm.WaitPolicy,
+    allocator: Allocator,
 
     pub fn next(self: *const AllocatingSource, _: r.VAddr) MapItem {
-        const pa = mi.phys.alloc_opts(.{ .policy = self.policy }) orelse @panic("Out of memory");
+        const pa = if (self.allocator == .Boot) mi.phys.early_alloc() else mi.phys.alloc();
         return .{
             .pa = pa,
             .len = mm.page_size,
@@ -65,26 +76,64 @@ const AllocatingSource = struct {
     }
 };
 
-/// Generic wrapper over the platform-specific pmap implementation.
+fn level_shift(level: usize) u6 {
-/// Provides some convenient methods for common mapping patterns.
+    return mi.impl.levels[level].shift;
+}
+
+fn page_size(level: usize) usize {
+    return @as(usize, 1) << level_shift(level);
+}
+
+fn index_for_level(va: usize, level: usize) usize {
+    return (va >> level_shift(level)) & mi.impl.levels[level].mask;
+}
+
 pub const PMap = struct {
     const Self = @This();
 
-    impl: mi.impl.PMap,
+    const num_levels = mi.impl.levels.len;
+    const entries_per_table = 512;
+    const index_mask = entries_per_table - 1;
 
-    /// Map a virtual address range to physical addresses provided by the source.
+    root_pa: r.PAddr,
-    pub fn map_from(self: *Self, va_start: r.VAddr, size: usize, source: anytype) void {
-        const Source = @TypeOf(source);
 
-        comptime {
+    fn is_leaf_level(level: usize) bool {
-            if (!@hasDecl(Source, "next")) @compileError("source must provide next(self: *Source, va: b.VAddr) MapItem");
+        return mi.impl.levels[level].leaf;
     }
 
-        self.impl.map_from(va_start, size, source);
+    fn leaf_level_enabled(level: usize) bool {
+        return is_leaf_level(level) and mi.impl.is_leaf_level_enabled(level);
+    }
+
+    fn cursor(self: *Self, va: r.VAddr) Cursor {
+        var c = Cursor{
+            .pmap = self,
+            .va = va,
+            .tables = undefined,
+            .top_level = num_levels - 1,
+        };
+
+        c.tables[num_levels - 1] = @ptrFromInt(mm.p2v(self.root_pa));
+        return c;
+    }
+
+    pub fn activate(self: *Self) void {
+        mi.impl.activate(self.root_pa);
+    }
+
+    pub fn map_from(self: *Self, va: r.VAddr, size: usize, source: anytype, comptime allocator: Allocator) void {
+        var c = self.cursor(va);
+        var remain = size;
+
+        while (remain > 0) {
+            const item = source.next(c.va);
+            c.map_range(item.pa, item.len, item.flags, allocator);
+            remain -|= item.len;
+        }
     }
 
     /// Map a contiguous virtual address range to a contiguous physical address range.
-    pub fn map_contiguous_range(self: *Self, va: r.VAddr, pa: r.PAddr, size: usize, flags: mm.MapFlags) void {
+    pub fn map_contiguous_range(self: *Self, va: r.VAddr, pa: r.PAddr, size: usize, flags: mm.MapFlags, comptime allocator: Allocator) void {
         const src = ContiguousSource{
             .flags = flags,
             .size = size,
@@ -92,42 +141,238 @@ pub const PMap = struct {
             .base_pa = pa,
         };
 
-        self.map_from(va, size, src);
+        self.map_from(va, size, src, allocator);
     }
 
-    /// Map a virtual address range to physical addresses allocated on demand.
+    /// Map a virtual address range to physical addresses allocated on demand by the boot allocator.
-    pub fn map_range_allocating(self: *Self, va: r.VAddr, size: usize, flags: mm.MapFlags, policy: mm.WaitPolicy) void {
+    pub fn map_range_allocating(self: *Self, va: r.VAddr, size: usize, flags: mm.MapFlags, comptime allocator: Allocator) void {
         const src = AllocatingSource{
+            .allocator = allocator,
             .flags = flags,
-            .policy = policy,
         };
 
-        self.map_from(va, size, src);
+        self.map_from(va, size, src, allocator);
+    }
+
+    pub fn query(self: *Self, va: r.VAddr) ?r.PAddr {
+        var c = self.cursor(va);
+
+        var level: usize = num_levels - 1;
+        while (true) {
+            const idx = index_for_level(va, level);
+            const pte = c.tables[level][idx];
+
+            if (!pte.present) return null;
+
+            if (level == 0 or (leaf_level_enabled(level) and pte.huge_page)) {
+                return pte.address() + (va & (page_size(level) - 1));
+            }
+
+            if (!c.walk_down_resolve(level - 1)) return null;
+            level -= 1;
+        }
     }
 
     /// Unmap a contiguous range of virtual pages.
-    /// Returns a list threaded through PFNs of the unmapped physical pages if any.
+    /// Only small pages are supported.
+    /// Return a pfn that points to the physical pages that were unmapped.
     pub fn unmap(self: *Self, va: r.VAddr, size: usize) ?mi.PfnList {
-        return self.impl.unmap(va, size);
+        std.debug.assert(std.mem.isAligned(va, mm.page_size));
+
+        var c = self.cursor(va);
+
+        const npages: usize = size / mm.page_size;
+        var list: rtl.List = undefined;
+        list.init();
+
+        for (0..npages) |_| {
+            var level: usize = num_levels - 1;
+            blk: while (true) {
+                const idx = index_for_level(c.va, level);
+                const pte = c.tables[level][idx];
+
+                if (!pte.present) break :blk;
+
+                if (level == 0) {
+                    // Found a present entry, add it to the list.
+                    const addr = pte.address();
+                    const pfn = mm.page_to_pfn(addr);
+                    const page = mm.pfn_to_struct_page(pfn);
+                    list.insert_tail(&page.free.link);
+
+                    c.tables[level][idx] = std.mem.zeroes(mi.impl.Pte);
+                    break;
                 }
 
-    /// Flush the TLB entry of a single page.
+                if (!c.walk_down_resolve(level - 1)) return null;
-    pub fn flush(self: *Self, va: r.VAddr) void {
+                level -= 1;
-        self.impl.flush(va);
             }
 
-    /// Map a single page. Convenience wrapper around map_contiguous_range.
+            c.advance(mm.page_size);
-    pub fn map_page(self: *Self, va: r.VAddr, pa: r.PAddr, flags: mm.MapFlags) void {
-        self.map_contiguous_range(va, pa, mm.page_size, flags);
         }
 
-    /// Activate this pagemap.
+        if (list.is_empty()) return null;
-    pub fn activate(self: *Self) void {
+
-        self.impl.activate();
+        const head_free: *mm.PageFree = @fieldParentPtr("link", list.first());
+        const tail_free: *mm.PageFree = @fieldParentPtr("link", list.last());
+        const head: *mm.Page = @ptrCast(head_free);
+        const tail: *mm.Page = @ptrCast(tail_free);
+
+        return .{ .head = mm.struct_page_to_pfn(head), .tail = mm.struct_page_to_pfn(tail) };
     }
 
-    /// Given a virtual page, return its physical page, if present.
+    pub const Cursor = struct {
-    pub fn query(self: *Self, va: r.VAddr) ?r.PAddr {
+        pmap: *Self,
-        return self.impl.query(va);
+        va: r.VAddr,
+        /// Cached pointers to each level's page table, filled as the cursor
+        /// descends. Only indices in `[0, top_level]` are valid.
+        tables: [num_levels][*]mi.impl.Pte,
+        /// Index of the highest level for which `tables` currently holds a
+        /// valid pointer. Descends toward 0 (the leaf) as `walk_down` runs.
+        /// Reset upward in `advance` when the VA crosses a level boundary.
+        top_level: usize,
+
+        fn walk_down(self: *Cursor, target_level: usize, allocate: bool, comptime allocator: Allocator) error{PageNotMapped}!void {
+            var current_level = self.top_level;
+
+            while (current_level > target_level) {
+                const idx = index_for_level(self.va, current_level);
+                var pte = mi.impl.Pte.load(&self.tables[current_level][idx]);
+
+                if (!pte.present) {
+                    if (!allocate) {
+                        return error.PageNotMapped;
                     }
+
+                    const new_table_pa = if (allocator == .Boot) mi.phys.early_alloc() else mi.phys.alloc();
+
+                    const table_ptr: [*]mi.impl.Pte = @ptrFromInt(mm.p2v(new_table_pa));
+                    @memset(table_ptr[0..entries_per_table], std.mem.zeroes(mi.impl.Pte));
+
+                    pte = mi.impl.make_table_pte(new_table_pa);
+                    self.tables[current_level][idx] = pte;
+                }
+
+                const child_va = mm.p2v(pte.address());
+                self.tables[current_level - 1] = @ptrFromInt(child_va);
+                self.top_level -= 1;
+                current_level -= 1;
+            }
+        }
+
+        pub fn walk_down_resolve(self: *Cursor, target_level: usize) bool {
+            self.walk_down(target_level, false, .Normal) catch return false;
+            return true;
+        }
+
+        /// Choose the largest page size that satisfies the alignment and
+        /// size constraints of the current VA, PA, and remaining byte count.
+        /// Tries large pages first; falls back to 4K.
+        fn choose_target_level(self: *Cursor, pa: usize, remain: usize) usize {
+            var level: usize = num_levels;
+            while (level > 0) {
+                level -= 1;
+                if (!leaf_level_enabled(level)) continue;
+
+                const pg_size = page_size(level);
+                const align_mask = pg_size - 1;
+
+                if (remain >= pg_size and
+                    (self.va & align_mask) == 0 and
+                    (pa & align_mask) == 0)
+                {
+                    return level;
+                }
+            }
+
+            @panic("no leaf level can map current alignment/size");
+        }
+
+        pub fn map_range(self: *Cursor, pa: r.PAddr, size: usize, flags: mm.MapFlags, comptime allocator: Allocator) void {
+            var remain = size;
+            var current_pa = pa;
+
+            while (remain > 0) {
+                const target_level = self.choose_target_level(current_pa, remain);
+
+                self.walk_down(target_level, true, allocator) catch @panic("walk_down failed during map_range");
+
+                const table = self.tables[target_level];
+                table[index_for_level(self.va, target_level)] =
+                    mi.impl.make_leaf_pte(current_pa, flags, target_level);
+
+                const step = page_size(target_level);
+                current_pa += step;
+                remain -= step;
+
+                self.advance(step);
+            }
+        }
+
+        /// Advance the cursor's virtual address by `offset` bytes and
+        /// invalidate any cached table pointers that no longer apply.
+        /// The cursor tracks the highest level whose index is still valid
+        /// for the new VA, `walk_down` will re-descend from there on the
+        /// next mapping operation.
+        pub fn advance(self: *Cursor, offset: usize) void {
+            const old_va = self.va;
+            self.va +%= offset;
+
+            var level: usize = num_levels - 1;
+            while (level > 0) : (level -= 1) {
+                const shift = level_shift(level);
+                if ((old_va >> shift) != (self.va >> shift)) {
+                    self.top_level = level;
+                    return;
+                }
+            }
+        }
+    };
 };
+
+pub fn wire_pte(space: *mi.Space, va: r.VAddr, policy: mm.WaitPolicy) mm.Error!*mi.impl.Pte {
+    std.debug.assert(space.lock.is_locked());
+
+    var table: [*]mi.impl.Pte = @ptrFromInt(mm.p2v(mi.kernel_space.pmap.root_pa));
+    var level = mi.impl.levels.len - 1;
+
+    while (true) {
+        const idx = index_for_level(va, level);
+        var entry = table[idx];
+
+        if (level == 0) {
+            std.debug.assert(!entry.present);
+            return &table[idx];
+        }
+
+        if (!entry.present) {
+            // Release the lock while allocating a new page table, since that may block.
+            space.lock.release();
+
+            const new_table_pa = mi.phys.alloc_opts(.{ .policy = policy }) orelse {
+                space.lock.acquire();
+                return mm.Error.OutOfMemory;
+            };
+
+            // Re-acquire the lock and re-validate the entry, since another thread may have raced to allocate it.
+            space.lock.acquire();
+
+            entry = table[idx];
+
+            if (entry.present) {
+                // Another thread beat us to it, free the page we just allocated and continue.
+                mi.phys.free(new_table_pa);
+            } else {
+                const new_table_ptr: [*]u64 = @ptrFromInt(mm.p2v(new_table_pa));
+                @memset(new_table_ptr[0..512], 0);
+
+                table[idx] = mi.impl.make_table_pte(new_table_pa);
+                entry = table[idx];
+            }
+        }
+
+        const child_va = mm.p2v(entry.address());
+        table = @ptrFromInt(child_va);
+        level -= 1;
+    }
+}

src/mm/private.zig

@@ -1,6 +1,5 @@
 const config = @import("config");
 pub const phys = @import("phys.zig");
-pub const radix_pmap = @import("radix_pmap.zig");
 pub const pmap = @import("pmap.zig");
 pub const init = @import("init.zig");
 pub const vmem = @import("vmem.zig");
@@ -22,6 +21,12 @@ pub const PfnList = extern struct {
     tail: u32,
 };
 
+pub const PMapLevel = struct {
+    shift: u6,
+    mask: usize,
+    leaf: bool,
+};
+
 pub const impl = switch (config.arch) {
     .amd64 => @import("amd64/impl.zig"),
     .um => @import("um/impl.zig"),

src/mm/radix_pmap.zig

@@ -1,288 +0,0 @@
-//! Generic radix-tree page map implementation.
-//!
-//! `RadixPmap` is a comptime-parameterized page table walker that implements
-//! virtual-to-physical address mapping over a multi-level radix tree (e.g.
-//! x86-64's four-level PT). The caller supplies an
-//! `Impl` type that describes the tree's level structure and provides
-//! platform-specific PTE constructors.
-//!
-//! The `Impl` type passed to `RadixPmap` must provide:
-//!
-//!   - `levels: []const Level`: ordered from the leaf up to the root,
-//!     each entry carrying the VA index shift and whether that level can
-//!     hold a leaf (large-page) mapping.
-//!   - `is_leaf_level_enabled(level: usize) bool`:runtime gate that
-//!     controls whether large pages are used at a given level (e.g. to
-//!     disable 1 GiB pages on hardware that lacks them).
-//!   - `Pte`: the PTE type. Must be zero-initializable and expose a
-//!     `present: bool` field and an `address() PAddr` method.
-//!   - `make_table_pte(pa: PAddr) Pte`: construct a non-leaf PTE pointing
-//!     to a child page table at `pa`.
-//!   - `make_leaf_pte(pa: PAddr, flags: MapFlags, level: usize) Pte`:
-//!     construct a leaf PTE mapping `pa` with the given flags at the
-//!     chosen level.
-//!   - `activate(root_pa: PAddr) void`: install the root page table into
-//!     the CPU (e.g. write CR3 or satp).
-
-const r = @import("root");
-const rtl = @import("rtl");
-const std = @import("std");
-const mm = r.mm;
-const mi = mm.private;
-
-pub const Level = struct {
-    shift: u6,
-    leaf: bool,
-};
-
-pub fn RadixPmap(comptime Impl: type) type {
-    return struct {
-        const Self = @This();
-
-        const num_levels = Impl.levels.len;
-        const entries_per_table = 512;
-        const index_mask = entries_per_table - 1;
-
-        root_pa: r.PAddr,
-
-        fn level_shift(level: usize) u6 {
-            return Impl.levels[level].shift;
-        }
-
-        fn page_size(level: usize) usize {
-            return @as(usize, 1) << level_shift(level);
-        }
-
-        fn index_for_level(va: usize, level: usize) usize {
-            return (va >> level_shift(level)) & index_mask;
-        }
-
-        fn is_leaf_level(level: usize) bool {
-            return Impl.levels[level].leaf;
-        }
-
-        fn leaf_level_enabled(level: usize) bool {
-            return is_leaf_level(level) and Impl.is_leaf_level_enabled(level);
-        }
-
-        fn cursor(self: *Self, start_va: usize) Cursor {
-            var tables: [num_levels][*]Impl.Pte = undefined;
-            tables[num_levels - 1] = @ptrFromInt(mm.p2v(self.root_pa));
-            return .{
-                .pmap = self,
-                .va = start_va,
-                .tables = tables,
-                .top_level = num_levels - 1,
-            };
-        }
-
-        pub fn activate(self: *Self) void {
-            Impl.activate(self.root_pa);
-        }
-
-        pub fn map_from(self: *Self, va: r.VAddr, size: usize, source: anytype) void {
-            var c = self.cursor(va);
-            var remain = size;
-
-            while (remain > 0) {
-                const item = source.next(c.va);
-                c.map_range(item.pa, item.len, item.flags);
-                remain -|= item.len;
-            }
-        }
-
-        pub fn query(self: *Self, va: r.VAddr) ?r.PAddr {
-            var c = self.cursor(va);
-
-            var level: usize = num_levels - 1;
-            while (true) {
-                const idx = index_for_level(va, level);
-                const pte = c.tables[level][idx];
-
-                if (!pte.present) return null;
-
-                if (level == 0 or (leaf_level_enabled(level) and pte.huge_page)) {
-                    return pte.address() + (va & (page_size(level) - 1));
-                }
-
-                if (!c.walk_down_resolve(level - 1)) return null;
-                level -= 1;
-            }
-        }
-
-        /// Unmap a contiguous range of virtual pages.
-        /// Only small pages are supported.
-        /// Return a pfn that points to the physical pages that were unmapped.
-        pub fn unmap(self: *Self, va: r.VAddr, size: usize) ?mi.PfnList {
-            std.debug.assert(std.mem.isAligned(va, mm.page_size));
-
-            var c = self.cursor(va);
-
-            const npages: usize = size / mm.page_size;
-            var list: rtl.List = undefined;
-            list.init();
-
-            for (0..npages) |_| {
-                var level: usize = num_levels - 1;
-                blk: while (true) {
-                    const idx = index_for_level(c.va, level);
-                    const pte = c.tables[level][idx];
-
-                    if (!pte.present) break :blk;
-
-                    if (level == 0) {
-                        // Found a present entry, add it to the list.
-                        const addr = pte.address();
-                        const pfn = mm.page_to_pfn(addr);
-                        const page = mm.pfn_to_struct_page(pfn);
-                        list.insert_tail(&page.free.link);
-
-                        c.tables[level][idx] = std.mem.zeroes(Impl.Pte);
-                        break;
-                    }
-
-                    if (!c.walk_down_resolve(level - 1)) return null;
-                    level -= 1;
-                }
-
-                c.advance(mm.page_size);
-            }
-
-            if (list.is_empty()) return null;
-
-            const head_free: *mm.PageFree = @fieldParentPtr("link", list.first());
-            const tail_free: *mm.PageFree = @fieldParentPtr("link", list.last());
-            const head: *mm.Page = @ptrCast(head_free);
-            const tail: *mm.Page = @ptrCast(tail_free);
-
-            return .{ .head = mm.struct_page_to_pfn(head), .tail = mm.struct_page_to_pfn(tail) };
-        }
-
-        /// Ensure all intermediate page tables down to `target_level` exist for
-        /// the given virtual address range, allocating them as needed. Useful for
-        /// pre-populating a shared region (e.g. the kernel half of a user pmap)
-        /// before any leaf mappings are installed.
-        pub fn preallocate(self: *Self, va: r.VAddr, size: usize, target_level: usize) void {
-            var c = self.cursor(va);
-            var remain = size;
-
-            // Step by one table at `target_level + 1` at a time so each
-            // `walk_down` call sees a fresh index at that level.
-            const step = page_size(target_level + 1);
-
-            while (remain > 0) {
-                c.walk_down(target_level, true) catch @panic("walk_down failed during preallocate");
-                const advance_amount = @min(remain, step);
-                c.advance(advance_amount);
-                remain -= advance_amount;
-            }
-        }
-
-        pub const Cursor = struct {
-            pmap: *Self,
-            va: r.VAddr,
-            /// Cached pointers to each level's page table, filled as the cursor
-            /// descends. Only indices in `[0, top_level]` are valid.
-            tables: [num_levels][*]Impl.Pte,
-            /// Index of the highest level for which `tables` currently holds a
-            /// valid pointer. Descends toward 0 (the leaf) as `walk_down` runs.
-            /// Reset upward in `advance` when the VA crosses a level boundary.
-            top_level: usize,
-
-            fn walk_down(self: *Cursor, target_level: usize, allocate: bool) error{PageNotMapped}!void {
-                var current_level = self.top_level;
-
-                while (current_level > target_level) {
-                    const idx = index_for_level(self.va, current_level);
-                    var pte = self.tables[current_level][idx];
-
-                    if (!pte.present) {
-                        if (!allocate) {
-                            return error.PageNotMapped;
-                        }
-
-                        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));
-
-                        pte = Impl.make_table_pte(new_table_pa);
-                        self.tables[current_level][idx] = pte;
-                    }
-
-                    const child_va = mm.p2v(pte.address());
-                    self.tables[current_level - 1] = @ptrFromInt(child_va);
-                    self.top_level -= 1;
-                    current_level -= 1;
-                }
-            }
-
-            pub fn walk_down_resolve(self: *Cursor, target_level: usize) bool {
-                self.walk_down(target_level, false) catch return false;
-                return true;
-            }
-
-            /// Choose the largest page size that satisfies the alignment and
-            /// size constraints of the current VA, PA, and remaining byte count.
-            /// Tries large pages first; falls back to 4K.
-            fn choose_target_level(self: *Cursor, pa: usize, remain: usize) usize {
-                var level: usize = num_levels;
-                while (level > 0) {
-                    level -= 1;
-                    if (!leaf_level_enabled(level)) continue;
-
-                    const pg_size = page_size(level);
-                    const align_mask = pg_size - 1;
-
-                    if (remain >= pg_size and
-                        (self.va & align_mask) == 0 and
-                        (pa & align_mask) == 0)
-                    {
-                        return level;
-                    }
-                }
-
-                @panic("no leaf level can map current alignment/size");
-            }
-
-            pub fn map_range(self: *Cursor, pa: r.PAddr, size: usize, flags: mm.MapFlags) void {
-                var remain = size;
-                var current_pa = pa;
-
-                while (remain > 0) {
-                    const target_level = self.choose_target_level(current_pa, remain);
-
-                    self.walk_down(target_level, true) catch @panic("walk_down failed during map_range");
-
-                    const table = self.tables[target_level];
-                    table[index_for_level(self.va, target_level)] =
-                        Impl.make_leaf_pte(current_pa, flags, target_level);
-
-                    const step = page_size(target_level);
-                    current_pa += step;
-                    remain -= step;
-
-                    self.advance(step);
-                }
-            }
-
-            /// Advance the cursor's virtual address by `offset` bytes and
-            /// invalidate any cached table pointers that no longer apply.
-            /// The cursor tracks the highest level whose index is still valid
-            /// for the new VA, `walk_down` will re-descend from there on the
-            /// next mapping operation.
-            pub fn advance(self: *Cursor, offset: usize) void {
-                const old_va = self.va;
-                self.va +%= offset;
-
-                var level: usize = num_levels - 1;
-                while (level > 0) : (level -= 1) {
-                    const shift = level_shift(level);
-                    if ((old_va >> shift) != (self.va >> shift)) {
-                        self.top_level = level;
-                        return;
-                    }
-                }
-            }
-        };
-    };
-}

src/mm/tlb.zig

@@ -10,6 +10,7 @@ const ps = r.ps;
 
 pub var sync_shootdowns: std.atomic.Value(usize) = .init(0);
 pub var async_shootdowns: std.atomic.Value(usize) = .init(0);
+pub var sent_shootdowns: std.atomic.Value(usize) = .init(0);
 
 /// Reclaim the memory associated with a state.
 /// This can only be done when all CPUs have flushed their TLB.
@@ -26,11 +27,11 @@ fn reclaim_state(state: *ke.ShootdownState) void {
 
     space.lock.release();
 
-    _ = async_shootdowns.fetchAdd(1, .monotonic);
-
     // Free the physical pages
     mm.phys.free_batch(mm.pfn_to_struct_page(pfn_list.head), mm.pfn_to_struct_page(pfn_list.tail), npages);
 
+    _ = async_shootdowns.fetchAdd(1, .monotonic);
+
     // Release the slot after all copied state has been consumed.
     state.release();
 }
@@ -40,7 +41,6 @@ fn worker_thread(_: ?*anyopaque) void {
         const entry = ke.shootdown.shootdowns.remove();
         const state: *ke.ShootdownState = @fieldParentPtr("link", entry);
 
-        _ = async_shootdowns.fetchAdd(1, .monotonic);
         reclaim_state(state);
     }
 }

src/platform/pc/smp.zig

@@ -55,7 +55,7 @@ pub fn init() linksection(r.init) void {
         .read = true,
         .write = true,
         .execute = true,
-    });
+    }, .Normal);
 
     // From Linux, on modern CPUs we can skip the long delay after INIT.
     const skip_delay = switch (amd64.cpu_features.vendor) {