3 files changed, 122 insertions, 22 deletions

diff --git a/src/host/buildvm_lib.c b/src/host/buildvm_lib.c
index 182ab90f..dcd3ca41 100644
--- a/src/host/buildvm_lib.c
+++ b/src/host/buildvm_lib.c
@@ -5,6 +5,7 @@
 
 #include "buildvm.h"
 #include "lj_obj.h"
+#include "lj_bc.h"
 #include "lj_lib.h"
 #include "buildvm_libbc.h"
 
@@ -152,28 +153,36 @@ static void libdef_func(BuildCtx *ctx, char *p, int arg)
   regfunc = REGFUNC_OK;
 }
 
-static uint32_t libdef_uleb128(uint8_t **pp)
+static uint8_t *libdef_uleb128(uint8_t *p, uint32_t *vv)
 {
-  uint8_t *p = *pp;
   uint32_t v = *p++;
   if (v >= 0x80) {
     int sh = 0; v &= 0x7f;
     do { v |= ((*p & 0x7f) << (sh += 7)); } while (*p++ >= 0x80);
   }
-  *pp = p;
-  return v;
+  *vv = v;
+  return p;
 }
 
-static void libdef_swapbc(uint8_t *p)
+static void libdef_fixupbc(uint8_t *p)
 {
   uint32_t i, sizebc;
   p += 4;
-  libdef_uleb128(&p);
-  libdef_uleb128(&p);
-  sizebc = libdef_uleb128(&p);
+  p = libdef_uleb128(p, &sizebc);
+  p = libdef_uleb128(p, &sizebc);
+  p = libdef_uleb128(p, &sizebc);
   for (i = 0; i < sizebc; i++, p += 4) {
-    uint8_t t = p[0]; p[0] = p[3]; p[3] = t;
-    t = p[1]; p[1] = p[2]; p[2] = t;
+    uint8_t op = p[libbc_endian ? 3 : 0];
+    uint8_t ra = p[libbc_endian ? 2 : 1];
+    uint8_t rc = p[libbc_endian ? 1 : 2];
+    uint8_t rb = p[libbc_endian ? 0 : 3];
+    if (!LJ_DUALNUM && op == BC_ISTYPE && rc == ~LJ_TNUMX+1) {
+      op = BC_ISNUM; rc++;
+    }
+    p[LJ_ENDIAN_SELECT(0, 3)] = op;
+    p[LJ_ENDIAN_SELECT(1, 2)] = ra;
+    p[LJ_ENDIAN_SELECT(2, 1)] = rc;
+    p[LJ_ENDIAN_SELECT(3, 0)] = rb;
   }
 }
 
@@ -190,8 +199,7 @@ static void libdef_lua(BuildCtx *ctx, char *p, int arg)
         *optr++ = LIBINIT_LUA;
         libdef_name(p, 0);
         memcpy(optr, libbc_code + ofs, len);
-        if (libbc_endian != LJ_BE)
-          libdef_swapbc(optr);
+        libdef_fixupbc(optr);
         optr += len;
         return;
       }
