feat(lua): make vim.mpack support vim.NIL and vim.empty_dict()

4 files changed, 87 insertions, 16 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index 5d600eae99..3fd3875557 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -687,6 +687,19 @@ vim.diff({a}, {b}, {opts})                                        *vim.diff()*
                     See {opts.result_type}. nil if {opts.on_hunk} is given.
 
 ------------------------------------------------------------------------------
+VIM.MPACK                                                          *lua-mpack*
+
+The *vim.mpack* module provides packing and unpacking of lua objects to
+msgpack encoded strings. |vim.NIL| and |vim.empty_dict()| are supported.
+
+vim.mpack.pack({obj})					*vim.mpack.pack*
+        Packs a lua object {obj} and returns the msgpack representation as
+        a string
+
+vim.mpack.unpack({str})					*vim.mpack.unpack*
+        Unpacks the msgpack encoded {str} and returns a lua object
+
+------------------------------------------------------------------------------
 VIM							*lua-builtin*
 
 vim.api.{func}({...})					*vim.api*
diff --git a/src/mpack/lmpack.c b/src/mpack/lmpack.c
index a66caf20ca..99207246c8 100644
--- a/src/mpack/lmpack.c
+++ b/src/mpack/lmpack.c
@@ -34,6 +34,7 @@
 #define PACKER_META_NAME "mpack.Packer"
 #define SESSION_META_NAME "mpack.Session"
 #define NIL_NAME "mpack.NIL"
+#define EMPTY_DICT_NAME "mpack.empty_dict"
 
 /* 
  * TODO(tarruda): When targeting lua 5.3 and being compiled with `long long`
@@ -55,14 +56,14 @@
 typedef struct {
   lua_State *L;
   mpack_parser_t *parser;
-  int reg, ext, unpacking;
+  int reg, ext, unpacking, mtdict;
   char *string_buffer;
 } Unpacker;
 
 typedef struct {
   lua_State *L;
   mpack_parser_t *parser;
-  int reg, ext, root, packing;
+  int reg, ext, root, packing, mtdict;
   int is_bin, is_bin_fn;
 } Packer;
 
@@ -234,7 +235,10 @@ static mpack_uint32_t lmpack_objlen(lua_State *L, int *is_array)
     len++;
   }
 
-  *is_array = isarr && max == len;
+  // when len==0, the caller should guess the type!
+  if (len > 0) {
+    *is_array = isarr && max == len;
+  }
 
 end:
   if ((size_t)-1 > (mpack_uint32_t)-1 && len > (mpack_uint32_t)-1)
@@ -268,6 +272,9 @@ static int lmpack_unpacker_new(lua_State *L)
 #endif
   rv->ext = LUA_NOREF;
 
+  lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME);
+  rv->mtdict = lmpack_ref(L, rv->reg);
+
   if (lua_istable(L, 1)) {
     /* parse options */
     lua_getfield(L, 1, "ext");
@@ -377,6 +384,11 @@ static void lmpack_parse_exit(mpack_parser_t *parser, mpack_node_t *node)
     case MPACK_TOKEN_MAP:
       lmpack_geti(L, unpacker->reg, (int)node->data[0].i);
       lmpack_unref(L, unpacker->reg, (int)node->data[0].i);
+      if (node->key_visited == 0 && node->tok.type == MPACK_TOKEN_MAP) {
+        lmpack_geti(L, unpacker->reg, unpacker->mtdict); // [table, mtdict]
+        lua_setmetatable(L, -2); // [table]
+      }
+
       break;
     default:
       break;
@@ -506,6 +518,9 @@ static int lmpack_packer_new(lua_State *L)
 #endif
   rv->ext = LUA_NOREF;
 
