interface.zig

const builtin = @import("builtin");
const TypeId = builtin.TypeId;
const TypeInfo = builtin.TypeInfo;
const std = @import("std");
const mem = std.mem;
const assert = std.debug.assert;

// @TODO Add const SelfType variant, perhaps SelfType(bool)? (to actually distinguish between const and non const methods)

pub const SelfType = *@OpaqueType();

fn is_vtable(comptime T: type) bool {
    comptime {
        const info = @typeInfo(T);

        if (TypeId(info) != TypeId.Struct)
            return false;

        for (info.Struct.fields) |*field| {
            var field_type = field.field_type;
            var field_type_info = @typeInfo(field_type);

            if (!mem.eql(u8, "Fn", field.name[field.name.len - 2..]))
                return false;

            if (TypeId(field_type_info) == TypeId.Optional) {
                field_type = field_type_info.Optional.child;
                field_type_info = @typeInfo(field_type);
            }

            if (TypeId(field_type_info) != TypeId.Fn)
                return false;

            if (field_type_info.Fn.is_generic or field_type_info.Fn.is_var_args)
                return false;

            if (field_type_info.Fn.args.len == 0)
                return false;

            if (field_type_info.Fn.args[0].arg_type != SelfType)
                return false;
        }

        return true;
    }
}

fn vtable_has_method(comptime VTableType: type, comptime name: []const u8, is_optional: ?*bool) bool {
    comptime {
        const info = @typeInfo(VTableType);

        for (info.Struct.fields) |*field| {
            if (mem.eql(u8, name, field.name[0..field.name.len - 2])) {
                if (is_optional) |io| {
                    const fn_info = @typeInfo(field.field_type);
                    io.* = TypeId(fn_info) == TypeId.Optional;
                }

                return true;
            }
        }

        return false;
    }
}

fn type_has_self_method(comptime T: type, comptime name: []const u8, is_const: ?*bool) bool {
    comptime {
        const defs = switch (@typeInfo(T)) {
            TypeId.Struct => |*s| s.defs,
            TypeId.Enum => |*e| e.defs,
            TypeId.Union => |*u| u.defs,
            else => return false,
        };

        for (defs) |*def| {
            if (mem.eql(u8, name, def.name)) {
                switch (def.data) {
                    TypeInfo.Definition.Data.Fn => |*fn_def| {
                        const fn_info = @typeInfo(fn_def.fn_type);

                        if (fn_info.Fn.args.len == 0)
                            return false;

                        if (fn_info.Fn.args[0].arg_type == *T) {
                            if (is_const) |ic| {
                                ic.* = false;
                            }
                            return true;
                        } else if (fn_info.Fn.args[0].arg_type == *const T) {
                            if (is_const) |ic| {
                                ic.* = true;
                            }
                            return true;
                        } else {
                            return false;
                        }
                    },
                    else => return false,
                }
            }
        }

        return false;
    }
}

fn pull_fn(comptime fn_name: []const u8, comptime FnType: type, comptime ImplType: type,
           comptime fn_info: *const TypeInfo.Fn, comptime is_const: bool) FnType {
    comptime {
        const zero_bit_self = @sizeOf(ImplType) == 0;
        const self_ptr_type = if (zero_bit_self) b: {
            break :b if (is_const) *const ImplType else *ImplType;
        } else b: {
            break :b if (is_const) *align(1) const ImplType else *align(1) ImplType;
        };

        return switch (fn_info.args.len - 1) {
            0 => struct {
                fn func(self_erased: SelfType) fn_info.return_type {
                    var self = if (!zero_bit_self) @alignCast(@alignOf(ImplType), @ptrCast(self_ptr_type, self_erased)) else (self_ptr_type)(undefined);
                    return @inlineCall(@field(self, fn_name));
                }
            }.func,
            1 => struct {
                fn func(self_erased: SelfType, arg0: fn_info.args[1].arg_type) fn_info.return_type {
                    var self = if (!zero_bit_self) @alignCast(@alignOf(ImplType), @ptrCast(self_ptr_type, self_erased)) else (self_ptr_type)(undefined);
                    return @inlineCall(@field(self, fn_name), arg0);
                }
            }.func,
            2 => struct {
                fn func(self_erased: SelfType, arg0: fn_info.args[1].arg_type,
                        arg1: fn_info.args[2].arg_type) fn_info.return_type {
                    var self = if (!zero_bit_self) @alignCast(@alignOf(ImplType), @ptrCast(self_ptr_type, self_erased)) else (self_ptr_type)(undefined);
                    return @inlineCall(@field(self, fn_name), arg0, arg1);
                }
            }.func,
            else => @compileError("Unsupported number of arguments."),
        };
    }
}

