Revisiting alignedAlloc and alignedFree

[Sze]

In the process of writing an answer for this topic about free with alignment, I was revisiting this older topic and trying to adapt @kALLEBALIK’s solution to recent Zig:

U8 alignment and storing data types as bytes

Seems like my “hacky solution” is quite ok. I tested my solution in godbolt to see the generated code. Since zig doesn’t handle alignment over the minimum page size we can use std.mem.page_size as the maximum alignment. // The Zig Allocator interface is not intended to solve alignments beyond // the minimum OS page size. For these use cases, the caller must use OS // APIs directly. Link to comment code (godbolt test) export fn alignedAlloc(alignment: u32) usize { const Log2PageSizeType…

And here is the variation I created (based on both topics):

const std = @import("std");
const builtin = @import("builtin");

pub fn aligned(alignment: u29) std.math.IntFittingRange(0, std.math.log2(std.heap.page_size_max)) {
    std.debug.assert(std.mem.isValidAlign(alignment));
    std.debug.assert(alignment > 0);
    std.debug.assert(alignment <= std.heap.page_size_max);
    return @intCast(std.math.log2(alignment));
}

pub fn alignedAlloc(allocator: std.mem.Allocator, T: type, alignment: u29, n: usize) ![]u8 {
    const last = comptime aligned(std.heap.page_size_max);
    switch (aligned(alignment)) {
        inline 0...last => |a| {
            return try allocator.alignedAlloc(T, 1 << a, n);
        },
        else => unreachable,
    }
}

pub fn alignedFree(allocator: std.mem.Allocator, alignment: u29, memory: anytype) void {
    const last = comptime aligned(std.heap.page_size_max);
    switch (aligned(alignment)) {
        inline 0...last => |a| {
            const T = @typeInfo(@TypeOf(memory)).pointer.child;
            return allocator.free(@as([]align(1 << a) T, @alignCast(memory)));
        },
        else => unreachable,
    }
}

pub inline fn pageAlign() u29 {
    return @intCast(std.heap.pageSize());
}

// pub const std_options: std.Options = .{
//     .page_size_max = 65536,
// };

const MyStruct = struct {
    allocator: std.mem.Allocator,
    bytes: []u8,

    pub fn init(allocator: std.mem.Allocator, bytes: []const u8) !@This() {
        const buffer = try alignedAlloc(allocator, u8, pageAlign(), bytes.len);
        @memcpy(buffer, bytes);
        return .{
            .allocator = allocator,
            .bytes = buffer,
        };
    }

    pub fn deinit(self: *@This()) void {
        alignedFree(self.allocator, pageAlign(), self.bytes);
    }
};

export fn example() [*]u8 {
    const allocator = std.heap.page_allocator;

    var myStruct = MyStruct.init(allocator, &[_]u8{ 1, 2, 3, 4 }) catch unreachable;
    defer myStruct.deinit();
    return myStruct.bytes.ptr;
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var myStruct = try MyStruct.init(allocator, &[_]u8{ 1, 2, 3, 4 });
    defer myStruct.deinit();
}

I tried to compare it a bit with the godbolt output of the original, but I don’t have enough assembly experience, to be able to tell whether the code produced by my version is good.

On a high level it seems good to me that my version directly switches on the powers of 2, but I don’t know whether that matters in a practical way. (Quite a few versions seem to result in similar code being generated and it is easy to get lost between different targets and build modes)

Any ideas about changing the code further?