vtable.zig

const std = @import("std");

/// Given an example struct with methods, create a function that produces vtables for similarly shaped structs.
/// A vtable contains function pointers to all the struct methods, with the struct type erased.
/// All structs should have an init function that takes an allocator as the first parameter and stores it in a field called "allocator".
/// At most 2 parameters are supported for methods (in addition to the initial self parameter).
pub fn VTableShape(comptime Template: type) type {
    var fields = @typeInfo(struct { //
        deinit: *const fn (*anyopaque) void,
    }).Struct.fields;

    // Loop through method declarations for the template struct and define a vtable type with matching fields,
    // the self parameter type changed to *anyopaque and error union type changed to anyerror.
    comptime for (std.meta.declarations(Template)) |decl| {
        const name = decl.name;
        if (std.mem.eql(u8, name, "deinit")) continue;

        const info = @typeInfo(@TypeOf(@field(Template, name)));
        if (info != .Fn) continue;

        var types_erased = info.Fn;
        const return_type = @typeInfo(types_erased.return_type.?);

        if (std.mem.eql(u8, name, "init")) {
            // Change init to return *anyopaque instead of actual type.
            // An error may always be returned, because init through vtable allocates an opaque object in heap.
            types_erased.return_type = anyerror!*anyopaque;
        } else {
            // Replace error union return type with anyerror.
            if (return_type == .ErrorUnion) {
                types_erased.return_type = anyerror!return_type.ErrorUnion.payload;
            }

            // Replace first parameter with *anyopaque.
            var param = types_erased.params[0];
            param.type = *anyopaque;
            types_erased.params = &[_]std.builtin.Type.Fn.Param{param} ++ types_erased.params[1..];
        }

        const Method = *const @Type(.{ .Fn = types_erased });

        // Add field to vtable type.
        fields = fields ++ &[_]std.builtin.Type.StructField{.{ //
            .name = name,
            .type = Method,
            .default_value = null,
            .is_comptime = false,
            .alignment = @alignOf(Method),
        }};
    };

    // Create vtable type.
    var struct_info = @typeInfo(struct {});
    struct_info.Struct.fields = fields;

    const VTable = @Type(struct_info);

    return struct { //
        /// Create a vtable for a class, compatible with given template.
        pub inline fn createVTable(comptime Class: type) VTable {
            var vtable: VTable = undefined;

            if (@hasField(VTable, "init")) {
                const info = @typeInfo(@TypeOf(Class.init)).Fn;
                const params = info.params;
                const has_error = @typeInfo(info.return_type.?) == .ErrorUnion;

                vtable.init = &switch (params.len) {
                    1 => struct {
                        pub fn call(allocator: std.mem.Allocator) !*anyopaque {
                            var self = try allocator.create(Class);
                            errdefer allocator.destroy(self);

                            self.* = if (has_error) try Class.init(allocator) else Class.init(allocator);

                            return @alignCast(@ptrCast(self));
                        }
                    },
                    2 => struct {
                        pub fn call(allocator: std.mem.Allocator, arg_1: params[1].type.?) !*anyopaque {
                            var self = try allocator.create(Class);
                            errdefer allocator.destroy(self);

                            self.* = if (has_error) try Class.init(allocator, arg_1) else Class.init(allocator, arg_1);

                            return @alignCast(@ptrCast(self));
                        }
                    },
                    3 => struct {
                        pub fn call(allocator: std.mem.Allocator, arg_1: params[1].type.?, arg_2: params[2].type.?) !*anyopaque {
                            var self = try allocator.create(Class);
                            errdefer allocator.destroy(self);

                            self.* = if (has_error) try Class.init(allocator, arg_1, arg_2) else Class.init(allocator, arg_1, arg_2);

                            return @alignCast(@ptrCast(self));
                        }
                    },
                    else => @compileError("Unsupported number of arguments for init"),
                }.call;
            }

            vtable.deinit = struct {
                pub fn call(ctx: *anyopaque) void {
                    const self: *Class = @ptrCast(@alignCast(ctx));

                    if (@hasDecl(Class, "deinit")) {
                        self.deinit();
                    }

                    if (@hasField(Class, "allocator")) {
                        self.allocator.destroy(self);
                    }
                }
            }.call;

            inline for (std.meta.fields(VTable)) |field| {
                const name = field.name;
                comptime if (std.mem.eql(u8, name, "init") or std.mem.eql(u8, name, "deinit")) continue;

                const info = @typeInfo(@typeInfo(@TypeOf(@field(vtable, name))).Pointer.child).Fn;
                const params = info.params;
                const Return = info.return_type.?;
                const conv = info.calling_convention;

                @field(vtable, name) = &switch (params.len) {
                    1 => struct { //
                        pub fn call(ctx: *anyopaque) callconv(conv) Return {
                            const self: *Class = @ptrCast(@alignCast(ctx));
                            return @call(.always_inline, @field(Class, name), .{self});
                        }
                    },
                    2 => struct { //
                        pub fn call(ctx: *anyopaque, arg_1: params[1].type.?) callconv(conv) Return {
                            const self: *Class = @ptrCast(@alignCast(ctx));
                            return @call(.always_inline, @field(Class, name), .{ self, arg_1 });
                        }
                    },
                    3 => struct { //
                        pub fn call(ctx: *anyopaque, arg_1: params[1].type.?, arg_2: params[2].type.?) callconv(conv) Return {
                            const self: *Class = @ptrCast(@alignCast(ctx));
                            return @call(.always_inline, @field(Class, name), .{ self, arg_1, arg_2 });
                        }
                    },
                    else => @compileError("Unsupported number of arguments for " ++ name),
                }.call;
            }

            return vtable;
        }
    };
}

test "VTableShape" {
    const Base = struct {
        const Self = @This();

        allocator: std.mem.Allocator,
        value: *u8,

        pub fn init(allocator: std.mem.Allocator) !Self {
            const ptr = try allocator.create(u8);
            ptr.* = 1;

            return .{ .allocator = allocator, .value = ptr };
        }

        pub fn query(self: *const Self) u8 {
            return self.value.*;
        }

        pub fn deinit(self: *Self) void {
            self.allocator.destroy(self.value);
        }
    };

    const Other = struct {
        const Self = @This();

        allocator: std.mem.Allocator,

        pub fn init(allocator: std.mem.Allocator) Self {
            return .{ .allocator = allocator };
        }

        pub fn query(_: *const Self) u8 {
            return 2;
        }
    };

    const Test = struct {
        fn testVTable(comptime Creator: type) !void {
            const allocator = std.testing.allocator;

            const createVTable = Creator.createVTable;
            const vtable_base = createVTable(Base);
            const vtable_other = createVTable(Other);

            const ptr_base: *anyopaque = try vtable_base.init(allocator);
            const ptr_other: *anyopaque = try vtable_other.init(allocator);

            try std.testing.expectEqual(@as(u8, 1), vtable_base.query(ptr_base));
            try std.testing.expectEqual(@as(u8, 2), vtable_other.query(ptr_other));

            vtable_base.deinit(ptr_base);
            vtable_other.deinit(ptr_other);
        }
    };

    try Test.testVTable(VTableShape(Base));
    try Test.testVTable(VTableShape(Other));
}