fn make_vtable(comptime VTableType: type, comptime ImplType: type) VTableType {
    var vtable: VTableType = undefined;
    comptime {
        const vtable_info = @typeInfo(VTableType);

        for (vtable_info.Struct.fields) |*field| {
            const fn_name = field.name[0..field.name.len - 2];

            var fn_type = field.field_type;
            var fn_info = @typeInfo(field.field_type);
            const is_nullable = TypeId(fn_info) == TypeId.Optional;
            if (is_nullable) {
                fn_type = fn_info.Optional.child;
                fn_info = @typeInfo(fn_type);
            }

            var is_const = false;
            const has_method = type_has_self_method(ImplType, fn_name, &is_const);

            if (!has_method and !is_nullable) {
                @compileError("Could not make vtable (field missing).");
            }

            // If we have no method, initialize the func ptr to null.
            if (!has_method) {
                @field(vtable, field.name) = null;
                continue;
            }

            @field(vtable, field.name) = pull_fn(fn_name, fn_type, ImplType, &fn_info.Fn, is_const);
        }
    }
    return vtable;
}

fn unwrap(comptime T: type) type {
    const info = @typeInfo(T);

    if (TypeId(info) == TypeId.Pointer) {
        return info.Pointer.child;
    }

    return T;
}

pub fn sbo_storage(comptime BuffSize: usize) type {
    return struct {
        const Self = this;

        data: extern union {
            small_buffer: packed struct {
                mem: [BuffSize - 1]u8,
                flag_byte: u8,
            },
            heap_ptr: packed struct {
                ptr: *u8,
                alloc: *mem.Allocator,
                unused_mem: [BuffSize - @sizeOf(*u8) - @sizeOf(*mem.Allocator) - 1]u8,
                flag_byte: u8,
            },
        },

        fn is_stored_inline(self: *const Self) bool {
            return self.data.small_buffer.flag_byte != 0;
        }

        fn init(obj: var, args: ...) Self {
            const ImplType = unwrap(@typeOf(obj));
            const ImplSize = @sizeOf(ImplType);

            var self: Self = undefined;
            comptime assert(@sizeOf(@typeOf(self.data)) == BuffSize * @sizeOf(u8));
        
            if (ImplSize >= BuffSize) {
                if (args.len != 1) {
                    @compileError("sbo_storage requires exactly one argument (&mem.Allocator) to be passed");
                }

                self.data.heap_ptr.alloc = args[0];
                self.data.heap_ptr.ptr = @ptrCast(*u8, try args[0].create(ImplType));

                mem.copy(u8, self.data.heap_ptr.ptr[0..ImplSize], @ptrCast(*const u8, obj)[0..ImplSize]);
                self.data.heap_ptr.flag_byte = 0;
                return self;
            }

            self.data.small_buffer.flag_byte = 1;
            if (ImplSize > 0) {
                mem.copy(u8, self.data.small_buffer.mem[0..], @ptrCast([*]const u8, obj)[0..ImplSize]);
            }

            return self;
        }

        fn erased_ptr(self: *Self) SelfType {
            if (!self.is_stored_inline()) {
                return @ptrCast(SelfType, self.data.heap_ptr.ptr);
            }

            return @ptrCast(SelfType, &self.data.small_buffer.mem[0]);
        }
    
        fn deinit(self: *Self) void {
            if (!self.is_stored_inline()) {
                self.data.heap_ptr.alloc.destroy(self.data.heap_ptr.ptr);
            }
        }
    };
}

