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

2 files changed, 14 insertions, 4 deletions

diff --git a/src/nvim/getchar.c b/src/nvim/getchar.c
index 472bc3a850..c346bce0b7 100644
--- a/src/nvim/getchar.c
+++ b/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);
 
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 9392765f40..27ebfacc5f 100644
--- a/src/nvim/lua/executor.c
+++ b/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.