Add allocator example and custom allocator example

.github/workflows/ci.yml

@@ -31,6 +31,8 @@ jobs:
         run: |
           pushd examples/basic && zig build && popd
           pushd examples/init && zig build && popd
+          pushd examples/allocator-custom && zig build && popd
+          pushd examples/allocator-builtin && zig build && popd
 
   arkvm-tests:
     runs-on: ubuntu-latest

examples/allocator-builtin/build.zig

@@ -0,0 +1,36 @@
+const std = @import("std");
+const napi_build = @import("zig-napi").napi_build;
+
+pub fn build(b: *std.Build) !void {
+    const target = b.standardTargetOptions(.{});
+    const optimize = b.standardOptimizeOption(.{});
+
+    const zig_napi = b.dependency("zig-napi", .{});
+    const napi = zig_napi.module("napi");
+
+    const result = try napi_build.nativeAddonBuild(b, .{
+        .name = "hello",
+        .root_module_options = .{
+            .root_source_file = b.path("./src/hello.zig"),
+            .target = target,
+            .optimize = optimize,
+        },
+    });
+
+    if (result.arm64) |arm64| {
+        arm64.root_module.addImport("napi", napi);
+    }
+    if (result.arm) |arm| {
+        arm.root_module.addImport("napi", napi);
+    }
+    if (result.x64) |x64| {
+        x64.root_module.addImport("napi", napi);
+    }
+
+    const dts = try napi_build.generateTypeDefinition(b, .{
+        .root_source_file = b.path("./src/hello.zig"),
+        .output = b.path("index.d.ts"),
+        .napi_module = napi,
+    });
+    b.getInstallStep().dependOn(&dts.step);
+}

examples/allocator-builtin/build.zig.zon

@@ -0,0 +1,8 @@
+.{
+    .name = .allocator_builtin,
+    .version = "0.0.1",
+    .minimum_zig_version = "0.16.0",
+    .fingerprint = 0xad1ec621c68afe17,
+    .dependencies = .{ .@"zig-napi" = .{ .path = "../.." } },
+    .paths = .{ "build.zig", "build.zig.zon", "src" },
+}

examples/allocator-builtin/index.d.ts

@@ -0,0 +1,7 @@
+/* auto-generated by zig-addon */
+/* eslint-disable */
+export declare function allocator_kind(): string
+export declare function manual_allocation_roundtrip(len: number): boolean
+export declare function make_js_owned_buffer(len: number): Buffer
+export declare function make_copied_buffer(): Buffer
+export declare function input_sum(input: Buffer): number

examples/allocator-builtin/src/hello.zig

@@ -0,0 +1,42 @@
+const napi = @import("napi");
+
+pub fn allocator_kind() []const u8 {
+    return "builtin-page";
+}
+
+pub fn manual_allocation_roundtrip(len: u32) bool {
+    const allocator = napi.globalAllocator();
+    const bytes = allocator.alloc(u8, len) catch return false;
+    defer allocator.free(bytes);
+
+    @memset(bytes, 0x2a);
+    return bytes.len == len and (len == 0 or bytes[0] == 0x2a);
+}
+
+pub fn make_js_owned_buffer(env: napi.Env, len: u32) !napi.Buffer {
+    const allocator = napi.globalAllocator();
+    const bytes = try allocator.alloc(u8, len);
+    errdefer allocator.free(bytes);
+
+    for (bytes, 0..) |*byte, index| {
+        byte.* = @intCast((index + 3) % 251);
+    }
+    return try napi.Buffer.from(env, bytes);
+}
+
+pub fn make_copied_buffer(env: napi.Env) !napi.Buffer {
+    const bytes = [_]u8{ 2, 4, 6, 8 };
+    return try napi.Buffer.copy(env, &bytes);
+}
+
+pub fn input_sum(input: napi.Buffer) u32 {
+    var sum: u32 = 0;
+    for (input.asConstSlice()) |byte| {
+        sum += byte;
+    }
+    return sum;
+}
+
+comptime {
+    napi.NODE_API_MODULE("hello", @This());
+}

examples/allocator-custom/build.zig

@@ -0,0 +1,36 @@
+const std = @import("std");
+const napi_build = @import("zig-napi").napi_build;
+
+pub fn build(b: *std.Build) !void {
+    const target = b.standardTargetOptions(.{});
+    const optimize = b.standardOptimizeOption(.{});
+
+    const zig_napi = b.dependency("zig-napi", .{});
+    const napi = zig_napi.module("napi");
+
+    const result = try napi_build.nativeAddonBuild(b, .{
+        .name = "hello",
+        .root_module_options = .{
+            .root_source_file = b.path("./src/hello.zig"),
+            .target = target,
+            .optimize = optimize,
+        },
+    });
+
+    if (result.arm64) |arm64| {
+        arm64.root_module.addImport("napi", napi);
+    }
+    if (result.arm) |arm| {
+        arm.root_module.addImport("napi", napi);
+    }
+    if (result.x64) |x64| {
+        x64.root_module.addImport("napi", napi);
+    }
+
+    const dts = try napi_build.generateTypeDefinition(b, .{
+        .root_source_file = b.path("./src/hello.zig"),
+        .output = b.path("index.d.ts"),
+        .napi_module = napi,
+    });
+    b.getInstallStep().dependOn(&dts.step);
+}