pub const non_owning_storage = struct {
    const Self = this;

    data: SelfType,

    fn init(obj: var, args: ...) Self {
        if (args.len > 0) {
            @compileError("Non owning storage expects no extra initialization arguments.");
        }

        if (@sizeOf(@typeOf(obj)) == 0) {
            return undefined;
        }

        return Self { .data=@ptrCast(SelfType, obj), };
    }

    fn erased_ptr(self: *Self) SelfType {
        return self.data;
    }

    fn deinit(self: *Self) void {}
};

fn vtable_return_type(comptime VTableType: type, name: []const u8) type {
    comptime {
        var is_optional = true;
        if (!vtable_has_method(VTableType, name, &is_optional)) {
            @compileError("Invalid interface method call.");
        }

        const info = @typeInfo(VTableType);

        for (info.Struct.fields) |*field| {
            if (mem.eql(u8, name, field.name[0..field.name.len - 2])) {
                const ret_type = if (is_optional) @typeInfo(@typeInfo(field.field_type).Optional.child).Fn.return_type
                                 else @typeInfo(field.field_type).Fn.return_type;

                if (is_optional) {
                    return ?ret_type;
                } else {
                    return ret_type;
                }
            }
        }

        unreachable;
    }
}

pub fn Interface(comptime VTableType: type, comptime StoragePolicy: type) type {
    if (!is_vtable(VTableType)) {
        @compileError("Invalid vtable type.");
    }

    return struct {
        const Self = this;

        vtable: *const VTableType,
        storage: StoragePolicy,

        pub fn init(obj: var, args: ...) Self {
            const ImplType = unwrap(@typeOf(obj));

            return Self {
                .vtable = &comptime make_vtable(VTableType, ImplType),
                .storage = StoragePolicy.init(obj, args),
            };
        }

        pub fn deinit(self: *Self) void {
            comptime var is_optional = true;
            // Call our virtual destructor, if we have one.
            if (vtable_has_method(VTableType, "deinit", &is_optional)) {
                if (is_optional) {
                    if (self.vtable.deinitFn) |f| {
                        f();
                    }
                } else {
                    self.vtable.deinitFn();
                }
            }
            // Release memory held by the storage.
            self.storage.deinit();
        }
    
        pub fn call(self: *Self, comptime name: []const u8, args: ...) vtable_return_type(VTableType, name) {
            comptime var is_optional = true;
            comptime assert(vtable_has_method(VTableType, name, &is_optional));

            const fn_ptr = if (is_optional) blk: {
                const val = @field(self.vtable, name ++ "Fn");
                if (val) |v| break :blk v;

                return null;
            } else blk: {
                break :blk @field(self.vtable, name ++ "Fn");
            };

            const erased_obj_ptr = self.storage.erased_ptr();

            switch (args.len) {
                0 => return fn_ptr(erased_obj_ptr),
                1 => return fn_ptr(erased_obj_ptr, args[0]),
                2 => return fn_ptr(erased_obj_ptr, args[0], args[1]),
                else => @compileError("Unsupported number of arguments."),
            }
        }
    };
}

test "is_vtable" {
    const VTableType = struct {
        printFn: fn(SelfType)void,
        pushFn: fn(SelfType, u8)bool,
    };

    assert(is_vtable(VTableType));

    const AlmostVTable = struct {
        print: fn(SelfType)void,
    };

    assert(!is_vtable(AlmostVTable));
}

test "simple interface" {
    const MyVtable = struct {
        readFn: fn(SelfType, []u8)void,
    };

    // In a world with @reify, this whole struct could be generated.
    const MyInterface = struct {
        const Self = this;
        const IFace = Interface(MyVtable, sbo_storage(24));

        iface: IFace,

        fn init(obj: var, args: ...) Self {
            return Self {
                .iface = IFace.init(obj, args),
            };
        }

        fn read(self: *Self, arg: []u8) void {
            return self.iface.call("read", arg);
        }

        fn deinit(self: *Self) void {
            self.iface.deinit();
        }
    };

    const FooReader = struct {
        const Self = this;

        dummy: u8,

        fn read(self: *const Self, data: []u8) void {
            for (data) |*c| {
                c.* = 42;
            }
        }
    };

    var instance = MyInterface.init(FooReader {.dummy = 0});

    var data: [1024]u8 = undefined;
    instance.read(data[0..]);
    var ok = true;
    for (data) |d| {
        if (d != 42)
            ok = false;
    }

    assert(ok);
}

