How Can I Specify an Allocator at Comptime in Zig?

549 Views Asked by At

In Zig, I want specify an allocator at comptime in order to remove the added overhead of passing around an allocator to each function that requires it. In the following MWE, I have a type generator that takes in a comptime ArenaAllocator and returns a type with that allocator as a "static" member. The type is instantiated, called MyType, and then an object of MyType is initialized using the previously specified allocator.

const std = @import("std");

fn MyTypeGen(comptime arena: *std.heap.ArenaAllocator) type {
    return struct {
        var allocator = arena.allocator();

        data: []const u8,

        pub fn init(data: []const u8) !@This() {
            const temp = try allocator.alloc(u8, data.len);
            @memcpy(temp, data);

            return .{ .data = temp };
        }
    };
}

pub fn main() !void {
    comptime var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    const MyType = MyTypeGen(&arena);

    const val = try MyType.init(&[_]u8{ 1, 2, 3 });

    std.debug.print("{}\n", .{val});
}

The problem is the above code fails with a segmenation fault:

Segmentation fault at address 0x7ff63e46b628
C:\Program Files\zig\lib\std\heap\arena_allocator.zig:171:39: 0x7ff63e3f2788 in createNode (Calculator.exe.obj)
        self.state.buffer_list.prepend(buf_node);
                                      ^
C:\Program Files\zig\lib\std\heap\arena_allocator.zig:184:29: 0x7ff63e3f216a in alloc (Calculator.exe.obj)
            (self.createNode(0, n + ptr_align) orelse return null);
                            ^
C:\Program Files\zig\lib\std\mem\Allocator.zig:225:53: 0x7ff63e4242dd in allocBytesWithAlignment__anon_6797 (Calculator.exe.obj)
    const byte_ptr = self.rawAlloc(byte_count, log2a(alignment), return_address) orelse return Error.OutOfMemory;
                                                    ^
C:\Program Files\zig\lib\std\mem\Allocator.zig:211:40: 0x7ff63e3f2e37 in allocWithSizeAndAlignment__anon_3321 (Calculator.exe.obj)
    return self.allocBytesWithAlignment(alignment, byte_count, return_address);
                                       ^
C:\Program Files\zig\lib\std\mem\Allocator.zig:129:41: 0x7ff63e3f11f3 in alloc__anon_3262 (Calculator.exe.obj)
    return self.allocAdvancedWithRetAddr(T, null, n, @returnAddress());
                                        ^
C:\...\src\main.zig:10:45: 0x7ff63e3f1044 in init (Calculator.exe.obj)
            const temp = try allocator.alloc(u8, data.len);
                                            ^
C:\...\src\main.zig:24:32: 0x7ff63e3f1486 in main (Program.exe.obj)
    const val = try MyType.init(&[_]u8{ 1, 2, 3 });
                               ^
C:\Program Files\zig\lib\std\start.zig:348:65: 0x7ff63e3f175c in WinStartup (Program.exe.obj)
    std.os.windows.kernel32.ExitProcess(initEventLoopAndCallMain());
                                                                ^
???:?:?: 0x7ffee8867343 in ??? (KERNEL32.DLL)
???:?:?: 0x7ffeea8426b0 in ??? (ntdll.dll)
run Program: error: the following command exited with error code 3:
C:\...\zig-out\bin\Program.exe
Build Summary: 3/5 steps succeeded; 1 failed (disable with --summary none)
run transitive failure
+- run Program failure
error: the following build command failed with exit code 1:
C:\...\zig-cache\o\03618dc1d07de7a2b7ce2381e69121e1\build.exe C:\Program Files\zig\zig.exe C:\... C:\...\zig-cache C:\...\AppData\Local\zig run

It appears the issue is that Windows specifies the heap allocator, which backs the ArenaAllocator, at runtime meaning that comptime var arena is initialized erroneously. Zig must just not be complaining about the ostensibly runtime dependency at compile time, hence the segfault.

How can I initialize the arena allocator at comptime successfully?

Alternate methods for specifying an allocator at comptime which don't require initializing an allocator in the above manner are welcome, too.

1

There are 1 best solutions below

0
On BEST ANSWER

Wrapping the allocator in its own global type seems to work.

const std = @import("std");

const GlobalArena = struct {
    var global_arena: std.heap.ArenaAllocator = undefined;

    fn init(arena: std.heap.ArenaAllocator) void {
        global_arena = arena;
    }

    fn deinit() void {
        global_arena.deinit();
        global_arena = undefined;
    }

    fn allocator() std.mem.Allocator {
        return global_arena.allocator();
    }
};

fn MyTypeGen(comptime GlobalAllocator: type) type {
    return struct {
        var allocator = GlobalAllocator.allocator();

        data: []const u8,

        pub fn init(data: []const u8) !@This() {
            const temp = try allocator.alloc(u8, data.len);
            @memcpy(temp, data);

            return .{ .data = temp };
        }
    };
}

pub fn main() !void {
    GlobalArena.init(std.heap.ArenaAllocator.init(std.heap.page_allocator));
    defer GlobalArena.deinit();

    const MyType = MyTypeGen(GlobalArena);

    const val = try MyType.init(&[_]u8{ 1, 2, 3 });

    std.debug.print("{}\n", .{val});
}
$ zig version
0.12.0-dev.587+eb072fa52

$ zig build run
main.MyTypeGen(main.GlobalArena){ .data = { 1, 2, 3 } }