Add support for user atoms and namecalls

Namecall is a mechanism in Luau to speed up method invocations.
The basic idea is that the VM can cache method names (strings)
to integer indices the first time it executes a method call.
At this point it calls the "user atom callback" with the string.
The user callback is responsible for mapping the method string
to a unique 16-bit index that's returned to the VM.

Next time the VM encounters the same string, it already knows
how to map the string to an index as, so it will reuse the user's
16-bit index.

The above is the mechanism for quickly resolving function name
strings to integers.  The other part of the API is using the
indices.  This part is the __namecall function that's attached
to a (userdata) object's metatable.  On a method call, the VM
knows that the userdata has a registered __namecall, and calls
that to dispatch to the actual user's native function to handle
the native method.  The namecall dispatch routine uses lua.namecallAtom()
to retrieve the method name/index, which is used to select which
actual native method is called.

It's not very simple but it should be fast as the VM doesn't
need to do a string->function hash table lookup on every
method invocation.

I'm not 100% sure of the details, but I suspect that the VM may
also patch the bytecode (or some internal representation of it)
directly with the namecall indices rather than looking them up
from some string hash table.

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);
+}