test "non owning zero bit implementation type" {
    const VTable = struct {
        fooFn: fn(SelfType, u64) u64,
    };

    const Fooer = struct {
        const Self = this;
        const IFace = Interface(VTable, non_owning_storage);

        iface: IFace,

        fn init(obj: var, args: ...) Self {
            return Self {
                .iface = IFace.init(obj, args),
            };
        }

        fn foo(self: *Self, arg: u64) u64 {
            return self.iface.call("foo", arg);
        }

        fn deinit(self: *Self) void {
            self.iface.deinit();
        }
    };

    const Doubler = struct {
        const Self = this;

        fn foo(self: *const Self, x: u64) u64 {
            return x * 2;
        }
    };

    var instance = Fooer.init(Doubler {});
    assert(instance.foo(21) == 42);
}

test "sbo storage zero bit implementation type" {
    const VTable = struct {
        fooFn: fn(SelfType, u64) u64,
    };

    const Fooer = struct {
        const Self = this;
        const IFace = Interface(VTable, sbo_storage(24));

        iface: IFace,

        fn init(obj: var, args: ...) Self {
            return Self {
                .iface = IFace.init(obj, args),
            };
        }

        fn foo(self: *Self, arg: u64) u64 {
            return self.iface.call("foo", arg);
        }

        fn deinit(self: *Self) void {
            self.iface.deinit();
        }
    };

    const Doubler = struct {
        const Self = this;

        fn foo(self: *const Self, x: u64) u64 {
            return x * 2;
        }
    };

    var instance = Fooer.init(Doubler {});
    assert(instance.foo(21) == 42);
}

test "multiple virtual methods" {
    const VTable = struct {
        // @TODO Nullable functions lead to LLVM IR error atm.
        fooFn: fn(SelfType) u8,
        barFn: fn(SelfType, usize) u64,
    };

    const Tester = struct {
        const Self = this;
        const IFace = Interface(VTable, non_owning_storage);

        iface: IFace,

        fn init(obj: var, args: ...) Self {
            return Self {
                .iface = IFace.init(obj, args),
            };
        }

        fn foo(self: *Self) u8 {
            return self.iface.call("foo");
        }

        fn bar(self: *Self, arg: usize) u64 {
            return self.iface.call("bar", arg);
        }

        fn deinit(self: *Self) void {
            self.iface.deinit();
        }
    };

    const TheBarOnlyTester = struct {
        const Self = this;

        data: []const u64,

        fn bar(self: *const Self, index: usize) u64 {
            return self.data[index];
        }
    };

    const TheFooBarTester = struct {
        const Self = this;

        fn foo(self: *const Self) u8 {
            return 2;
        }

        fn bar(self: *const Self, arg: usize) u64 {
            return 21;
        }
    };

    // const some_data = []u64 { 0, 1, 0, 2, 0, 3, };
    // // Passing this without the ref makes it const.
    // var bar_only_instance = Tester.init(&(TheBarOnlyTester { .data=some_data[0..], }));
    // 
    // assert(bar_only_instance.foo() == null);
    // assert(bar_only_instance.bar(0) == 0 and bar_only_instance.bar(5) == 3);

    var complete_instance = Tester.init(&(TheFooBarTester {}));
    assert(complete_instance.foo() * complete_instance.bar(0) == 42);
}

Thanks for fixing this up! I tried to add a second method to MyInterface, but it was failing. It seems like make_vtable doesn't work if there are multiple methods? It expects the subsequent methods to have the same function signature as the first one. Maybe the closures around line 150 aren't working.

Updated the code, atm nullable methods are broken (invalid LLVM IR is generated) though.