[Lua] Add LuaJIT support (#6584)

* [Lua] Add LuaJIT support

Here is the output of LuaTest.sh:

Run with LuaJIT:
built 100 512-byte flatbuffers in 0.16sec: 0.63/msec, 0.31MB/sec
built 1000 512-byte flatbuffers in 0.08sec: 12.06/msec, 5.89MB/sec
built 10000 512-byte flatbuffers in 0.80sec: 12.44/msec, 6.07MB/sec
built 10000 512-byte flatbuffers in 0.33sec: 30.58/msec, 14.93MB/sec
traversed 100 592-byte flatbuffers in 0.04sec: 2.51/msec, 1.42MB/sec
traversed 1000 592-byte flatbuffers in 0.03sec: 31.52/msec, 17.79MB/sec
traversed 10000 592-byte flatbuffers in 0.21sec: 48.77/msec, 27.53MB/sec
Run with Lua 5.3:
built 100 512-byte flatbuffers in 0.02sec: 5.44/msec, 2.66MB/sec
built 1000 512-byte flatbuffers in 0.17sec: 5.74/msec, 2.80MB/sec
built 10000 512-byte flatbuffers in 1.75sec: 5.72/msec, 2.79MB/sec
built 10000 512-byte flatbuffers in 1.38sec: 7.26/msec, 3.55MB/sec
traversed 100 592-byte flatbuffers in 0.00sec: 27.64/msec, 15.60MB/sec
traversed 1000 592-byte flatbuffers in 0.03sec: 30.46/msec, 17.20MB/sec
traversed 10000 592-byte flatbuffers in 0.34sec: 29.62/msec, 16.72MB/sec

* [Lua] Better usage description

* update according to the review

9 files changed, 253 insertions, 26 deletions

diff --git a/docs/source/LuaUsage.md b/docs/source/LuaUsage.md
index 75b1f3b8..43c370f5 100644
--- a/docs/source/LuaUsage.md
+++ b/docs/source/LuaUsage.md
@@ -29,8 +29,8 @@ flatbuffers/blob/master/tests/luatest.lua).
 To run the tests, use the [LuaTest.sh](https://github.com/google/flatbuffers/
 blob/master/tests/LuaTest.sh) shell script.
 
