feat(lsp): improve json deserialization performance (#15854)

mjlbach · web-flow · commit 912a6e5a9c58 · 2021-10-05T08:37:20.000-07:00
* Add optional second table argument to vim.json.decode which takes
  a table 'luanil' which can include the 'object' and/or 'array' keys. These
  options use luanil when converting NULL in json objects and arrays
  respectively. The default behavior matches the original lua-cjson.
* Remove recursive_convert_NIL function from rpc.lua, use
  vim.json.decode with luanil = { object = true } instead. This removes a hotpath
  in the json deserialization pipeline by dropping keys with json NULL
  values throughout the deserialized table.
diff --git a/runtime/lua/vim/lsp/rpc.lua b/runtime/lua/vim/lsp/rpc.lua
@@ -13,36 +13,6 @@ local function is_dir(filename)
   return stat and stat.type == 'directory' or false
 end
 
-local NIL = vim.NIL
-
----@private
-local recursive_convert_NIL
-recursive_convert_NIL = function(v, tbl_processed)
-  if v == NIL then
-    return nil
-  elseif not tbl_processed[v] and type(v) == 'table' then
-    tbl_processed[v] = true
-    local inside_list = vim.tbl_islist(v)
-    return vim.tbl_map(function(x)
-      if not inside_list or (inside_list and type(x) == "table") then
-        return recursive_convert_NIL(x, tbl_processed)
-      else
-        return x
-      end
-    end, v)
-  end
-
-  return v
-end
-
----@private
---- Returns its argument, but converts `vim.NIL` to Lua `nil`.
----@param v (any) Argument
----@returns (any)
-local function convert_NIL(v)
-  return recursive_convert_NIL(v, {})
-end
-
 ---@private
 --- Merges current process env with the given env and returns the result as
 --- a list of "k=v" strings.
@@ -457,7 +427,7 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
 
   ---@private
   local function handle_body(body)
-    local ok, decoded = pcall(vim.json.decode, body)
+    local ok, decoded = pcall(vim.json.decode, body, { luanil = { object = true } })
     if not ok then
       on_error(client_errors.INVALID_SERVER_JSON, decoded)
       return
@@ -466,8 +436,6 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
 
     if type(decoded.method) == 'string' and decoded.id then
       local err
-      -- Server Request
-      decoded.params = convert_NIL(decoded.params)
       -- Schedule here so that the users functions don't trigger an error and
       -- we can still use the result.
       schedule(function()
@@ -494,9 +462,6 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
       end)
     -- This works because we are expecting vim.NIL here
     elseif decoded.id and (decoded.result ~= vim.NIL or decoded.error ~= vim.NIL) then
-      -- Server Result
-      decoded.error = convert_NIL(decoded.error)
-      decoded.result = convert_NIL(decoded.result)
 
       -- We sent a number, so we expect a number.
       local result_id = tonumber(decoded.id)
@@ -544,7 +509,6 @@ local function start(cmd, cmd_args, dispatchers, extra_spawn_params)
       end
     elseif type(decoded.method) == 'string' then
       -- Notification
-      decoded.params = convert_NIL(decoded.params)
       try_call(client_errors.NOTIFICATION_HANDLER_ERROR,
           dispatchers.notification, decoded.method, decoded.params)
     else
diff --git a/src/cjson/lua_cjson.c b/src/cjson/lua_cjson.c
@@ -170,11 +170,19 @@ typedef struct {
     int decode_array_with_array_mt;
 } json_config_t;
 
+typedef struct {
+    /* convert null in json objects to lua nil instead of vim.NIL */
+    int luanil_object;
+    /* convert null in json arrays to lua nil instead of vim.NIL */
+    int luanil_array;
+} json_options_t;
+
 typedef struct {
     const char *data;
     const char *ptr;
     strbuf_t *tmp;    /* Temporary storage for strings */
     json_config_t *cfg;
+    json_options_t *options;
     int current_depth;
 } json_parse_t;
 
@@ -865,7 +873,7 @@ static int json_encode(lua_State *l)
 /* ===== DECODING ===== */
 
 static void json_process_value(lua_State *l, json_parse_t *json,
-                               json_token_t *token);
+                               json_token_t *token, bool use_luanil);
 
 static int hexdigit2int(char hex)
 {
@@ -1296,7 +1304,7 @@ static void json_parse_object_context(lua_State *l, json_parse_t *json)
 
         /* Fetch value */
         json_next_token(json, &token);
-        json_process_value(l, json, &token);
+        json_process_value(l, json, &token, json->options->luanil_object);
 
         /* Set key = value */
         lua_rawset(l, -3);
@@ -1343,7 +1351,7 @@ static void json_parse_array_context(lua_State *l, json_parse_t *json)
     }
 
     for (i = 1; ; i++) {
-        json_process_value(l, json, &token);
+        json_process_value(l, json, &token, json->options->luanil_array);
         lua_rawseti(l, -2, i);            /* arr[i] = value */
 
         json_next_token(json, &token);
@@ -1362,7 +1370,7 @@ static void json_parse_array_context(lua_State *l, json_parse_t *json)
 
 /* Handle the "value" context */
 static void json_process_value(lua_State *l, json_parse_t *json,
-                               json_token_t *token)
+                               json_token_t *token, bool use_luanil)
 {
     switch (token->type) {
     case T_STRING:
@@ -1381,7 +1389,11 @@ static void json_process_value(lua_State *l, json_parse_t *json,
         json_parse_array_context(l, json);
         break;;
     case T_NULL:
-        nlua_pushref(l, nlua_nil_ref);
+        if (use_luanil) {
+            lua_pushnil(l);
+        } else {
+            nlua_pushref(l, nlua_nil_ref);
+        }
         break;;
     default:
         json_throw_parse_error(l, json, "value", token);
@@ -1392,12 +1404,46 @@ static int json_decode(lua_State *l)
 {
     json_parse_t json;
     json_token_t token;
+    json_options_t options = { .luanil_object = false, .luanil_array = false };
+
+
     size_t json_len;
 
-    luaL_argcheck(l, lua_gettop(l) == 1, 1, "expected 1 argument");
+    switch (lua_gettop(l)) {
+    case 1:
+        break;
+    case 2:
+        luaL_checktype(l, 2, LUA_TTABLE);
+        lua_getfield(l, 2, "luanil");
+
+        /* We only handle the luanil option for now */
+        if (lua_isnil(l, -1)) { 
+            lua_pop(l, 1); 
+            break; 
+        }
+
+        luaL_checktype(l, -1, LUA_TTABLE);
+
+        lua_getfield(l, -1, "object");
+        if (!lua_isnil(l, -1)) {
+            options.luanil_object = true;
+        }
+        lua_pop(l, 1);
+
+        lua_getfield(l, -1, "array");
+        if (!lua_isnil(l, -1)) {
+            options.luanil_array = true;
+        }
+        /* Also pop the luanil table */
+        lua_pop(l, 2);
+        break;
+    default:
+        return luaL_error (l, "expected 1 or 2 arguments");
+    }
 
     json.cfg = json_fetch_config(l);
     json.data = luaL_checklstring(l, 1, &json_len);
+    json.options = &options;
     json.current_depth = 0;
     json.ptr = json.data;
 
@@ -1415,7 +1461,7 @@ static int json_decode(lua_State *l)
     json.tmp = strbuf_new(json_len);
 
     json_next_token(&json, &token);
-    json_process_value(l, &json, &token);
+    json_process_value(l, &json, &token, json.options->luanil_object);
 
     /* Ensure there is no more input left */
     json_next_token(&json, &token);