examples/allocator-custom/build.zig.zon

@@ -0,0 +1,8 @@
+.{
+    .name = .allocator_custom,
+    .version = "0.0.1",
+    .minimum_zig_version = "0.16.0",
+    .fingerprint = 0x3eb28e374c2ad8c5,
+    .dependencies = .{ .@"zig-napi" = .{ .path = "../.." } },
+    .paths = .{ "build.zig", "build.zig.zon", "src" },
+}

examples/allocator-custom/index.d.ts

@@ -0,0 +1,17 @@
+/* auto-generated by zig-addon */
+/* eslint-disable */
+
+export interface Stats {
+  alloc_calls: number
+  free_calls: number
+  active_allocations: number
+  active_bytes: number
+}
+
+
+export declare function allocator_kind(): string
+export declare function allocator_stats(): Stats
+export declare function custom_allocation_roundtrip(len: number): boolean
+export declare function make_js_owned_buffer(len: number): Buffer
+export declare function make_copied_buffer(): Buffer
+export declare function input_sum(input: Buffer): number

examples/allocator-custom/src/counting_allocator.zig

@@ -0,0 +1,81 @@
+const std = @import("std");
+
+pub const Stats = struct {
+    alloc_calls: usize,
+    free_calls: usize,
+    active_allocations: isize,
+    active_bytes: isize,
+};
+
+pub const CountingAllocator = struct {
+    backing: std.mem.Allocator,
+    alloc_calls: std.atomic.Value(usize) = .init(0),
+    free_calls: std.atomic.Value(usize) = .init(0),
+    active_allocations: std.atomic.Value(isize) = .init(0),
+    active_bytes: std.atomic.Value(isize) = .init(0),
+
+    const Self = @This();
+
+    pub fn init(backing: std.mem.Allocator) Self {
+        return .{ .backing = backing };
+    }
+
+    pub fn allocator(self: *Self) std.mem.Allocator {
+        return .{
+            .ptr = self,
+            .vtable = &.{
+                .alloc = alloc,
+                .resize = resize,
+                .remap = remap,
+                .free = free,
+            },
+        };
+    }
+
+    pub fn stats(self: *Self) Stats {
+        return .{
+            .alloc_calls = self.alloc_calls.load(.monotonic),
+            .free_calls = self.free_calls.load(.monotonic),
+            .active_allocations = self.active_allocations.load(.monotonic),
+            .active_bytes = self.active_bytes.load(.monotonic),
+        };
+    }
+
+    fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
+        const self: *Self = @ptrCast(@alignCast(ctx));
+        const ptr = self.backing.rawAlloc(len, alignment, ret_addr) orelse return null;
+
+        _ = self.alloc_calls.fetchAdd(1, .monotonic);
+        _ = self.active_allocations.fetchAdd(1, .monotonic);
+        _ = self.active_bytes.fetchAdd(@intCast(len), .monotonic);
+        return ptr;
+    }
+
+    fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
+        const self: *Self = @ptrCast(@alignCast(ctx));
+        if (!self.backing.rawResize(memory, alignment, new_len, ret_addr)) {
+            return false;
+        }
+
+        _ = self.active_bytes.fetchAdd(@as(isize, @intCast(new_len)) - @as(isize, @intCast(memory.len)), .monotonic);
+        return true;
+    }
+
+    fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
+        const self: *Self = @ptrCast(@alignCast(ctx));
+        const ptr = self.backing.rawRemap(memory, alignment, new_len, ret_addr) orelse return null;
+
+        _ = self.active_bytes.fetchAdd(@as(isize, @intCast(new_len)) - @as(isize, @intCast(memory.len)), .monotonic);
+        return ptr;
+    }
+
+    fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
+        const self: *Self = @ptrCast(@alignCast(ctx));
+
+        _ = self.free_calls.fetchAdd(1, .monotonic);
+        _ = self.active_allocations.fetchSub(1, .monotonic);
+        _ = self.active_bytes.fetchSub(@intCast(memory.len), .monotonic);
+
+        self.backing.rawFree(memory, alignment, ret_addr);
+    }
+};

examples/allocator-custom/src/hello.zig

