Add support for user atoms and namecalls

src/libluau.zig

@@ -42,6 +42,9 @@ pub const CFn = *const fn (state: ?*LuaState) callconv(.C) c_int;
 /// Type for C userdata destructors
 pub const CUserdataDtorFn = *const fn (userdata: *anyopaque) callconv(.C) void;
 
+/// Type for C useratom callback
+pub const CUserAtomCallbackFn = *const fn (str: [*c]const u8, len: usize) callconv(.C) i16;
+
 /// The internal Lua debug structure
 /// See https://www.lua.org/manual/5.1/manual.html#lua_Debug
 const Debug = c.lua_Debug;
@@ -326,6 +329,14 @@ pub const Lua = struct {
         c.lua_createtable(lua.state, num_arr, num_rec);
     }
 
+    pub fn setReadonly(lua: *Lua, idx: i32, enabled: bool) void {
+        c.lua_setreadonly(lua.state, idx, @intFromBool(enabled));
+    }
+
+    pub fn getReadonly(lua: *Lua, idx: i32) bool {
+        return c.lua_getreadonly(lua.state, idx) != 0;
+    }
+
     /// Returns true if the two values at the indexes are equal following the semantics of the
     /// Lua == operator.
     /// See https://www.lua.org/manual/5.1/manual.html#lua_equal
@@ -961,6 +972,26 @@ pub const Lua = struct {
         return error.Fail;
     }
 
+    /// Converts the Lua string at the given `index` to a string atom.
+    /// The Lua value must be a string.
+    pub fn toStringAtom(lua: *Lua, index: i32) !struct { i32, [:0]const u8 } {
+        var atom: c_int = undefined;
+        if (c.lua_tostringatom(lua.state, index, &atom)) |ptr| {
+            return .{ atom, std.mem.span(ptr) };
+        }
+        return error.Fail;
+    }
+
+    /// Retrieve the user atom index and name for the method being
+    /// invoked in a namecall.
+    pub fn namecallAtom(lua: *Lua) !struct { i32, [:0]const u8 } {
+        var atom: c_int = undefined;
+        if (c.lua_namecallatom(lua.state, &atom)) |ptr| {
+            return .{ atom, std.mem.span(ptr) };
+        }
+        return error.Fail;
+    }
+
     /// Returns the `LuaType` of the value at the given index
     /// Note that this is equivalent to lua_type but because type is a Zig primitive it is renamed to `typeOf`
     /// See https://www.lua.org/manual/5.1/manual.html#lua_type
@@ -1069,6 +1100,12 @@ pub const Lua = struct {
         return error.Fail;
     }
 
+    pub fn setUserAtomCallbackFn(lua: *Lua, cb: CUserAtomCallbackFn) void {
+        if (c.lua_callbacks(lua.state)) |cb_struct| {
+            cb_struct.*.useratom = cb;
+        }
+    }
+
     // Auxiliary library functions
     //
     // Auxiliary library functions are included in alphabetical order.
@@ -1089,6 +1126,13 @@ pub const Lua = struct {
         unreachable;
     }
 
+    /// Raises a type error for the argument arg of the C function that called it, using a standard message; tname is a "name" for the expected type. This function never returns.
+    /// See https://www.lua.org/manual/5.4/manual.html#luaL_typeerror
+    pub fn typeError(lua: *Lua, arg: i32, type_name: [:0]const u8) noreturn {
+        _ = c.luaL_typeerror(lua.state, arg, type_name.ptr);
+        unreachable;
+    }
+
     /// Calls a metamethod
     /// See https://www.lua.org/manual/5.1/manual.html#luaL_callmeta
     pub fn callMeta(lua: *Lua, obj: i32, field: [:0]const u8) !void {
@@ -1182,6 +1226,14 @@ pub const Lua = struct {
         return @as([*]T, @ptrCast(@alignCast(ptr)))[0..size];
     }
 
+    /// Checks whether the function argument `arg` is a vector and returns the vector as a floating point slice.
+    pub fn checkVector(lua: *Lua, arg: i32) [luau_vector_size]f32 {
+        const vec = lua.toVector(arg) catch {
+            lua.typeError(arg, lua.typeName(LuaType.vector));
+        };
+        return vec;
+    }
+
     /// Loads and runs the given string
     /// See https://www.lua.org/manual/5.1/manual.html#luaL_dostring
     /// TODO: does it make sense to have this in Luau?
@@ -1498,6 +1550,7 @@ pub const ZigContFn = fn (lua: *Lua, status: Status, ctx: i32) i32;
 pub const ZigReaderFn = fn (lua: *Lua, data: *anyopaque) ?[]const u8;
 pub const ZigWriterFn = fn (lua: *Lua, buf: []const u8, data: *anyopaque) bool;
 pub const ZigUserdataDtorFn = fn (data: *anyopaque) void;
+pub const ZigUserAtomCallbackFn = fn (str: []const u8) i16;
 
 fn TypeOfWrap(comptime T: type) type {
     return switch (T) {
@@ -1506,6 +1559,7 @@ fn TypeOfWrap(comptime T: type) type {
         ZigReaderFn => CReaderFn,
         ZigWriterFn => CWriterFn,
         ZigUserdataDtorFn => CUserdataDtorFn,
+        ZigUserAtomCallbackFn => CUserAtomCallbackFn,
         else => @compileError("unsupported type given to wrap: '" ++ @typeName(T) ++ "'"),
     };
 }
@@ -1521,6 +1575,7 @@ pub fn wrap(comptime value: anytype) TypeOfWrap(@TypeOf(value)) {
         ZigReaderFn => wrapZigReaderFn(value),
         ZigWriterFn => wrapZigWriterFn(value),
         ZigUserdataDtorFn => wrapZigUserdataDtorFn(value),
+        ZigUserAtomCallbackFn => wrapZigUserAtomCallbackFn(value),
         else => @compileError("unsupported type given to wrap: '" ++ @typeName(T) ++ "'"),
     };
 }
@@ -1545,6 +1600,19 @@ fn wrapZigUserdataDtorFn(comptime f: ZigUserdataDtorFn) CUserdataDtorFn {
     }.inner;
 }
 
