os: mach taggedPageAllocator

src/os/mach.zig

@@ -1,4 +1,7 @@
 const std = @import("std");
+const assert = @import("../quirks.zig").inlineAssert;
+const mem = std.mem;
+const Allocator = std.mem.Allocator;
 
 /// macOS virtual memory tags for use with mach_vm_map/mach_vm_allocate.
 /// These identify memory regions in tools like vmmap and Instruments.
@@ -30,6 +33,118 @@ pub const VMTag = enum(u8) {
     }
 };
 
+/// Creates a page allocator that tags all allocated memory with the given
+/// VMTag.
+pub fn taggedPageAllocator(tag: VMTag) Allocator {
+    return .{
+        // We smuggle the tag in as the context pointer.
+        .ptr = @ptrFromInt(@as(usize, @intFromEnum(tag))),
+        .vtable = &TaggedPageAllocator.vtable,
+    };
+}
+
+/// This is based heavily on the Zig 0.15.2 PageAllocator implementation,
+/// with only the posix implementation. Zig 0.15.2 is MIT licensed.
+const TaggedPageAllocator = struct {
+    pub const vtable: Allocator.VTable = .{
+        .alloc = alloc,
+        .resize = resize,
+        .remap = remap,
+        .free = free,
+    };
+
+    fn alloc(context: *anyopaque, n: usize, alignment: mem.Alignment, ra: usize) ?[*]u8 {
+        _ = ra;
+        assert(n > 0);
+        const tag: VMTag = @enumFromInt(@as(u8, @truncate(@intFromPtr(context))));
+        return map(n, alignment, tag);
+    }
+
+    fn resize(context: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, return_address: usize) bool {
+        _ = context;
+        _ = alignment;
+        _ = return_address;
+        return realloc(memory, new_len, false) != null;
+    }
+
+    fn remap(context: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, return_address: usize) ?[*]u8 {
+        _ = context;
+        _ = alignment;
+        _ = return_address;
+        return realloc(memory, new_len, true);
+    }
+
+    fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, return_address: usize) void {
+        _ = context;
+        _ = alignment;
+        _ = return_address;
+        return unmap(@alignCast(memory));
+    }
+
+    pub fn map(n: usize, alignment: mem.Alignment, tag: VMTag) ?[*]u8 {
+        const page_size = std.heap.pageSize();
+        if (n >= std.math.maxInt(usize) - page_size) return null;
+        const alignment_bytes = alignment.toByteUnits();
+
+        const aligned_len = mem.alignForward(usize, n, page_size);
+        const max_drop_len = alignment_bytes - @min(alignment_bytes, page_size);
+        const overalloc_len = if (max_drop_len <= aligned_len - n)
+            aligned_len
+        else
+            mem.alignForward(usize, aligned_len + max_drop_len, page_size);
+        const hint = @atomicLoad(@TypeOf(std.heap.next_mmap_addr_hint), &std.heap.next_mmap_addr_hint, .unordered);
+        const slice = std.posix.mmap(
+            hint,
+            overalloc_len,
+            std.posix.PROT.READ | std.posix.PROT.WRITE,
+            .{ .TYPE = .PRIVATE, .ANONYMOUS = true },
+            tag.make(),
+            0,
+        ) catch return null;
+        const result_ptr = mem.alignPointer(slice.ptr, alignment_bytes) orelse return null;
+        // Unmap the extra bytes that were only requested in order to guarantee
+        // that the range of memory we were provided had a proper alignment in it
+        // somewhere. The extra bytes could be at the beginning, or end, or both.
+        const drop_len = result_ptr - slice.ptr;
+        if (drop_len != 0) std.posix.munmap(slice[0..drop_len]);
+        const remaining_len = overalloc_len - drop_len;
+        if (remaining_len > aligned_len) std.posix.munmap(@alignCast(result_ptr[aligned_len..remaining_len]));
+        const new_hint: [*]align(std.heap.page_size_min) u8 = @alignCast(result_ptr + aligned_len);
+        _ = @cmpxchgStrong(@TypeOf(std.heap.next_mmap_addr_hint), &std.heap.next_mmap_addr_hint, hint, new_hint, .monotonic, .monotonic);
+        return result_ptr;
+    }
+
+    pub fn unmap(memory: []align(std.heap.page_size_min) u8) void {
+        const page_aligned_len = mem.alignForward(usize, memory.len, std.heap.pageSize());
+        std.posix.munmap(memory.ptr[0..page_aligned_len]);
+    }
+
+    pub fn realloc(uncasted_memory: []u8, new_len: usize, may_move: bool) ?[*]u8 {
+        const memory: []align(std.heap.page_size_min) u8 = @alignCast(uncasted_memory);
+        const page_size = std.heap.pageSize();
+        const new_size_aligned = mem.alignForward(usize, new_len, page_size);
+
+        const page_aligned_len = mem.alignForward(usize, memory.len, page_size);
+        if (new_size_aligned == page_aligned_len)
+            return memory.ptr;
+
+        if (std.posix.MREMAP != void) {
+            // TODO: if the next_mmap_addr_hint is within the remapped range, update it
+            const new_memory = std.posix.mremap(memory.ptr, memory.len, new_len, .{ .MAYMOVE = may_move }, null) catch return null;
+            return new_memory.ptr;
+        }
+
+        if (new_size_aligned < page_aligned_len) {
+            const ptr = memory.ptr + new_size_aligned;
+            // TODO: if the next_mmap_addr_hint is within the unmapped range, update it
+            std.posix.munmap(@alignCast(ptr[0 .. page_aligned_len - new_size_aligned]));
+            return memory.ptr;
+        }
+
+        return null;
+    }
+};
+
 test "VMTag.make" {
     try std.testing.expectEqual(@as(i32, @bitCast(@as(u32, 240) << 24)), VMTag.application_specific_1.make());
 }