@@ -0,0 +1,54 @@
+const std = @import("std");
+const napi = @import("napi");
+const CountingAllocator = @import("counting_allocator.zig").CountingAllocator;
+const Stats = @import("counting_allocator.zig").Stats;
+
+var custom_allocator_state = CountingAllocator.init(std.heap.page_allocator);
+pub const napi_allocator = custom_allocator_state.allocator();
+
+pub fn allocator_kind() []const u8 {
+    return "custom-counting";
+}
+
+pub fn allocator_stats() Stats {
+    return custom_allocator_state.stats();
+}
+
+pub fn custom_allocation_roundtrip(len: u32) bool {
+    const allocator = napi.globalAllocator();
+    const bytes = allocator.alloc(u8, len) catch return false;
+    defer allocator.free(bytes);
+
+    for (bytes, 0..) |*byte, index| {
+        byte.* = @intCast(index % 251);
+    }
+    return bytes.len == len and (len == 0 or bytes[0] == 0);
+}
+
+pub fn make_js_owned_buffer(env: napi.Env, len: u32) !napi.Buffer {
+    const allocator = napi.globalAllocator();
+    const bytes = try allocator.alloc(u8, len);
+    errdefer allocator.free(bytes);
+
+    for (bytes, 0..) |*byte, index| {
+        byte.* = @intCast((index + 7) % 251);
+    }
+    return try napi.Buffer.from(env, bytes);
+}
+
+pub fn make_copied_buffer(env: napi.Env) !napi.Buffer {
+    const bytes = [_]u8{ 11, 13, 17, 19 };
+    return try napi.Buffer.copy(env, &bytes);
+}
+
+pub fn input_sum(input: napi.Buffer) u32 {
+    var sum: u32 = 0;
+    for (input.asConstSlice()) |byte| {
+        sum += byte;
+    }
+    return sum;
+}
+
+comptime {
+    napi.NODE_API_MODULE("hello", @This());
+}

scripts/arkvm/run_arkvm_tests.sh

@@ -131,6 +131,8 @@ if [[ -n "${ARKVM_TEST_SUITE:-}" ]]; then
 else
   run_case "examples/basic" "test/basic.ts" "__ZIG_NAPI_TEST_RESULT__" "test/basic" "-Darkvm-test=true -Doptimize=ReleaseSafe" "arkvm-host"
   run_case "examples/init" "test/init.ts" "__ZIG_NAPI_INIT_TEST_RESULT__" "test/init" "-Darkvm-test=true -Doptimize=ReleaseSafe" "arkvm-host"
+  run_case "examples/allocator-custom" "test/allocator-custom.ts" "__ZIG_NAPI_ALLOCATOR_CUSTOM_RESULT__" "test/allocator-custom" "-Darkvm-test=true -Doptimize=ReleaseSafe" "arkvm-host"
+  run_case "examples/allocator-builtin" "test/allocator-builtin.ts" "__ZIG_NAPI_ALLOCATOR_BUILTIN_RESULT__" "test/allocator-builtin" "-Darkvm-test=true -Doptimize=ReleaseSafe" "arkvm-host"
 fi
 
 [[ "${KEEP_WORKDIR}" == "1" ]] || rm -rf "${WORK_ROOT}"

src/build/napi-tsgen.zig

@@ -1838,6 +1838,9 @@ fn generate(allocator: std.mem.Allocator, io: std.Io, root_source_path: []const
     }
 
     inline for (root_info.decls) |decl| {
+        if (comptime std.mem.eql(u8, decl.name, "napi_allocator")) {
+            continue;
+        }
         const value = @field(root, decl.name);
         const decl_type = @TypeOf(value);
         if (comptime @typeInfo(decl_type) == .@"fn") {

src/napi.zig

@@ -68,38 +68,19 @@ pub fn FunctionRef(comptime Args: type, comptime Return: type) type {
 }
 pub const ObjectRef = reference.Reference(value.Object);
 
-/// Set the allocator used by napi wrappers, including JS-runtime-owned state.
-pub fn setAllocator(new_allocator: std.mem.Allocator) void {
-    global_allocator.global_manager.set(new_allocator);
-    global_allocator.runtime_manager.set(new_allocator);
-}
-
-/// Reset the allocator used by napi wrappers to the default page allocator.
-pub fn resetAllocator() void {
-    global_allocator.global_manager.set(std.heap.page_allocator);
-    global_allocator.runtime_manager.set(std.heap.page_allocator);
-}
-
-pub fn setGlobalAllocator(new_allocator: std.mem.Allocator) void {
-    setAllocator(new_allocator);
-}
-
-pub fn resetGlobalAllocator() void {
-    resetAllocator();
-}
-
 pub fn globalAllocator() std.mem.Allocator {
     return global_allocator.globalAllocator();
 }
 
 /// Override only short-lived conversion/operation allocations.
-/// This is mainly useful for scoped allocator tests; applications should use setAllocator.
+/// This is mainly useful for scoped allocator tests; applications should use a
+/// root `napi_allocator` declaration instead.
 pub fn setOperationAllocator(new_allocator: std.mem.Allocator) void {
     global_allocator.global_manager.set(new_allocator);
 }
 
 pub fn resetOperationAllocator() void {
-    global_allocator.global_manager.set(std.heap.page_allocator);
+    global_allocator.global_manager.set(global_allocator.defaultAllocator());
 }
 
 pub fn AsyncContext(comptime Event: type) type {