feat(mapset): support restoring Lua callback (#20024)

vim-patch:9.0.0341: mapset() does not restore <Nop> mapping properly

Problem:    mapset() does not restore <Nop> mapping properly.
Solution:   Use an empty string for <Nop>. (closes vim/vim#11022)
https://github.com/vim/vim/commit/92a3d20682d46359bb50a452b4f831659e799155

4 files changed, 100 insertions, 30 deletions

diff --git a/runtime/pack/dist/opt/termdebug/plugin/termdebug.vim b/runtime/pack/dist/opt/termdebug/plugin/termdebug.vim
index 802ebd42b5..bfece6aa72 100644
--- a/runtime/pack/dist/opt/termdebug/plugin/termdebug.vim
+++ b/runtime/pack/dist/opt/termdebug/plugin/termdebug.vim
@@ -964,14 +964,7 @@ func s:DeleteCommands()
       nunmap K
     else
       " call mapset(s:k_map_saved)
-      let mode = s:k_map_saved.mode !=# ' ' ? s:k_map_saved.mode : ''
-      call nvim_set_keymap(mode, 'K', s:k_map_saved.rhs, {
-        \ 'expr': s:k_map_saved.expr ? v:true : v:false,
-        \ 'noremap': s:k_map_saved.noremap ? v:true : v:false,
-        \ 'nowait': s:k_map_saved.nowait ? v:true : v:false,
-        \ 'script': s:k_map_saved.script ? v:true : v:false,
-        \ 'silent': s:k_map_saved.silent ? v:true : v:false,
-        \ })
+      call mapset('n', 0, s:k_map_saved)
     endif
     unlet s:k_map_saved
   endif
diff --git a/src/nvim/mapping.c b/src/nvim/mapping.c
index ae37141eac..5c9903aa6b 100644
--- a/src/nvim/mapping.c
+++ b/src/nvim/mapping.c
@@ -277,21 +277,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)
@@ -308,7 +316,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.
@@ -2150,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 = "";  // need non-NULL for strlen()
+      callback_obj.data.luaref = LUA_NOREF;
+    }
+    api_free_object(callback_obj);
+  }
+  if (lhs == NULL || lhsraw == NULL || (orig_rhs == NULL && rhs_lua == LUA_NOREF)) {
     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;
@@ -2183,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);
diff --git a/src/nvim/testdir/test_maparg.vim b/src/nvim/testdir/test_maparg.vim
index 17c5f433ac..f903f5b934 100644
--- a/src/nvim/testdir/test_maparg.vim
+++ b/src/nvim/testdir/test_maparg.vim
@@ -158,11 +158,11 @@ func Test_range_map()
   call assert_equal("abcd", getline(1))
 endfunc
 
-func One_mapset_test(keys)
-  exe 'nnoremap ' .. a:keys .. ' original<CR>'
+func One_mapset_test(keys, rhs)
+  exe 'nnoremap ' .. a:keys .. ' ' .. a:rhs
   let orig = maparg(a:keys, 'n', 0, 1)
   call assert_equal(a:keys, orig.lhs)
-  call assert_equal('original<CR>', orig.rhs)
+  call assert_equal(a:rhs, orig.rhs)
   call assert_equal('n', orig.mode)
 
   exe 'nunmap ' .. a:keys
@@ -172,15 +172,16 @@ func One_mapset_test(keys)
   call mapset('n', 0, orig)
   let d = maparg(a:keys, 'n', 0, 1)
   call assert_equal(a:keys, d.lhs)
-  call assert_equal('original<CR>', d.rhs)
+  call assert_equal(a:rhs, d.rhs)
   call assert_equal('n', d.mode)
 
   exe 'nunmap ' .. a:keys
 endfunc
 
 func Test_mapset()
-  call One_mapset_test('K')
-  call One_mapset_test('<F3>')
+  call One_mapset_test('K', 'original<CR>')
+  call One_mapset_test('<F3>', 'original<CR>')
+  call One_mapset_test('<F3>', '<lt>Nop>')
 
   " Check <> key conversion
   new
