Simple reflection system for zig

rtti.zig

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

pub const Reflection = struct {
    internal: []const type,
    types: []const Type,
    lookup: std.StaticStringMap(LocalId), // todo: use HashMap for this ?

    pub const LocalId = u32;
    pub const invalid_id = std.math.maxInt(LocalId);

    pub const Int = struct {
        signedness: std.builtin.Signedness,
        size: u32,
    };

    pub const Float = struct {
        size: u32,
    };

    pub const Vec = struct {
        child: LocalId,
        len: u16,
        element_size: u16,
    };

    pub const Field = struct {
        type: LocalId,
        index: u16,
        offset: u16,
        // bit_offset: u16,
        // bit_size: u16,
    };

    pub const Map = struct {
        name: [*:0]const u8,
        fields: *const std.StaticStringMap(Field),
    };

    pub const Ptr = struct {
        child: LocalId,
    };

    pub const Slice = struct {
        child: LocalId,
        // ptr: u32,
        // len: u32,
    };

    pub const Opt = struct {
        child: LocalId,
        // flag: u32,
    };

    pub const Type = union(enum) {
        @"opaque": void,
        bool: void,
        int: Int,
        float: Float,
        vec: Vec,
        map: Map,
        ptr: Ptr,
        slice: Slice,
        opt: Opt,

        // pub fn format(self: *const @This(), _: []const u8, _: std.fmt.FormatOptions, writer: anytype) @TypeOf(writer).Error!void {
        //     var buf: [64]u8 = undefined;
        //     switch (self.*) {
        //         .bool => |_| try writer.writeAll("bool"),
        //         .int => |i| try writer.writeAll(
        //             if (i.signedness == .unsigned)
        //                 try std.fmt.bufPrint(&buf, "u{d}", .{i.size * 8})
        //             else
        //                 try std.fmt.bufPrint(&buf, "i{d}", .{i.size * 8}),
        //         ),
        //         .float => |f| try writer.writeAll(try std.fmt.bufPrint(&buf, "f{d}", .{f.size * 8})),
        //         .vec => |v| try writer.writeAll(try std.fmt.bufPrint(&buf, "@Vector({d}, TypeId({d}))", .{ v.len, v.child })),
        //         .map => |m| try writer.writeAll(m.name),
        //         .ptr => |p| try writer.writeAll(try std.fmt.bufPrint(&buf, "?[*]TypeId({d})", .{p.child})),
        //         .slice => |s| try writer.writeAll(try std.fmt.bufPrint(&buf, "?[]TypeId({d})", .{s.child})),
        //         .opt => |s| try writer.writeAll(try std.fmt.bufPrint(&buf, "?TypeId({d})", .{s.child})),
        //         else => try writer.writeAll("opaque"),
        //     }
        // }
    };

    // pub const invalid = std.math.maxInt(usize);

    pub fn init(comptime base_types: []const type) @This() {
        @setEvalBranchQuota(10_000);

        const internal = blk: {
            var tmp: [256]type = undefined;
            var i: usize = 0;
            var len: usize = base_types.len;
            @memcpy(tmp[0..len], base_types);
            while (i < len) : (i += 1) {
                const child = switch (@typeInfo(tmp[i])) {
                    .@"struct" => |s| {
                        if (@hasDecl(tmp[i], "no_reflect") and tmp[i].no_reflect) {
                            continue;
                        }
                        for (s.fields) |field| {
                            if (std.mem.indexOfScalar(type, tmp[0..len], field.type) != null) continue;
                            tmp[len] = field.type;
                            len += 1;
                        }
                        continue;
                    },
                    .array => |a| a.child,
                    .vector => |v| v.child,
                    .pointer => |p| p.child,
                    .optional => |o| o.child,
                    else => continue,
                };
                // default is to add the child type
                if (std.mem.indexOfScalar(type, tmp[0..len], child) != null) continue;
                tmp[len] = child;
                len += 1;
            }
            const out = tmp[0..len].*;
            break :blk &out;
        };

        const types = blk: {
            var tmp: [internal.len]Type = undefined;
            for (internal, 0..) |T, t| {
                switch (@typeInfo(T)) {
                    .bool => |_| tmp[t] = .{ .bool = undefined },
                    .int => |int| tmp[t] = .{ .int = .{ .signedness = int.signedness, .size = @sizeOf(T) } },
                    .@"enum" => |e| tmp[t] = .{ .int = .{ .signedness = .unsigned, .size = @sizeOf(e.tag_type) } },
                    .float => |_| tmp[t] = .{ .float = .{ .size = @sizeOf(T) } },
                    .@"struct" => |s| {
                        if (@hasDecl(T, "no_reflect") and T.no_reflect) {
                            tmp[t] = .{ .@"opaque" = undefined };
                            continue;
                        }
                        var kv: [s.fields.len]struct { []const u8, Field } = undefined;
                        for (s.fields, 0..) |field, f| {
                            kv[f] = .{
                                field.name,
                                Field{
                                    .type = std.mem.indexOfScalar(type, internal, field.type) orelse unreachable,
                                    .index = f,
                                    .offset = @offsetOf(T, field.name),
                                    // .bit_offset = @bitOffsetOf(T, field.name),
                                    // .bit_size = @bitSizeOf(field.type),
                                },
                            };
                        }
                        tmp[t] = .{ .map = .{
                            .name = typeName(T),
                            .fields = &std.StaticStringMap(Field).initComptime(kv),
                        } };
                    },
                    .vector => |v| {
                        tmp[t] = .{ .vec = .{
                            .child = std.mem.indexOfScalar(type, internal, v.child) orelse unreachable,
                            .len = v.len,
                            .element_size = @sizeOf(v.child),
                        } };
                    },
                    .array => |a| {
                        tmp[t] = .{ .vec = .{
                            .child = std.mem.indexOfScalar(type, internal, a.child) orelse unreachable,
                            .len = a.len,
                            .element_size = @sizeOf(a.child),
                        } };
                    },
                    .pointer => |p| {
                        const child = std.mem.indexOfScalar(type, internal, p.child) orelse unreachable;
                        switch (p.size) {
                            .One, .Many, .C => tmp[t] = .{ .ptr = .{ .child = child } },
                            .Slice => tmp[t] = .{
                                .slice = .{
                                    .child = child,
                                    // todo: not guaranteed to be correct
                                    // .ptr = 0,
                                    // .len = @sizeOf(usize),
                                },
                            },
                        }
                    },
                    .optional => |o| {
                        switch (@typeInfo(o.child)) {
                            // collapse optional pointers
                            .pointer => |p| {
                                const child = std.mem.indexOfScalar(type, internal, p.child) orelse unreachable;
                                switch (p.size) {
                                    .One, .Many, .C => tmp[t] = .{ .ptr = .{ .child = child } },
                                    .Slice => tmp[t] = .{
                                        .slice = .{
                                            .child = child,
                                            // todo: not guaranteed to be correct
                                            // .ptr = 0,
                                            // .len = @sizeOf(usize),
                                        },
                                    },
                                }
                                continue;
                            },
                            else => {
                                const child = std.mem.indexOfScalar(type, internal, o.child) orelse unreachable;
                                tmp[t] = .{
                                    .opt = .{
                                        .child = child,
                                        // todo: not guaranteed to be correct
                                        // .flag = @sizeOf(o.child),
                                    },
                                };
                            },
                        }
                    },
                    else => tmp[t] = .{ .@"opaque" = undefined },
                }
            }
            const out = tmp;
            break :blk &out;
        };

        const lookup = blk: {
            const KV = struct { []const u8, LocalId };
            var kv: [types.len]KV = undefined;
            var len: usize = 0;
            for (types, 0..) |t, i| {
                switch (t) {
                    .map => |map| {
                        kv[len] = .{ std.mem.span(map.name), i };
                        len += 1;
                    },
                    else => {
                        // ignore any other types
                        // kv[len] = .{ @typeName(internal[i]), i };
                        // len += 1;
                    },
                }
            }
            const out = kv[0..len].*;
            break :blk std.StaticStringMap(LocalId).initComptime(&out);
        };

        return .{
            .internal = internal,
            .types = types,
            .lookup = lookup,
        };
    }

    // /// check assumptions about the types, runs even on ReleaseFast builds
    // pub fn assert(self: *const @This()) !void {
    //     inline for (self.internal, 0..) |T, t| {
    //         switch (@typeInfo(T)) {
    //             .pointer => |p| {
    //                 switch (p.size) {
    //                     .One, .Many, .C => {},
    //                     .Slice => {
    //                         const slice: T = undefined;
    //                         const ptr = @intFromPtr(&slice.ptr) - @intFromPtr(&slice);
    //                         if (ptr != self.types[t].slice.ptr) return error.InvalidSlice;
    //                     },
    //                 }
    //             },
    //             .optional => |o| {
    //                 switch (@typeInfo(o.child)) {
    //                     // collapse optional pointers
    //                     .pointer => |p| {
    //                         switch (p.size) {
    //                             .One, .Many, .C => {},
    //                             .Slice => {
    //                                 const slice: T = undefined;
    //                                 const ptr = @intFromPtr(&slice.ptr) - @intFromPtr(&slice);
    //                                 if (ptr != self.types[t].slice.ptr) return error.InvalidOptSlice;
    //                             },
    //                         }
    //                         continue;
    //                     },
    //                     else => {
    //                         var value: o.child = undefined;
    //                         @memset(std.mem.asBytes(&value), 0);
    //                         const opt: T = value;
    //                         if (std.mem.indexOfScalar(u8, std.mem.asBytes(&opt), 1) != self.types[t].opt.flag) {
    //                             return error.InvalidOpt;
    //                         }
    //                     },
    //                 }
    //             },
    //             else => {},
    //         }
    //     }
    // }

    /// `TypeId` of a comptime type
    pub fn typeId(self: *const @This(), comptime T: type) LocalId {
        return std.mem.indexOfScalar(type, self.internal, T) orelse
            unreachable;
    }

    /// finds a `TypeId` based on type name
    pub fn findId(self: *const @This(), name: []const u8) ?LocalId {
        return self.lookup.get(name);
    }

    pub const Prop = struct {
        type: LocalId = invalid_id,
        offsets: [3]u16 = undefined,
        len: u8 = 0,

        pub fn get(self: *const @This(), comptime T: type, base_ptr: ?[*]u8) ?*T {
            var ptr = base_ptr orelse return null;
            if (self.len > 0) {
                ptr = ptr + self.offsets[0];
            }
            var i: usize = 1;
            while (i < self.len) : (i += 1) {
                ptr = @as(*?[*]u8, @ptrCast(@alignCast(ptr))).* orelse return null;
                ptr = ptr + self.offsets[i];
            }
            return @as(*T, @ptrCast(@alignCast(ptr)));
        }
    };

    pub fn findProp(self: *const @This(), type_id: LocalId, path: []const u8) ?Prop {
        var prop = Prop{
            .type = type_id,
            .offsets = undefined,
            .len = 1,
        };
        prop.offsets[0] = 0;

        var i: usize = 0;
        while (i < path.len) {
            const j: usize = std.mem.indexOfAnyPos(u8, path, i, ".[]") orelse path.len;
            const sub_path = path[i..j];
            while (sub_path.len != 0) {
                switch (self.types[prop.type]) {
                    .vec => |v| {
                        var k: u16 = undefined;
                        if (sub_path.len == 1 and 'w' <= sub_path[0] and sub_path[0] <= 'z') {
                            // support x, y, z, w indexing on vec
                            k = (sub_path[0] -% 'w' -% 1) & 0x03;
                        } else {
                            k = std.fmt.parseInt(u16, sub_path, 10) catch return null;
                        }
                        if (k >= v.len) return null;
                        prop.offsets[prop.len - 1] += v.element_size * k;
                        prop.type = v.child;
                    },
                    .map => |m| {
                        const field = m.fields.get(sub_path) orelse return null;
                        prop.offsets[prop.len - 1] += field.offset;
                        prop.type = field.type;
                    },
                    .ptr => |p| {
                        prop.offsets[prop.len] = 0;
                        prop.len += 1;
                        prop.type = p.child;
                        continue;
                    },
                    else => return null,
                }
                break;
            }
            i = j + 1;
        }

        return prop;
    }
};

