1 files changed, 34 insertions, 19 deletions

diff --git a/src/nvim/mapping.c b/src/nvim/mapping.c
index a3705d5c6f..923eea145c 100644
--- a/src/nvim/mapping.c
+++ b/src/nvim/mapping.c
@@ -276,21 +276,29 @@ static bool set_maparg_lhs_rhs(const char *const orig_lhs, const size_t orig_lhs
     mapargs->alt_lhs_len = 0;
   }
 
+  set_maparg_rhs(orig_rhs, orig_rhs_len, rhs_lua, cpo_flags, mapargs);
+
+  return true;
+}
+
+/// @see set_maparg_lhs_rhs
+static void set_maparg_rhs(const char *const orig_rhs, const size_t orig_rhs_len,
+                           const LuaRef rhs_lua, const int cpo_flags, MapArguments *const mapargs)
+{
   mapargs->rhs_lua = rhs_lua;
 
   if (rhs_lua == LUA_NOREF) {
     mapargs->orig_rhs_len = orig_rhs_len;
     mapargs->orig_rhs = xcalloc(mapargs->orig_rhs_len + 1, sizeof(char_u));
     STRLCPY(mapargs->orig_rhs, orig_rhs, mapargs->orig_rhs_len + 1);
-
     if (STRICMP(orig_rhs, "<nop>") == 0) {  // "<Nop>" means nothing
-      mapargs->rhs = xcalloc(1, sizeof(char_u));  // single null-char
+      mapargs->rhs = xcalloc(1, sizeof(char_u));  // single NUL-char
       mapargs->rhs_len = 0;
       mapargs->rhs_is_noop = true;
     } else {
       char *rhs_buf = NULL;
-      replaced = replace_termcodes(orig_rhs, orig_rhs_len, &rhs_buf, REPTERM_DO_LT, NULL,
-                                   cpo_flags);
+      char *replaced = replace_termcodes(orig_rhs, orig_rhs_len, &rhs_buf, REPTERM_DO_LT, NULL,
+                                         cpo_flags);
       mapargs->rhs_len = STRLEN(replaced);
       // NB: replace_termcodes may produce an empty string even if orig_rhs is non-empty
       // (e.g. a single ^V, see :h map-empty-rhs)
@@ -307,7 +315,6 @@ static bool set_maparg_lhs_rhs(const char *const orig_lhs, const size_t orig_lhs
                                             (char_u)KS_EXTRA, KE_LUA, rhs_lua);
     mapargs->rhs = vim_strsave((char_u *)tmp_buf);
   }
-  return true;
 }
 
 /// Parse a string of |:map-arguments| into a @ref MapArguments struct.
@@ -2000,8 +2007,7 @@ static Dictionary mapblock_fill_dict(const mapblock_T *const mp, const char *lhs
   FUNC_ATTR_NONNULL_ARG(1)
 {
   Dictionary dict = ARRAY_DICT_INIT;
-  char *const lhs = str2special_save((const char *)mp->m_keys,
-                                     compatible, !compatible);
+  char *const lhs = str2special_save((const char *)mp->m_keys, compatible, !compatible);
   char *const mapmode = map_mode_to_chars(mp->m_mode);
   varnumber_T noremap_value;
 
@@ -2114,13 +2120,15 @@ static void get_maparg(typval_T *argvars, typval_T *rettv, int exact)
     }
   } else {
     // Return a dictionary.
-    tv_dict_alloc_ret(rettv);
     if (mp != NULL && (rhs != NULL || rhs_lua != LUA_NOREF)) {
       Dictionary dict = mapblock_fill_dict(mp,
                                            did_simplify ? (char *)keys_simplified : NULL,
                                            buffer_local, true);
       (void)object_to_vim(DICTIONARY_OBJ(dict), rettv, NULL);
       api_free_dictionary(dict);
+    } else {
+      // Return an empty dictionary.
+      tv_dict_alloc_ret(rettv);
     }
   }
 
@@ -2149,29 +2157,36 @@ void f_mapset(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
   char *lhs = tv_dict_get_string(d, "lhs", false);
   char *lhsraw = tv_dict_get_string(d, "lhsraw", false);
   char *lhsrawalt = tv_dict_get_string(d, "lhsrawalt", false);
-  char *rhs = tv_dict_get_string(d, "rhs", false);
-  if (lhs == NULL || lhsraw == NULL || rhs == NULL) {
+  char *orig_rhs = tv_dict_get_string(d, "rhs", false);
+  LuaRef rhs_lua = LUA_NOREF;
+  dictitem_T *callback_di = tv_dict_find(d, S_LEN("callback"));
+  if (callback_di != NULL) {
+    Object callback_obj = vim_to_object(&callback_di->di_tv);
+    if (callback_obj.type == kObjectTypeLuaRef && callback_obj.data.luaref != LUA_NOREF) {
+      rhs_lua = callback_obj.data.luaref;
+      orig_rhs = "";
+      callback_obj.data.luaref = LUA_NOREF;
+    }
+    api_free_object(callback_obj);
+  }
+  if (lhs == NULL || lhsraw == NULL || orig_rhs == NULL) {
     emsg(_("E460: entries missing in mapset() dict argument"));
+    api_free_luaref(rhs_lua);
     return;
   }
-  char *orig_rhs = rhs;
-  char *arg_buf = NULL;
-  rhs = replace_termcodes(rhs, STRLEN(rhs), &arg_buf, REPTERM_DO_LT, NULL, CPO_TO_CPO_FLAGS);
 
-  int noremap = tv_dict_get_number(d, "noremap") ? REMAP_NONE : 0;
+  int noremap = tv_dict_get_number(d, "noremap") != 0 ? REMAP_NONE : 0;
   if (tv_dict_get_number(d, "script") != 0) {
     noremap = REMAP_SCRIPT;
   }
-  MapArguments args = {  // TODO(zeertzjq): support restoring "callback"?
-    .rhs = (char_u *)rhs,
-    .rhs_lua = LUA_NOREF,
-    .orig_rhs = vim_strsave((char_u *)orig_rhs),
+  MapArguments args = {
     .expr = tv_dict_get_number(d, "expr") != 0,
     .silent = tv_dict_get_number(d, "silent") != 0,
     .nowait = tv_dict_get_number(d, "nowait") != 0,
     .replace_keycodes = tv_dict_get_number(d, "replace_keycodes") != 0,
     .desc = tv_dict_get_string(d, "desc", false),
   };
+  set_maparg_rhs(orig_rhs, strlen(orig_rhs), rhs_lua, CPO_TO_CPO_FLAGS, &args);
   scid_T sid = (scid_T)tv_dict_get_number(d, "sid");
   linenr_T lnum = (linenr_T)tv_dict_get_number(d, "lnum");
   bool buffer = tv_dict_get_number(d, "buffer") != 0;
@@ -2182,7 +2197,7 @@ void f_mapset(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
 
   // Delete any existing mapping for this lhs and mode.
   MapArguments unmap_args = MAP_ARGUMENTS_INIT;
-  set_maparg_lhs_rhs(lhs, strlen(lhs), rhs, strlen(rhs), LUA_NOREF, 0, &unmap_args);
+  set_maparg_lhs_rhs(lhs, strlen(lhs), "", 0, LUA_NOREF, 0, &unmap_args);
   unmap_args.buffer = buffer;
   buf_do_map(MAPTYPE_UNMAP, &unmap_args, mode, false, curbuf);
   xfree(unmap_args.rhs);