-*Note: This script requires [Lua 5.3](https://www.lua.org/) to be
-installed.*
+*Note: This script requires [Lua 5.3](https://www.lua.org/) and
+[LuaJIT](http://luajit.org/) to be installed.*
 
 ## Using the FlatBuffers Lua library
 
diff --git a/lua/flatbuffers/binaryarray.lua b/lua/flatbuffers/binaryarray.lua
index bf728cdd..3d16c36e 100644
--- a/lua/flatbuffers/binaryarray.lua
+++ b/lua/flatbuffers/binaryarray.lua
@@ -1,3 +1,9 @@
+local compat = require("flatbuffers.compat")
+-- locals for slightly faster access
+local string_pack = compat.string_pack
+local string_unpack = compat.string_unpack
+
+
 local m = {} -- the module table
 
 local mt = {} -- the module metatable
@@ -104,19 +110,15 @@ function mt:Set(value, position)
     self.data[position] = value
 end
 
--- locals for slightly faster access
-local sunpack = string.unpack
-local spack = string.pack
-
 -- Pack the data into a binary representation
 function m.Pack(fmt, ...)
-    return spack(fmt, ...)
+    return string_pack(fmt, ...)
 end
 
 -- Unpack the data from a binary representation in
 -- a Lua value
 function m.Unpack(fmt, s, pos)
-    return sunpack(fmt, s.str, pos + 1)
+    return string_unpack(fmt, s.str, pos + 1)
 end
 
 -- Return the binary array module
diff --git a/lua/flatbuffers/builder.lua b/lua/flatbuffers/builder.lua
index 4a0c4f6a..bb807232 100644
--- a/lua/flatbuffers/builder.lua
+++ b/lua/flatbuffers/builder.lua
@@ -1,6 +1,7 @@
 local N = require("flatbuffers.numTypes")
 local ba = require("flatbuffers.binaryarray")
 local compat = require("flatbuffers.compat")
+local string_unpack = compat.string_unpack
 
 local m = {}
 
@@ -34,7 +35,7 @@ local function vtableEqual(a, objectStart, b)
     end
 
     for i, elem in ipairs(a) do
-        local x = string.unpack(VOffsetT.packFmt, b, 1 + (i - 1) * 2)
+        local x = string_unpack(VOffsetT.packFmt, b, 1 + (i - 1) * 2)
         if x ~= 0 or elem ~= 0 then
             local y = objectStart - elem
             if x ~= y then
@@ -119,7 +120,7 @@ function mt:WriteVtable()
         local vt2Offset = self.vtables[i]
         local vt2Start = #self.bytes - vt2Offset
         local vt2lenstr = self.bytes:Slice(vt2Start, vt2Start+1)
-        local vt2Len = string.unpack(VOffsetT.packFmt, vt2lenstr, 1)
+        local vt2Len = string_unpack(VOffsetT.packFmt, vt2lenstr, 1)
 
         local metadata = VtableMetadataFields * 2
         local vt2End = vt2Start + vt2Len
@@ -210,7 +211,7 @@ function mt:Prep(size, additionalBytes)
     local h = self.head
 
     local k = #self.bytes - h + additionalBytes
-    local alignsize = ((~k) + 1) & (size - 1) -- getAlignSize(k, size)
+    local alignsize = getAlignSize(k, size)
 
     local desiredSize = alignsize + size + additionalBytes
 
diff --git a/lua/flatbuffers/compat.lua b/lua/flatbuffers/compat.lua
index 89c992be..fbce4578 100644
--- a/lua/flatbuffers/compat.lua
+++ b/lua/flatbuffers/compat.lua
@@ -1,17 +1,12 @@
-local m = {}
-
-local getAlignSize
+local m
 if _VERSION == "Lua 5.3" then
-    getAlignSize = function(k, size)
-            return ((~k) + 1) & (size - 1)
-        end    
+    m = require("flatbuffers.compat_5_3")
 else
-    getAlignSize = function(self, size, additionalBytes)        
-        local alignsize = bit32.bnot(#self.bytes-self:Head() + additionalBytes) + 1
-        return bit32.band(alignsize,(size - 1))
+    local ok = pcall(require, "jit")
+    if not ok then
+        error("Only Lua 5.3 or LuaJIT is supported")
+    else
+        m = require("flatbuffers.compat_luajit")
     end
 end
-    
-m.GetAlignSize = getAlignSize
-
 return m
\ No newline at end of file
diff --git a/lua/flatbuffers/compat_5_3.lua b/lua/flatbuffers/compat_5_3.lua
new file mode 100644
index 00000000..762bf073
--- /dev/null
+++ b/lua/flatbuffers/compat_5_3.lua
@@ -0,0 +1,14 @@
+-- We need to put it into a separate file to avoid syntax error like `unexpected symbol near '~'`
+local m = {}
+
+
+m.GetAlignSize = function(k, size)
+    return ((~k) + 1) & (size - 1)
+end
+
+
+m.string_pack = string.pack
+m.string_unpack = string.unpack
+
+
+return m
diff --git a/lua/flatbuffers/compat_luajit.lua b/lua/flatbuffers/compat_luajit.lua
new file mode 100644
index 00000000..cb8be71e
--- /dev/null
+++ b/lua/flatbuffers/compat_luajit.lua
@@ -0,0 +1,206 @@
+local bit = require("bit")
+local ffi = require("ffi")
+local band = bit.band
+local bnot = bit.bnot
+
+
+local m = {}
+local Uint8Bound = 256 -- bound is the max uintN + 1
+local Uint16Bound = 65536
+local Uint32Bound = 4294967296
+
+
+m.GetAlignSize = function(k, size)
+    return band((bnot(k) + 1), (size - 1))
+end
+
+
+local function pack_I1(n)
+    return string.char(n)
+end
+local function pack_i1(n)
+    if n < 0 then
+        n = Uint8Bound + n
+    end
+    return pack_I1(n)
+end
+
+local function unpack_I1(n, pos)
+    return string.byte(n, pos)
+end
+local function unpack_i1(n, pos)
+    local res = unpack_I1(n, pos)
+    if res >= Uint8Bound / 2 then
+        return res - Uint8Bound
+    end
+    return res
+end
+
+local b2 = ffi.new("unsigned char[2]")
+local function pack_I2(n)
+    for i = 0, 1 do
+        b2[i] = bit.band(n, 255)
+        n = bit.rshift(n, 8)
+    end
+    return ffi.string(b2, 2)
+end
+local function pack_i2(n)
+    if n < 0 then
+        n = Uint16Bound + n
+    end
+    return pack_I2(n)
+end
+
+local function unpack_I2(n, pos)
+    local a, b = string.byte(n, pos, pos + 1)
+    return b * Uint8Bound + a
+end
+local function unpack_i2(n, pos)
+    local res = unpack_I2(n, pos)
+    if res >= Uint16Bound / 2 then
+        return res - Uint16Bound
+    end
+    return res
+end
+
+local b4 = ffi.new("unsigned char[4]")
+local function pack_I4(n)
+    for i = 0, 3 do
+        b4[i] = bit.band(n, 255)
+        n = bit.rshift(n, 8)
+    end
+    return ffi.string(b4, 4)
+end
+local function pack_i4(n)
+    if n < 0 then
+        n = Uint32Bound + n
+    end
+    return pack_I4(n)
+end
+
+local function unpack_I4(n, pos)
+    local a, b, c, d = string.byte(n, pos, pos + 3)
+    return Uint8Bound * (Uint8Bound * ((Uint8Bound * d) + c) + b) + a
+end
+local function unpack_i4(n, pos)
+    local res = unpack_I4(n, pos)
+    if res >= Uint32Bound / 2 then
+        return res - Uint32Bound
+    end
+    return res
+end
+
+local b8 = ffi.new("unsigned char[8]")
+local function pack_I8(n)
+    n = ffi.cast("unsigned long long", n)
+    local hi = math.floor(tonumber(n / Uint32Bound))
+    local li = n % Uint32Bound
+    for i = 0, 3 do
+        b8[i] = bit.band(li, 255)
+        li = bit.rshift(li, 8)
+    end
+    for i = 4, 7 do
+        b8[i] = bit.band(hi, 255)
+        hi = bit.rshift(hi, 8)
+    end
+    return ffi.string(b8, 8)
+end
+local function pack_i8(n)
+    n = ffi.cast("signed long long", n)
+    return pack_I8(n)
+end
+
+local function unpack_I8(n, pos)
+    local a, b, c, d = string.byte(n, pos, pos + 3)
+    local li = Uint8Bound * (Uint8Bound * ((Uint8Bound * d) + c) + b) + a
+    local a, b, c, d = string.byte(n, pos + 4, pos + 7)
+    local hi = Uint8Bound * (Uint8Bound * ((Uint8Bound * d) + c) + b) + a
+    return ffi.cast("unsigned long long", hi) * Uint32Bound + li
+end
+local function unpack_i8(n, pos)
+    local res = unpack_I8(n, pos)
+    return ffi.cast("signed long long", res)
+end
+
+local bf = ffi.new("float[1]")
+local function pack_f(n)
+    bf[0] = n
+    return ffi.string(bf, 4)
+end
+
+local function unpack_f(n, pos)
+    ffi.copy(bf, ffi.cast("char *", n) + pos - 1, 4)
+    return tonumber(bf[0])
+end
+
+local bd = ffi.new("double[1]")
+local function pack_d(n)
+    bd[0] = n
+    return ffi.string(bd, 8)
+end
+
+local function unpack_d(n, pos)
+    ffi.copy(bd, ffi.cast("char *", n) + pos - 1, 8)
+    return tonumber(bd[0])
+end
+
+
+m.string_pack = function(fmt, i, ...)
+    if fmt == "<I1" then
+        return pack_I1(i)
+    elseif fmt == "<I2" then
+        return pack_I2(i)
+    elseif fmt == "<I4" then
+        return pack_I4(i)
+    elseif fmt == "<I8" then
+        return pack_I8(i)
+    elseif fmt == "<i1" then
+        return pack_i1(i)
+    elseif fmt == "<i2" then
+        return pack_i2(i)
+    elseif fmt == "<i4" then
+        return pack_i4(i)
+    elseif fmt == "<i8" then
+        return pack_i8(i)
+    elseif fmt == "<f" then
+        return pack_f(i)
+    elseif fmt == "<d" then
+        return pack_d(i)
+    else
+        error(string.format("FIXME: support fmt %s", fmt))
+    end
+end
+
+
+m.string_unpack = function(fmt, s, pos)
+    if not pos then
+        pos = 1
+    end
+
+    if fmt == "<I1" then
+        return unpack_I1(s, pos)
+    elseif fmt == "<I2" then
+        return unpack_I2(s, pos)
+    elseif fmt == "<I4" then
+        return unpack_I4(s, pos)
+    elseif fmt == "<I8" then
+        return unpack_I8(s, pos)
+    elseif fmt == "<i1" then
+        return unpack_i1(s, pos)
+    elseif fmt == "<i2" then
+        return unpack_i2(s, pos)
+    elseif fmt == "<i4" then
+        return unpack_i4(s, pos)
+    elseif fmt == "<i8" then
+        return unpack_i8(s, pos)
+    elseif fmt == "<f" then
+        return unpack_f(s, pos)
+    elseif fmt == "<d" then
+        return unpack_d(s, pos)
+    else
+        error(string.format("FIXME: support fmt %s", fmt))
+    end
+end
+
+
+return m
diff --git a/lua/flatbuffers/view.lua b/lua/flatbuffers/view.lua
index fde15c3d..105e967d 100644
--- a/lua/flatbuffers/view.lua
+++ b/lua/flatbuffers/view.lua
@@ -1,3 +1,7 @@
+local compat = require("flatbuffers.compat")
+local string_unpack = compat.string_unpack
+
+
 local m = {}
 local mt = {}
 
@@ -12,13 +16,12 @@ local function enforceOffset(off)
     end
 end
 
-local unpack = string.unpack
 local function unPackUoffset(bytes, off)
-    return unpack("<I4", bytes.str, off + 1)
+    return string_unpack("<I4", bytes.str, off + 1)
 end
 
 local function unPackVoffset(bytes, off)
-    return unpack("<I2", bytes.str, off + 1)
+    return string_unpack("<I2", bytes.str, off + 1)
 end
 
 function m.New(buf, pos)
diff --git a/tests/LuaTest.bat b/tests/LuaTest.bat
index 2bbf1d0e..e4e3fd8d 100644
--- a/tests/LuaTest.bat
+++ b/tests/LuaTest.bat
@@ -3,4 +3,7 @@ if "%1"=="-b" set buildtype=%2
 
 ..\%buildtype%\flatc.exe --lua -I include_test monster_test.fbs
 
+echo Run with LuaJIT:
+luajit.exe luatest.lua
+echo Run with Lua 5.3:
 lua53.exe luatest.lua
\ No newline at end of file
diff --git a/tests/LuaTest.sh b/tests/LuaTest.sh
index eb70e8a1..fdf4fddb 100755
--- a/tests/LuaTest.sh
+++ b/tests/LuaTest.sh
@@ -19,4 +19,7 @@ test_dir="$(pwd)"
 
 ${test_dir}/../flatc --lua -I include_test monster_test.fbs
 
+echo "Run with LuaJIT:"
+luajit luatest.lua
+echo "Run with Lua 5.3:"
 lua5.3 luatest.lua