pub const TypeId = *const struct {
    _: u8,
};

// full explanation here https://github.com/ziglang/zig/issues/19858#issuecomment-2369861301
pub inline fn typeId(comptime T: type) TypeId {
    return &struct {
        comptime {
            _ = T;
        }
        var id: @typeInfo(TypeId).pointer.child = undefined;
    }.id;
}

pub fn typeName(comptime T: type) [:0]const u8 {
    // todo: isn't working for game.World.Prefab
    // todo: check type @TypeOf(T.name) == [:0]const u8) {
    if (@hasDecl(T, "name")) {
        return T.name;
    }
    return comptime blk: {
        const name = @typeName(T);
        var buffer: [name.len + 1]u8 = undefined;
        var len: usize = 0;
        var prev: usize = 0;
        var cursor: usize = 0;
        while (cursor < name.len) : (cursor += 1) {
            switch (name[cursor]) {
                '.' => prev = cursor + 1,
                '(', ',', ')' => {
                    for (name[prev .. cursor + 1]) |c| {
                        buffer[len] = c;
                        len += 1;
                    }
                    prev = cursor + 1;
                },
                '{' => {
                    cursor += 1;
                    var count: usize = 1;
                    while (cursor < name.len) : (cursor += 1) {
                        switch (name[cursor]) {
                            '{' => count += 1,
                            '}' => {
                                count -= 1;
                                if (count == 0) break;
                            },
                            else => {},
                        }
                    }
                    // note: matches the result if the input param is a function
                    for ("()") |c| {
                        buffer[len] = c;
                        len += 1;
                    }
                    prev = cursor;
                },
                else => {},
            }
        }
        // flush the remainig of the type name
        for (name[prev..@min(name.len, cursor + 1)]) |c| {
            buffer[len] = c;
            len += 1;
        }

        buffer[len] = 0;

        const out = buffer;
        var slice: [:0]const u8 = undefined;
        slice.ptr = @ptrCast(&out);
        slice.len = len;
        break :blk slice;
    };
}

