fix(input): discard following keys when discarding <Cmd>/K_LUA (#36498)

Technically the current behavior does match documentation. However, the
keys following <Cmd>/K_LUA aren't normally received by vim.on_key()
callbacks either, so it does makes sense to discard them along with the
preceding key.

One may also argue that vim.on_key() callbacks should instead receive
the following keys together with the <Cmd>/K_LUA, but doing that may
cause some performance problems, and even in that case the keys should
still be discarded together.

runtime/doc/lua.txt

@@ -1383,8 +1383,10 @@ vim.on_key({fn}, {ns_id}, {opts})                               *vim.on_key()*
                  applied. {typed} may be empty if {key} is produced by
                  non-typed key(s) or by the same typed key(s) that produced a
                  previous {key}. If {fn} returns an empty string, {key} is
-                 discarded/ignored. When {fn} is `nil`, the callback
-                 associated with namespace {ns_id} is removed.
+                 discarded/ignored, and if {key} is <Cmd> then the
+                 "<Cmd>…<CR>" sequence is discarded as a whole. When {fn} is
+                 `nil`, the callback associated with namespace {ns_id} is
+                 removed.
       • {ns_id}  (`integer?`) Namespace ID. If nil or 0, generates and returns
                  a new |nvim_create_namespace()| id.
       • {opts}   (`table?`) Optional parameters

runtime/lua/vim/_editor.lua

@@ -604,7 +604,8 @@ local on_key_cbs = {} --- @type table<integer,[function, table]>
 ---          are applied, and {typed} is the key(s) before mappings are applied.
 ---          {typed} may be empty if {key} is produced by non-typed key(s) or by the
 ---          same typed key(s) that produced a previous {key}.
----          If {fn} returns an empty string, {key} is discarded/ignored.
+---          If {fn} returns an empty string, {key} is discarded/ignored, and if {key}
+---          is [<Cmd>] then the "[<Cmd>]…[<CR>]" sequence is discarded as a whole.
 ---          When {fn} is `nil`, the callback associated with namespace {ns_id} is removed.
 ---@param ns_id integer? Namespace ID. If nil or 0, generates and returns a
 ---                      new |nvim_create_namespace()| id.

src/nvim/edit.c

@@ -986,7 +986,7 @@ static int insert_handle_key(InsertState *s)
     goto check_pum;
 
   case K_LUA:
-    map_execute_lua(false);
+    map_execute_lua(false, false);
 
 check_pum:
     // nvim_select_popupmenu_item() can be called from the handling of

src/nvim/ex_getln.c

@@ -1296,7 +1296,7 @@ static int command_line_execute(VimState *state, int key)
     } else if (s->c == K_COMMAND) {
       do_cmdline(NULL, getcmdkeycmd, NULL, DOCMD_NOWAIT);
     } else {
-      map_execute_lua(false);
+      map_execute_lua(false, false);
     }
     // If the window changed incremental search state is not valid.
     if (s->is_state.winid != curwin->handle) {

src/nvim/getchar.c

@@ -1798,6 +1798,16 @@ int vgetc(void)
   // Execute Lua on_key callbacks.
   kvi_push(on_key_buf, NUL);
   if (nlua_execute_on_key(c, on_key_buf.items)) {
+    // Keys following K_COMMAND/K_LUA/K_PASTE_START aren't normally received by
+    // vim.on_key() callbacks, so discard them along with the current key.
+    if (c == K_COMMAND) {
+      xfree(getcmdkeycmd(NUL, NULL, 0, false));
+    } else if (c == K_LUA) {
+      map_execute_lua(false, true);
+    } else if (c == K_PASTE_START) {
+      paste_repeat(0);
+    }
+    // Discard the current key.
     c = K_IGNORE;
   }
   kvi_destroy(on_key_buf);
@@ -3213,9 +3223,10 @@ char *getcmdkeycmd(int promptc, void *cookie, int indent, bool do_concat)
 /// Handle a Lua mapping: get its LuaRef from typeahead and execute it.
 ///
 /// @param may_repeat  save the LuaRef for redoing with "." later
+/// @param discard     discard the keys instead of executing the LuaRef
 ///
 /// @return  false if getting the LuaRef was aborted, true otherwise
-bool map_execute_lua(bool may_repeat)
+bool map_execute_lua(bool may_repeat, bool discard)
 {
   garray_T line_ga;
   int c1 = -1;
@@ -3241,9 +3252,9 @@ bool map_execute_lua(bool may_repeat)
 
   no_mapping--;
 
-  if (aborted) {
+  if (aborted || discard) {
     ga_clear(&line_ga);
-    return false;
+    return !aborted;
   }
 
   LuaRef ref = (LuaRef)atoi(line_ga.ga_data);

src/nvim/normal.c

@@ -3191,7 +3191,7 @@ static void nv_colon(cmdarg_T *cap)
   }
 
   if (is_lua) {
-    cmd_result = map_execute_lua(true);
+    cmd_result = map_execute_lua(true, false);
   } else {
     // get a command line and execute it
     cmd_result = do_cmdline(NULL, is_cmdkey ? getcmdkeycmd : getexline, NULL,

src/nvim/terminal.c

@@ -986,7 +986,7 @@ static int terminal_execute(VimState *state, int key)
     break;
 
   case K_LUA:
-    map_execute_lua(false);
+    map_execute_lua(false, false);
     break;
 
   case Ctrl_N:

test/functional/lua/vim_spec.lua

@@ -2041,28 +2041,110 @@ stack traceback:
     end)
 
     it('can discard input', function()
-      -- discard every other normal 'x' command
+      -- discard the first key produced by every other 'x' key typed
       exec_lua [[
         n_key = 0
 
         vim.on_key(function(buf, typed_buf)
           if typed_buf == 'x' then
             n_key = n_key + 1
+            return (n_key % 2 == 0) and '' or nil
           end
-          return (n_key % 2 == 0) and "" or nil
         end)
       ]]
 
       api.nvim_buf_set_lines(0, 0, -1, true, { '54321' })
 
-      feed('x')
+      feed('x') -- 'x' not discarded
       expect('4321')
-      feed('x')
+      feed('x') -- 'x' discarded
       expect('4321')
-      feed('x')
+      feed('x') -- 'x' not discarded
       expect('321')
-      feed('x')
+      feed('x') -- 'x' discarded
       expect('321')
+
+      api.nvim_buf_set_lines(0, 0, -1, true, { '54321' })
+
+      -- only the first key from the mapping is discarded
+      command('nnoremap x $x')
+      feed('0x') -- '$' not discarded
+      expect('5432')
+      feed('0x') -- '$' discarded
+      expect('432')
+      feed('0x') -- '$' not discarded
+      expect('43')
+      feed('0x') -- '$' discarded
+      expect('3')
+
+      feed('i')
+      -- when discarding <Cmd>, the following command is also discarded.
+      command([[inoremap x <Cmd>call append('$', 'foo')<CR>]])
+      feed('x') -- not discarded
+      expect('3\nfoo')
+      feed('x') -- discarded
+      expect('3\nfoo')
+      feed('x') -- not discarded
+      expect('3\nfoo\nfoo')
+      feed('x') -- discarded
+      expect('3\nfoo\nfoo')
+
+      -- K_LUA is handled similarly to <Cmd>
+      exec_lua([[vim.keymap.set('i', 'x', function() vim.fn.append('$', 'bar') end)]])
+      feed('x') -- not discarded
+      expect('3\nfoo\nfoo\nbar')
+      feed('x') -- discarded
+      expect('3\nfoo\nfoo\nbar')
+      feed('x') -- not discarded
+      expect('3\nfoo\nfoo\nbar\nbar')
+      feed('x') -- discarded
+      expect('3\nfoo\nfoo\nbar\nbar')
+    end)
+
+    it('behaves consistently with <Cmd>, K_LUA, nvim_paste', function()
+      exec_lua([[
+        vim.keymap.set('i', '<F2>', "<Cmd>call append('$', 'FOO')<CR>")
+        vim.keymap.set('i', '<F3>', function() vim.fn.append('$', 'BAR') end)
+      ]])
+
+      feed('qrafoo<F2><F3>')
+      api.nvim_paste('bar', false, -1)
+      feed('<Esc>q')
+      expect('foobar\nFOO\nBAR')
+
+      exec_lua([[
+        keys = {}
+        typed = {}
+
+        vim.on_key(function(buf, typed_buf)
+          table.insert(keys, buf)
+          table.insert(typed, typed_buf)
+        end)
+      ]])
+
+      feed('@r')
+      local keys = exec_lua('return keys')
+      eq('@r', exec_lua([[return table.concat(typed, '')]]))
+      expect('foobarfoobar\nFOO\nBAR\nFOO\nBAR')
+
+      -- Add a new callback that discards most special keys as well as 'f'.
+      -- The old callback is still active.
+      exec_lua([[
+        vim.on_key(function(buf, _)
+          if not buf:find('^[@rao\27]$') then
+            return ''
+          end
+        end)
+
+        keys = {}
+        typed = {}
+      ]])
+
+      feed('@r')
+      eq(keys, exec_lua('return keys'))
+      eq('@r', exec_lua([[return table.concat(typed, '')]]))
+      -- The "bar" paste is discarded as a whole.
+      expect('foobarfoobaroo\nFOO\nBAR\nFOO\nBAR')
     end)
 
     it('callback invalid return', function()