+/// Wrap a ZigFn in a CFn for passing to the API
+fn wrapZigUserAtomCallbackFn(comptime f: ZigUserAtomCallbackFn) CUserAtomCallbackFn {
+    return struct {
+        fn inner(str: [*c]const u8, len: usize) callconv(.C) i16 {
+            if (str) |s| {
+                const buf = s[0..len];
+                return @call(.always_inline, f, .{buf});
+            }
+            return -1;
+        }
+    }.inner;
+}
+
 /// Wrap a ZigReaderFn in a CReaderFn for passing to the API
 fn wrapZigReaderFn(comptime f: ZigReaderFn) CReaderFn {
     return struct {

src/tests.zig

@@ -2318,3 +2318,115 @@ test "luau 4-vectors" {
         try expectEqual([4]f32{ 6, 8, 10, 12 }, vec4);
     }
 }
+
+test "useratom" {
+    if (ziglua.lang != .luau) return;
+
+    const useratomCb = struct {
+        pub fn inner(str: []const u8) i16 {
+            if (std.mem.eql(u8, str, "method_one")) {
+                return 0;
+            } else if (std.mem.eql(u8, str, "another_method")) {
+                return 1;
+            }
+            return -1;
+        }
+    }.inner;
+
+    var lua = try Lua.init(testing.allocator);
+    defer lua.deinit();
+    lua.setUserAtomCallbackFn(ziglua.wrap(useratomCb));
+
+    _ = lua.pushString("unknownatom");
+    _ = lua.pushString("method_one");
+    _ = lua.pushString("another_method");
+
+    const atom_idx0, const str0 = try lua.toStringAtom(-2);
+    const atom_idx1, const str1 = try lua.toStringAtom(-1);
+    const atom_idx2, const str2 = try lua.toStringAtom(-3);
+    try testing.expect(std.mem.eql(u8, str0, "method_one"));
+    try testing.expect(std.mem.eql(u8, str1, "another_method"));
+    try testing.expect(std.mem.eql(u8, str2, "unknownatom")); // should work, but returns -1 for atom idx
+
+    try expectEqual(0, atom_idx0);
+    try expectEqual(1, atom_idx1);
+    try expectEqual(-1, atom_idx2);
+
+    lua.pushInteger(13);
+    try expectError(error.Fail, lua.toStringAtom(-1));
+}
+
+test "namecall" {
+    if (ziglua.lang != .luau) return;
+
+    const funcs = struct {
+        const dot_idx: i32 = 0;
+        const sum_idx: i32 = 1;
+
+        // The useratom callback to initially form a mapping from method names to
+        // integer indices. The indices can then be used to quickly dispatch the right
+        // method in namecalls without needing to perform string compares.
+        pub fn useratomCb(str: []const u8) i16 {
+            if (std.mem.eql(u8, str, "dot")) {
+                return dot_idx;
+            }
+            if (std.mem.eql(u8, str, "sum")) {
+                return sum_idx;
+            }
+            return -1;
+        }
+
+        pub fn vectorNamecall(l: *Lua) i32 {
+            const atom_idx, _ = l.namecallAtom() catch {
+                l.raiseErrorStr("%s is not a valid vector method", .{l.checkString(1)});
+            };
+            switch (atom_idx) {
+                dot_idx => {
+                    const a = l.checkVector(1);
+                    const b = l.checkVector(2);
+                    l.pushNumber(a[0] * b[0] + a[1] * b[1] + a[2] * b[2]); // vec3 dot
+                    return 1;
+                },
+                sum_idx => {
+                    const a = l.checkVector(1);
+                    l.pushNumber(a[0] + a[1] + a[2]);
+                    return 1;
+                },
+                else => unreachable,
+            }
+        }
+    };
+
+    var lua = try Lua.init(testing.allocator);
+    defer lua.deinit();
+    lua.setUserAtomCallbackFn(ziglua.wrap(funcs.useratomCb));
+
+    lua.register("vector", ziglua.wrap(vectorCtor));
+    lua.pushVector(0, 0, 0);
+
+    try lua.newMetatable("vector");
+    lua.pushString("__namecall");
+    lua.pushFunction(ziglua.wrap(funcs.vectorNamecall), "vector_namecall");
+    lua.setTable(-3);
+
+    lua.setReadonly(-1, true);
+    lua.setMetatable(-2);
+
+    // Vector setup, try some lua code on them.
+    try lua.doString(
+        \\local a = vector(1, 2, 3)
+        \\local b = vector(3, 2, 1)
+        \\return a:dot(b)
+    );
+    const d = try lua.toNumber(-1);
+    lua.pop(-1);
+    try expectEqual(10, d);
+
+    try lua.doString(
+        \\local a = vector(1, 2, 3)
+        \\return a:sum()
+    );
+    const s = try lua.toNumber(-1);
+    lua.pop(-1);
+    try expectEqual(6, s);
+}