pub fn fnName(comptime f: anytype) [:0]const u8 {
    const Inner = struct {
        fn Dummy(comptime func: anytype) type {
            return struct {
                fn warpper() void {
                    func();
                }
            };
        }
    };
    return comptime blk: {
        const name = @typeName(Inner.Dummy(f));
        const start = (std.mem.indexOfScalar(u8, name, '\'') orelse unreachable) + 1;
        const end = std.mem.indexOfScalarPos(u8, name, start, '\'') orelse unreachable;
        const len = end - start;
        var buffer: [len + 1]u8 = undefined;
        @memcpy(buffer[0..len], name[start..end]);
        buffer[len] = 0;
        const out = buffer;
        var slice: [:0]const u8 = undefined;
        slice.ptr = @ptrCast(&out);
        slice.len = len;
        break :blk slice;
    };
}

test {
    const T = struct {
        fn myFunction() void {}
    };
    try std.testing.expectEqualStrings("myFunction", fnName(T.myFunction));
    try std.testing.expectEqualStrings("print", fnName(std.debug.print));
}

// todo: @typeName isn't displaying the expected result
// test {
//     const namespace = struct {
//         const inner = struct {
//             const FooType = struct { u64, u32 };
//             const FooConfig = struct {};
//             fn Foo(comptime T: type, comptime config: anytype) type {
//                 _ = T;
//                 _ = config;
//                 return struct {};
//             }
//         };
//     };
//
//     const in = namespace.inner.FooConfig{};
//     const b: usize = 2;
//     const T1 = namespace.inner.Foo(u32, std.debug.print);
//     std.debug.print("{s}\n", .{@typeName(T1)});
//     std.debug.print("{s}\n\n", .{typeName(T1)});
//
//     const T2 = namespace.inner.Foo(namespace.inner.FooType, .{ in, b });
//     std.debug.print("{s}\n", .{@typeName(T2)});
//     std.debug.print("{s}\n\n", .{typeName(T2)});
// }