diff --git a/src/host/buildvm_libbc.h b/src/host/buildvm_libbc.h
index ee97836a..6b0e1d03 100644
--- a/src/host/buildvm_libbc.h
+++ b/src/host/buildvm_libbc.h
@@ -3,8 +3,8 @@
 static const int libbc_endian = 0;
 
 static const uint8_t libbc_code[] = {
-0,1,2,0,0,1,2,22,1,0,0,72,1,2,0,241,135,158,166,3,220,203,178,130,4,0,1,2,0,
-0,1,2,22,1,0,0,72,1,2,0,243,244,148,165,20,198,190,199,252,3,0
+0,1,2,0,0,1,2,24,1,0,0,76,1,2,0,241,135,158,166,3,220,203,178,130,4,0,1,2,0,
+0,1,2,24,1,0,0,76,1,2,0,243,244,148,165,20,198,190,199,252,3,0
 };
 
 static const struct { const char *name; int ofs; } libbc_map[] = {
diff --git a/src/host/genlibbc.lua b/src/host/genlibbc.lua
index bdcfb588..72c55d73 100644
--- a/src/host/genlibbc.lua
+++ b/src/host/genlibbc.lua
@@ -6,6 +6,15 @@
 -- Released under the MIT license. See Copyright Notice in luajit.h
 ----------------------------------------------------------------------------
 
+local ffi = require("ffi")
+local bit = require("bit")
+local vmdef = require("jit.vmdef")
+local bcnames = vmdef.bcnames
+
+local format = string.format
+
+local isbe = (string.byte(string.dump(function() end), 5) % 2 == 1)
+
 local function usage(arg)
   io.stderr:write("Usage: ", arg and arg[0] or "genlibbc",
                   " [-o buildvm_libbc.h] lib_*.c\n")
@@ -36,15 +45,100 @@ local function read_files(names)
   return src
 end
 
+local function transform_lua(code)
+  local fixup = {}
+  local n = -30000
+  code = string.gsub(code, "CHECK_(%w*)%((.-)%)", function(tp, var)
+    n = n + 1
+    fixup[n] = { "CHECK", tp }
+    return format("%s=%d", var, n)
+  end)
+  code = string.gsub(code, "PAIRS%((.-)%)", function(var)
+    fixup.PAIRS = true
+    return format("nil, %s, 0", var)
+  end)
+  return "return "..code, fixup
+end
+
+local function read_uleb128(p)
+  local v = p[0]; p = p + 1
+  if v >= 128 then
+    local sh = 7; v = v - 128
+    repeat
+      local r = p[0]
+      v = v + bit.lshift(bit.band(r, 128), sh)
+      sh = sh + 7
+      p = p + 1
+    until r < 128
+  end
+  return p, v
+end
+
+-- ORDER LJ_T
+local name2itype = {
+  str = 5, func = 9, tab = 12, int = 14, num = 15
+}
+
+local BC = {}
+for i=0,#bcnames/6-1 do
+  BC[string.gsub(string.sub(bcnames, i*6+1, i*6+6), " ", "")] = i
+end
+local xop, xra = isbe and 3 or 0, isbe and 2 or 1
+local xrc, xrb = isbe and 1 or 2, isbe and 0 or 3
+
+local function fixup_dump(dump, fixup)
+  local buf = ffi.new("uint8_t[?]", #dump+1, dump)
+  local p = buf+5
+  local n, sizebc
+  p, n = read_uleb128(p)
+  local start = p
+  p = p + 4
+  p = read_uleb128(p)
+  p = read_uleb128(p)
+  p, sizebc = read_uleb128(p)
+  local rawtab = {}
+  for i=0,sizebc-1 do
+    local op = p[xop]
+    if op == BC.KSHORT then
+      local rd = p[xrc] + 256*p[xrb]
+      rd = bit.arshift(bit.lshift(rd, 16), 16)
+      local f = fixup[rd]
+      if f then
+        if f[1] == "CHECK" then
+          local tp = f[2]
+          if tp == "tab" then rawtab[p[xra]] = true end
+          p[xop] = tp == "num" and BC.ISNUM or BC.ISTYPE
+          p[xrb] = 0
+          p[xrc] = name2itype[tp]
+        else
+          error("unhandled fixup type: "..f[1])
+        end
+      end
+    elseif op == BC.TGETV then
+      if rawtab[p[xrb]] then
+        p[xop] = BC.TGETR
+      end
+    elseif op == BC.TSETV then
+      if rawtab[p[xrb]] then
+        p[xop] = BC.TSETR
+      end
+    elseif op == BC.ITERC then
+      if fixup.PAIRS then
+        p[xop] = BC.ITERN
+      end
+    end
+    p = p + 4
+  end
+  return ffi.string(start, n)
+end
+
 local function find_defs(src)
   local defs = {}
   for name, code in string.gmatch(src, "LJLIB_LUA%(([^)]*)%)%s*/%*(.-)%*/") do
     local env = {}
-    local func = assert(load("return "..code, "", nil, env))()
-    local d = string.dump(func, true)
-    local ofs = 6
-    while string.byte(d, ofs) > 127 do ofs = ofs + 1 end
-    defs[name] = string.sub(d, ofs+1, -2)
+    local tcode, fixup = transform_lua(code)
+    local func = assert(load(tcode, "", nil, env))()
+    defs[name] = fixup_dump(string.dump(func, true), fixup)
     defs[#defs+1] = name
   end
   return defs
@@ -54,9 +148,7 @@ local function gen_header(defs)
   local t = {}
   local function w(x) t[#t+1] = x end
   w("/* This is a generated file. DO NOT EDIT! */\n\n")
-  w("static const int libbc_endian = ")
-  w(string.byte(string.dump(function() end), 5) % 2)
-  w(";\n\n")
+  w("static const int libbc_endian = ") w(isbe and 1 or 0) w(";\n\n")
   local s = ""
   for _,name in ipairs(defs) do
     s = s .. defs[name]