feat(lua): allow vim.on_key() callback to consume the key (#30939)

runtime/doc/lua.txt

@@ -1635,7 +1635,7 @@ vim.notify_once({msg}, {level}, {opts})                    *vim.notify_once()*
     Return: ~
         (`boolean`) true if message was displayed, else false
 
-vim.on_key({fn}, {ns_id})                                       *vim.on_key()*
+vim.on_key({fn}, {ns_id}, {opts})                               *vim.on_key()*
     Adds Lua function {fn} with namespace id {ns_id} as a listener to every,
     yes every, input key.
 
@@ -1649,17 +1649,19 @@ vim.on_key({fn}, {ns_id})                                       *vim.on_key()*
       • {fn} will not be cleared by |nvim_buf_clear_namespace()|
 
     Parameters: ~
-      • {fn}     (`fun(key: string, typed: string)?`) Function invoked for
-                 every input key, after mappings have been applied but before
-                 further processing. Arguments {key} and {typed} are raw
-                 keycodes, where {key} is the key after mappings 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}.
-                 When {fn} is `nil` and {ns_id} is specified, the callback
+      • {fn}     (`fun(key: string, typed: string): string??`) Function
+                 invoked for every input key, after mappings have been applied
+                 but before further processing. Arguments {key} and {typed}
+                 are raw keycodes, where {key} is the key after mappings 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. 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
 
     Return: ~
         (`integer`) Namespace id associated with {fn}. Or count of all

runtime/doc/news.txt

@@ -191,6 +191,7 @@ EVENTS
 
 • |CompleteDone| now sets the `reason` key in `v:event` which specifies the reason
   for completion being done.
+• |vim.on_key()| callbacks can consume the key by returning an empty string.
 
 LSP
 

runtime/lua/vim/_editor.lua

@@ -651,7 +651,7 @@ do
   end
 end
 
-local on_key_cbs = {} --- @type table<integer,function>
+local on_key_cbs = {} --- @type table<integer,[function, table]>
 
 --- Adds Lua function {fn} with namespace id {ns_id} as a listener to every,
 --- yes every, input key.
@@ -664,34 +664,37 @@ local on_key_cbs = {} --- @type table<integer,function>
 ---           it won't be invoked for those keys.
 ---@note {fn} will not be cleared by |nvim_buf_clear_namespace()|
 ---
----@param fn fun(key: string, typed: string)? Function invoked for every input key,
+---@param fn nil|fun(key: string, typed: string): string? Function invoked for every input key,
 ---          after mappings have been applied but before further processing. Arguments
 ---          {key} and {typed} are raw keycodes, where {key} is the key after mappings
 ---          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}.
----          When {fn} is `nil` and {ns_id} is specified, the callback associated with
----          namespace {ns_id} is removed.
+---          If {fn} returns an empty string, {key} is discarded/ignored.
+---          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.
+---@param opts table? Optional parameters
 ---
 ---@see |keytrans()|
 ---
 ---@return integer Namespace id associated with {fn}. Or count of all callbacks
 ---if on_key() is called without arguments.
-function vim.on_key(fn, ns_id)
+function vim.on_key(fn, ns_id, opts)
   if fn == nil and ns_id == nil then
     return vim.tbl_count(on_key_cbs)
   end
 
   vim.validate('fn', fn, 'callable', true)
   vim.validate('ns_id', ns_id, 'number', true)
+  vim.validate('opts', opts, 'table', true)
+  opts = opts or {}
 
   if ns_id == nil or ns_id == 0 then
     ns_id = vim.api.nvim_create_namespace('')
   end
 
-  on_key_cbs[ns_id] = fn
+  on_key_cbs[ns_id] = fn and { fn, opts }
   return ns_id
 end
 
@@ -700,12 +703,23 @@ end
 function vim._on_key(buf, typed_buf)
   local failed_ns_ids = {}
   local failed_messages = {}
+  local discard = false
   for k, v in pairs(on_key_cbs) do
-    local ok, err_msg = pcall(v, buf, typed_buf)
+    local ok, rv = pcall(v[1], buf, typed_buf)
+    if ok and rv ~= nil then
+      if type(rv) == 'string' and #rv == 0 then
+        discard = true
+        -- break   -- Without break deliver to all callbacks even when it eventually discards.
+        -- "break" does not make sense unless callbacks are sorted by ???.
+      else
+        ok = false
+        rv = 'return string must be empty'
+      end
+    end
     if not ok then
       vim.on_key(nil, k)
       table.insert(failed_ns_ids, k)