@@ -203,6 +204,26 @@ func Test_mapset()
 
   iunmap K
 
+  " Test that <Nop> is restored properly
+  inoremap K <Nop>
+  call feedkeys("SK\<Esc>", 'xt')
+  call assert_equal('', getline(1))
+
+  let orig = maparg('K', 'i', 0, 1)
+  call assert_equal('K', orig.lhs)
+  call assert_equal('<Nop>', orig.rhs)
+  call assert_equal('i', orig.mode)
+
+  inoremap K foo
+  call feedkeys("SK\<Esc>", 'xt')
+  call assert_equal('foo', getline(1))
+
+  call mapset('i', 0, orig)
+  call feedkeys("SK\<Esc>", 'xt')
+  call assert_equal('', getline(1))
+
+  iunmap K
+
   " Test literal <CR> using a backslash
   let cpo_save = &cpo
   set cpo-=B
diff --git a/test/functional/vimscript/map_functions_spec.lua b/test/functional/vimscript/map_functions_spec.lua
index aa64006de0..96b86d053e 100644
--- a/test/functional/vimscript/map_functions_spec.lua
+++ b/test/functional/vimscript/map_functions_spec.lua
@@ -3,6 +3,8 @@ local helpers = require('test.functional.helpers')(after_each)
 local clear = helpers.clear
 local eq = helpers.eq
 local eval = helpers.eval
+local exec = helpers.exec
+local exec_lua = helpers.exec_lua
 local expect = helpers.expect
 local feed = helpers.feed
 local funcs = helpers.funcs
@@ -10,6 +12,7 @@ local meths = helpers.meths
 local nvim = helpers.nvim
 local source = helpers.source
 local command = helpers.command
+local pcall_err = helpers.pcall_err
 
 describe('maparg()', function()
   before_each(clear)
@@ -194,4 +197,43 @@ describe('mapset()', function()
     feed('foo')
     expect('<<lt><')
   end)
+
+  it('can restore Lua callback from the dict returned by maparg()', function()
+    eq(0, exec_lua([[
+      GlobalCount = 0
+      vim.api.nvim_set_keymap('n', 'asdf', '', {callback = function() GlobalCount = GlobalCount + 1 end })
+      return GlobalCount
+    ]]))
+    feed('asdf')
+    eq(1, exec_lua([[return GlobalCount]]))
+
+    exec_lua([[
+      _G.saved_asdf_map = vim.fn.maparg('asdf', 'n', false, true)
+      vim.api.nvim_set_keymap('n', 'asdf', '', {callback = function() GlobalCount = GlobalCount + 10 end })
+    ]])
+    feed('asdf')
+    eq(11, exec_lua([[return GlobalCount]]))
+
+    exec_lua([[vim.fn.mapset('n', false, _G.saved_asdf_map)]])
+    feed('asdf')
+    eq(12, exec_lua([[return GlobalCount]]))
+
+    exec([[
+      let g:saved_asdf_map = maparg('asdf', 'n', v:false, v:true)
+      lua vim.api.nvim_set_keymap('n', 'asdf', '', {callback = function() GlobalCount = GlobalCount + 10 end })
+    ]])
+    feed('asdf')
+    eq(22, exec_lua([[return GlobalCount]]))
+
+    command([[call mapset('n', v:false, g:saved_asdf_map)]])
+    feed('asdf')
+    eq(23, exec_lua([[return GlobalCount]]))
+  end)
+
+  it('does not leak memory if lhs is missing', function()
+    eq('Error executing lua: Vim:E460: entries missing in mapset() dict argument',
+       pcall_err(exec_lua, [[vim.fn.mapset('n', false, {rhs = 'foo'})]]))
+    eq('Error executing lua: Vim:E460: entries missing in mapset() dict argument',
+       pcall_err(exec_lua, [[vim.fn.mapset('n', false, {callback = function() end})]]))
+  end)
 end)