Expose Luau built-in float vector support

The Luau VM supports native f32 3- or 4-vectors so that typical
linear algebra operations are fast in game code.

Both the 3- and 4-vector flavors are supported.  Use
-Dluau_vector_size=N to choose which.  This must be configured
at build time, as the native Luau VM must be re-compiled for this
setting.

Author: nurpax

build.zig

@@ -20,7 +20,7 @@ pub fn build(b: *Build) void {
 
     const lang = b.option(Language, "lang", "Lua language version to build") orelse .lua54;
     const shared = b.option(bool, "shared", "Build shared library instead of static") orelse false;
-
+    const luau_vector_size = b.option(usize, "luau_vector_size", "Luau built-in vector size: 3-vector or 4-vector") orelse 3;
     const upstream = b.dependency(@tagName(lang), .{});
 
     // Zig module
@@ -37,10 +37,11 @@ pub fn build(b: *Build) void {
     // Expose build configuration to the ziglua module
     const config = b.addOptions();
     config.addOption(Language, "lang", lang);
+    config.addOption(usize, "luau_vector_size", luau_vector_size);
     ziglua.addOptions("config", config);
 
     const lib = switch (lang) {
-        .luau => buildLuau(b, target, optimize, upstream, shared),
+        .luau => buildLuau(b, target, optimize, upstream, shared, luau_vector_size),
         else => buildLua(b, target, optimize, upstream, lang, shared),
     };
 
@@ -152,7 +153,7 @@ fn buildLua(b: *Build, target: Build.ResolvedTarget, optimize: std.builtin.Optim
 }
 
 /// Luau has diverged enough from Lua (C++, project structure, ...) that it is easier to separate the build logic
-fn buildLuau(b: *Build, target: Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, upstream: *Build.Dependency, shared: bool) *Step.Compile {
+fn buildLuau(b: *Build, target: Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, upstream: *Build.Dependency, shared: bool, luau_vector_len: usize) *Step.Compile {
     const lib_opts = .{
         .name = "luau",
         .target = target,
@@ -174,6 +175,7 @@ fn buildLuau(b: *Build, target: Build.ResolvedTarget, optimize: std.builtin.Opti
         "-DLUA_API=extern\"C\"",
         "-DLUACODE_API=extern\"C\"",
         "-DLUACODEGEN_API=extern\"C\"",
+        if (luau_vector_len == 4) "-DLUA_VECTOR_SIZE=4" else "-DLUA_VECTOR_SIZE=3",
     };
 
     for (luau_source_files) |file| {

src/libluau.zig

@@ -11,13 +11,19 @@ const c = @cImport({
 
 const config = @import("config");
 pub const lang = config.lang;
+/// True if Luau built for 4-vectors
+pub const luau_vector_size = config.luau_vector_size;
 
 /// This function is defined in luau.cpp and must be called to define the assertion printer
 extern "c" fn zig_registerAssertionHandler() void;
 
 /// This function is defined in luau.cpp and ensures Zig uses the correct free when compiling luau code
 extern "c" fn zig_luau_free(ptr: *anyopaque) void;
 
+// Workaround for Zig apparently always choosing the 3-vector variant of lua_pushvector(lua_State* L, float x, float y, float z)
+// despite compiling Luau with -DLUA_VECTOR_SIZE=4.
+extern "c" fn zig_lua_pushvector4(L: ?*LuaState, x: f32, y: f32, z: f32, w: f32) void;
+
 const Allocator = std.mem.Allocator;
 
 // Types
@@ -139,6 +145,7 @@ pub const LuaType = enum(i5) {
     boolean = c.LUA_TBOOLEAN,
     light_userdata = c.LUA_TLIGHTUSERDATA,
     number = c.LUA_TNUMBER,
+    vector = c.LUA_TVECTOR,
     string = c.LUA_TSTRING,
     table = c.LUA_TTABLE,
     function = c.LUA_TFUNCTION,
@@ -520,6 +527,11 @@ pub const Lua = struct {
         return c.lua_isuserdata(lua.state, index) != 0;
     }
 
+    /// Returns true if the value at the given index is a vector
+    pub fn isVector(lua: *Lua, index: i32) bool {
+        return c.lua_isvector(lua.state, index);
+    }
+
     /// Returns true if the value at index1 is smaller than the value at index2, following the
     /// semantics of the Lua < operator.
     /// See https://www.lua.org/manual/5.4/manual.html#lua_lessthan
@@ -713,6 +725,17 @@ pub const Lua = struct {
         c.lua_pushvalue(lua.state, index);
     }
 
+    fn pushVector3(lua: *Lua, x: f32, y: f32, z: f32) void {
+        c.lua_pushvector(lua.state, x, y, z);
+    }
+
+    fn pushVector4(lua: *Lua, x: f32, y: f32, z: f32, w: f32) void {
+        zig_lua_pushvector4(lua.state, x, y, z, w);
+    }
+
+    /// Pushes a floating point 3-vector (or 4-vector if configured) `v` onto the stack
+    pub const pushVector = if (luau_vector_size == 3) pushVector3 else pushVector4;
+
     /// Returns true if the two values in indices `index1` and `index2` are primitively equal
     /// Bypasses __eq metamethods
     /// Returns false if not equal, or if any index is invalid
@@ -927,6 +950,22 @@ pub const Lua = struct {
         return error.Fail;
     }
 
+    /// Converts the Lua value at the given `index` to a 3- or 4-vector.
+    /// The Lua value must be a vector.
+    pub fn toVector(lua: *Lua, index: i32) ![luau_vector_size]f32 {
+        const res = c.lua_tovector(lua.state, index);
+        if (res) |r| {
+            if (luau_vector_size == 3) {
+                return [_]f32{ r[0], r[1], r[2] };
+            } else if (luau_vector_size == 4) {
+                return [_]f32{ r[0], r[1], r[2], r[3] };
+            } else {
+                @compileError("luau_vector_size must be 3 or 4");
+            }
+        }
+        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

src/luau.cpp

@@ -1,5 +1,5 @@
 #include "Luau/Common.h"
-
+#include "lua.h"
 #include <cstdio>
 #include <cstdlib>
 
@@ -16,3 +16,9 @@ extern "C" void zig_registerAssertionHandler() {
 extern "C" void zig_luau_free(void *ptr) {
     free(ptr);
 }
+
+#if LUA_VECTOR_SIZE == 4
+extern "C" void zig_lua_pushvector4(lua_State* L, float x, float y, float z, float w) {
+    lua_pushvector(L, x, y, z, w);
+}
+#endif

src/tests.zig

@@ -381,6 +381,7 @@ test "typenames" {
     try expectEqualStrings("function", lua.typeName(.function));
     try expectEqualStrings("userdata", lua.typeName(.userdata));
     try expectEqualStrings("thread", lua.typeName(.thread));
+    try expectEqualStrings("vector", lua.typeName(.vector));
 }
 
 test "unsigned" {
@@ -2249,3 +2250,64 @@ test "tagged userdata" {
     lua.pushInteger(13);
     try expectError(error.Fail, lua.userdataTag(-1));
 }
+
+fn vectorCtor(l: *Lua) i32 {
+    const x = l.toNumber(1) catch unreachable;
+    const y = l.toNumber(2) catch unreachable;
+    const z = l.toNumber(3) catch unreachable;
+    if (ziglua.luau_vector_size == 4) {
+        const w = l.optNumber(4, 0);
+        l.pushVector(@floatCast(x), @floatCast(y), @floatCast(z), @floatCast(w));
+    } else {
+        l.pushVector(@floatCast(x), @floatCast(y), @floatCast(z));
+    }
+    return 1;
+}
+
+test "luau vectors" {
+    var lua = try Lua.init(testing.allocator);
+    defer lua.deinit();
+    lua.openLibs();
+    lua.register("vector", ziglua.wrap(vectorCtor));
+
+    try lua.doString(
+        \\function test()
+        \\  local a = vector(1, 2, 3)
+        \\  local b = vector(4, 5, 6)
+        \\  local c = (a + b) * vector(2, 2, 2)
+        \\  return c
+        \\end
+    );
+    _ = try lua.getGlobal("test");
+    try lua.protectedCall(0, 1, 0);
+    var v = try lua.toVector(-1);
+    try testing.expectEqualSlices(f32, &[3]f32{ 10, 14, 18 }, v[0..3]);
+
+    if (ziglua.luau_vector_size == 3) lua.pushVector(1, 2, 3) else lua.pushVector(1, 2, 3, 4);
+    try expect(lua.isVector(-1));
+    v = try lua.toVector(-1);
+    const expected = if (ziglua.luau_vector_size == 3) [3]f32{ 1, 2, 3 } else [4]f32{ 1, 2, 3, 4 };
+    try expectEqual(expected, v);
+    try expectEqualStrings("vector", lua.typeNameIndex(-1));
+
+    lua.pushInteger(5);
+    try expect(!lua.isVector(-1));
+}
+
+test "luau 4-vectors" {
+    var lua = try Lua.init(testing.allocator);
+    defer lua.deinit();
+    lua.openLibs();
+    lua.register("vector", ziglua.wrap(vectorCtor));
+
+    // More specific 4-vector tests
+    if (ziglua.luau_vector_size == 4) {
+        try lua.doString(
+            \\local a = vector(1, 2, 3, 4)
+            \\local b = vector(5, 6, 7, 8)
+            \\return a + b
+        );
+        const vec4 = try lua.toVector(-1);
+        try expectEqual([4]f32{ 6, 8, 10, 12 }, vec4);
+    }
+}