-      table.insert(failed_messages, err_msg)
+      table.insert(failed_messages, rv)
     end
   end
 
@@ -718,6 +732,7 @@ function vim._on_key(buf, typed_buf)
       )
     )
   end
+  return discard
 end
 
 --- Convert UTF-32, UTF-16 or UTF-8 {index} to byte index.

src/nvim/getchar.c

@@ -1772,7 +1772,9 @@ int vgetc(void)
 
   // Execute Lua on_key callbacks.
   kvi_push(on_key_buf, NUL);
-  nlua_execute_on_key(c, on_key_buf.items);
+  if (nlua_execute_on_key(c, on_key_buf.items)) {
+    c = K_IGNORE;
+  }
   kvi_destroy(on_key_buf);
   kvi_init(on_key_buf);
 

src/nvim/lua/executor.c

@@ -2063,12 +2063,13 @@ char *nlua_register_table_as_callable(const typval_T *const arg)
   return name;
 }
 
-void nlua_execute_on_key(int c, char *typed_buf)
+/// @return true to discard the key
+bool nlua_execute_on_key(int c, char *typed_buf)
 {
   static bool recursive = false;
 
   if (recursive) {
-    return;
+    return false;
   }
   recursive = true;
 
@@ -2097,9 +2098,15 @@ void nlua_execute_on_key(int c, char *typed_buf)
 
   int save_got_int = got_int;
   got_int = false;  // avoid interrupts when the key typed is Ctrl-C
-  if (nlua_pcall(lstate, 2, 0)) {
+  bool discard = false;
+  if (nlua_pcall(lstate, 2, 1)) {
     nlua_error(lstate,
                _("Error executing  vim.on_key Lua callback: %.*s"));
+  } else {
+    if (lua_isboolean(lstate, -1)) {
+      discard = lua_toboolean(lstate, -1);
+    }
+    lua_pop(lstate, 1);
   }
   got_int |= save_got_int;
 
@@ -2112,6 +2119,7 @@ void nlua_execute_on_key(int c, char *typed_buf)
 #endif
 
   recursive = false;
+  return discard;
 }
 
 // Sets the editor "script context" during Lua execution. Used by :verbose.

test/functional/lua/vim_spec.lua

@@ -28,6 +28,7 @@ local rmdir = n.rmdir
 local write_file = t.write_file
 local poke_eventloop = n.poke_eventloop
 local assert_alive = n.assert_alive
+local expect = n.expect
 
 describe('lua stdlib', function()
   before_each(clear)
@@ -3416,6 +3417,70 @@ describe('lua stdlib', function()
                                                                     |
       ]])
     end)
+
+    it('can discard input', function()
+      clear()
+      -- discard every other normal 'x' command
+      exec_lua [[
+        n_key = 0
+
+        vim.on_key(function(buf, typed_buf)
+          if typed_buf == 'x' then
+            n_key = n_key + 1
+          end
+          return (n_key % 2 == 0) and "" or nil
+        end)
+      ]]
+
+      api.nvim_buf_set_lines(0, 0, -1, true, { '54321' })
+
+      feed('x')
+      expect('4321')
+      feed('x')
+      expect('4321')
+      feed('x')
+      expect('321')
+      feed('x')
+      expect('321')
+    end)
+
+    it('callback invalid return', function()
+      clear()
+      -- second key produces an error which removes the callback
+      exec_lua [[
+        n_call = 0
+        vim.on_key(function(buf, typed_buf)
+          if typed_buf == 'x' then
+            n_call = n_call + 1
+          end
+          return n_call >= 2 and '!' or nil
+        end)
+      ]]
+
+      api.nvim_buf_set_lines(0, 0, -1, true, { '54321' })
+
+      local function cleanup_msg(msg)
+        return (remove_trace(msg):gsub('^Error.*\n *Messages: ', ''))
+      end
+
+      feed('x')
+      eq(1, exec_lua [[ return n_call ]])
+
+      eq(1, exec_lua [[ return vim.on_key(nil, nil) ]])
+
+      eq('', cleanup_msg(eval('v:errmsg')))
+      feed('x')
+      eq(2, exec_lua [[ return n_call ]])
+      eq('return string must be empty', cleanup_msg(eval('v:errmsg')))
+      command('let v:errmsg = ""')
+
+      eq(0, exec_lua [[ return vim.on_key(nil, nil) ]])
+
+      feed('x')
+      eq(2, exec_lua [[ return n_call ]])
+      expect('21')
+      eq('', cleanup_msg(eval('v:errmsg')))
+    end)
   end)
 
   describe('vim.wait', function()