+  lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME);
+  rv->mtdict = lmpack_ref(L, rv->reg);
+
   if (lua_istable(L, 1)) {
     /* parse options */
     lua_getfield(L, 1, "ext");
@@ -620,11 +635,11 @@ static void lmpack_unparse_enter(mpack_parser_t *parser, mpack_node_t *node)
       break;
     }
     case LUA_TTABLE: {
-      int is_array;
       mpack_uint32_t len;
       mpack_node_t *n;
 
-      if (packer->ext != LUA_NOREF && lua_getmetatable(L, -1)) {
+      int has_meta = lua_getmetatable(L, -1);
+      if (packer->ext != LUA_NOREF && has_meta) {
         /* check if there's a handler for this metatable */
         lmpack_geti(L, packer->reg, packer->ext);
         lua_pushvalue(L, -2);
@@ -674,13 +689,24 @@ static void lmpack_unparse_enter(mpack_parser_t *parser, mpack_node_t *node)
            * -2: metatable
            * -3: original table 
            *
-           * We want to leave only the original table since it will be handled
-           * below, so pop 2
+           * We want to leave only the original table and metatable since they
+           * will be handled below, so pop 1
            */
-          lua_pop(L, 2);
+          lua_pop(L, 1);
         }
       }
 
+      int is_array = 1;
+      if (has_meta) {
+        // stack: [table, metatable]
+        if (packer->mtdict != LUA_NOREF) {
+          lmpack_geti(L, packer->reg, packer->mtdict); // [table, metatable, mtdict]
+          is_array = !lua_rawequal(L, -1, -2);
+          lua_pop(L, 1); // [table, metatable];
+        }
+        lua_pop(L, 1); // [table]
+      }
+
       /* check for cycles */
       n = node;
       while ((n = MPACK_PARENT_NODE(n))) {
@@ -1031,6 +1057,9 @@ static int lmpack_unpack(lua_State *L)
   unpacker.string_buffer = NULL;
   unpacker.L = L;
 
+  lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME);
+  unpacker.mtdict = lmpack_ref(L, unpacker.reg);
+
   result = mpack_parse(&parser, &str, &len, lmpack_parse_enter,
       lmpack_parse_exit);
 
@@ -1072,6 +1101,10 @@ static int lmpack_pack(lua_State *L)
   packer.L = L;
   packer.root = lmpack_ref(L, packer.reg);
 
+  lua_getfield(L, LUA_REGISTRYINDEX, EMPTY_DICT_NAME);
+  packer.mtdict = lmpack_ref(L, packer.reg);
+
+
   luaL_buffinit(L, &buffer);
   b = luaL_prepbuffer(&buffer);
   bl = LUAL_BUFFERSIZE;
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 2facfa28cd..d071203db1 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -507,8 +507,18 @@ static int nlua_state_init(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_setfield(lstate, -2, "__tostring");
   lua_setmetatable(lstate, -2);
   nlua_nil_ref = nlua_ref(lstate, -1);
+  lua_pushvalue(lstate, -1);
+  lua_setfield(lstate, LUA_REGISTRYINDEX, "mpack.NIL");
   lua_setfield(lstate, -2, "NIL");
 
+  // vim._empty_dict_mt
+  lua_createtable(lstate, 0, 0);
+  lua_pushcfunction(lstate, &nlua_empty_dict_tostring);
+  lua_setfield(lstate, -2, "__tostring");
+  nlua_empty_dict_ref = nlua_ref(lstate, -1);
+  lua_pushvalue(lstate, -1);
+  lua_setfield(lstate, LUA_REGISTRYINDEX, "mpack.empty_dict");
+  lua_setfield(lstate, -2, "_empty_dict_mt");
 
   // vim.mpack
   luaopen_mpack(lstate);
@@ -523,14 +533,6 @@ static int nlua_state_init(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_setfield(lstate, -2, "mpack");
   lua_pop(lstate, 3);
 
-
-  // vim._empty_dict_mt
-  lua_createtable(lstate, 0, 0);
-  lua_pushcfunction(lstate, &nlua_empty_dict_tostring);
-  lua_setfield(lstate, -2, "__tostring");
-  nlua_empty_dict_ref = nlua_ref(lstate, -1);
-  lua_setfield(lstate, -2, "_empty_dict_mt");
-
   // internal vim._treesitter... API
   nlua_add_treesitter(lstate);
 
diff --git a/test/functional/lua/mpack_spec.lua b/test/functional/lua/mpack_spec.lua
new file mode 100644
index 0000000000..ef693f01f3
--- /dev/null
+++ b/test/functional/lua/mpack_spec.lua
@@ -0,0 +1,23 @@
+-- Test suite for testing interactions with API bindings
+local helpers = require('test.functional.helpers')(after_each)
+
+local clear = helpers.clear
+local eq = helpers.eq
+local exec_lua = helpers.exec_lua
+
+describe('lua vim.mpack', function()
+  before_each(clear)
+  it('can pack vim.NIL', function()
+    eq({true, true, true, true}, exec_lua [[
+      local var = vim.mpack.unpack(vim.mpack.pack({33, vim.NIL, 77}))
+      return {var[1]==33, var[2]==vim.NIL, var[3]==77, var[4]==nil}
+    ]])
+  end)
+
+  it('can pack vim.empty_dict()', function()
+    eq({{{}, "foo", {}}, true, false}, exec_lua [[
+      local var = vim.mpack.unpack(vim.mpack.pack({{}, "foo", vim.empty_dict()}))
+      return {var, vim.tbl_islist(var[1]), vim.tbl_islist(var[3])}
+    ]])
+  end)
+end)