test {
    const Vec2 = @Vector(2, f32);
    const Pos = struct { value: Vec2 = @splat(0.0) };
    const Vel = struct { value: Vec2 = @splat(0.0) };
    const Particle = struct { pos: Pos = .{}, vel: Vel = .{} };

    const reflect = comptime Reflection.init(&.{Particle});
    // try reflect.assert();

    const type_id = reflect.findId("Particle") orelse unreachable;
    const pos_x = reflect.findProp(type_id, "pos.value[0]") orelse unreachable;
    const pos_y = reflect.findProp(type_id, "pos.value[1]") orelse unreachable;

    var particle = Particle{ .pos = .{ .value = .{ 1.0, 2.0 } } };
    try std.testing.expectEqual(1.0, (pos_x.get(f32, @ptrCast(&particle)) orelse unreachable).*);
    try std.testing.expectEqual(2.0, (pos_y.get(f32, @ptrCast(&particle)) orelse unreachable).*);
}

test {
    const Name = struct { value: [*:0]const u8 };
    const A = struct { value: i32 };
    const B = struct { value: *const A };
    const Misc = struct { name: Name, a: *const A, b: *const B };

    const reflect = comptime Reflection.init(&.{Misc});
    // try reflect.assert();

    const type_id = reflect.findId("Misc") orelse unreachable;
    const namev = reflect.findProp(type_id, "name.value") orelse unreachable;
    const a = reflect.findProp(type_id, "a") orelse unreachable;
    const av = reflect.findProp(type_id, "a.value") orelse unreachable;
    const b = reflect.findProp(type_id, "b") orelse unreachable;
    const bv = reflect.findProp(type_id, "b.value") orelse unreachable;
    const bvv = reflect.findProp(type_id, "b.value.value") orelse unreachable;

    var misc = Misc{
        .name = .{ .value = "john4" },
        .a = &A{ .value = -1 },
        .b = &B{ .value = &A{ .value = 1 } },
    };
    try std.testing.expectEqualStrings("john4", std.mem.span((namev.get([*:0]u8, @ptrCast(&misc)) orelse unreachable).*));
    try std.testing.expectEqual(misc.a, (a.get(*const A, @ptrCast(&misc)) orelse unreachable).*);
    try std.testing.expectEqual(misc.a.value, (av.get(i32, @ptrCast(&misc)) orelse unreachable).*);
    try std.testing.expectEqual(misc.b, (b.get(*const B, @ptrCast(&misc)) orelse unreachable).*);
    try std.testing.expectEqual(misc.b.value, (bv.get(*const A, @ptrCast(&misc)) orelse unreachable).*);
    try std.testing.expectEqual(misc.b.value.value, (bvv.get(i32, @ptrCast(&misc)) orelse unreachable).*);
}

// test {
//     const reflect = comptime Reflection.init(&.{ []u8, []@Vector(2, f32) });
//     try reflect.assert();
// }

test {
    // optional pointers are just pointers that might be null
    try std.testing.expectEqual(@sizeOf(*u8), @sizeOf(?*u8));
    try std.testing.expectEqual(@sizeOf([]u8), @sizeOf(?[]u8));

    // // how optionals are expected to work
    // const reflect = comptime Reflection.init(&.{ ?u8, ??u8, ?u16, ??u16, ?u32, ??u32, ?u64, ??u64, ?[4]u8, ??[4]u8, ?struct { u8, u32 }, ??struct { u8, u32 } });
    // try reflect.assert();
}