lua: add vim.register_keystroke_callback (#12536)

* feat: Add vim.register_keystroke_callback

* fixup: Forgot to remove mention of old option

* fixup: Answer jamessan comments

* fixup: Answer norcalli comments

* fixup: portability

* Update runtime/doc/lua.txt

Co-authored-by: Ashkan Kiani <ashkan.k.kiani@gmail.com>

6 files changed, 235 insertions, 1 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index aa9addece8..f336ba0c36 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -928,6 +928,34 @@ vim.region({bufnr}, {pos1}, {pos2}, {type}, {inclusive})       *vim.region()*
         whether the selection is inclusive or not, into a zero-indexed table 
         of linewise selections of the form `{linenr = {startcol, endcol}}` .
 
+                                             *vim.register_keystroke_callback()*
+vim.register_keystroke_callback({fn}, {ns_id})
+        Register a lua {fn} with an {ns_id} to be run after every keystroke.
+
+        Parameters: ~
+            {fn}: (function): Function to call on keystroke.
+                    It should take one argument, which is a string.
+                    The string will contain the literal keys typed.
+                    See |i_CTRL-V|
+
+                    If {fn} is `nil`, it removes the callback for the
+                    associated {ns_id}.
+
+            {ns_id}: (number)  Namespace ID. If not passed or 0, will generate
+                     and return a new namespace ID from |nvim_create_namespace()|
+
+        Return: ~
+            (number) Namespace ID associated with {fn}
+
+        NOTE: {fn} will be automatically removed if an error occurs while
+        calling. This is to prevent the annoying situation of every keystroke
+        erroring while trying to remove a broken callback.
+
+        NOTE: {fn} will receive the keystrokes after mappings have been
+        evaluated
+
+        NOTE: {fn} will *NOT* be cleared from |nvim_buf_clear_namespace()|
+
 vim.rpcnotify({channel}, {method}[, {args}...])		    *vim.rpcnotify()*
 	Sends {event} to {channel} via |RPC| and returns immediately.
 	If {channel} is 0, the event is broadcast to all channels.
diff --git a/src/nvim/getchar.c b/src/nvim/getchar.c
index 5ab5a7db1b..dc11e4a232 100644
--- a/src/nvim/getchar.c
+++ b/src/nvim/getchar.c
@@ -27,6 +27,7 @@
 #include "nvim/ex_docmd.h"
 #include "nvim/ex_getln.h"
 #include "nvim/func_attr.h"
+#include "nvim/lua/executor.h"
 #include "nvim/main.h"
 #include "nvim/mbyte.h"
 #include "nvim/memline.h"
@@ -1535,6 +1536,9 @@ int vgetc(void)
    */
   may_garbage_collect = false;
 
+  // Exec lua callbacks for on_keystroke
+  nlua_execute_log_keystroke(c);
+
   return c;
 }
 
diff --git a/src/nvim/keymap.c b/src/nvim/keymap.c
index a553110552..2b6f022d9d 100644
--- a/src/nvim/keymap.c
+++ b/src/nvim/keymap.c
@@ -530,13 +530,24 @@ unsigned int trans_special(const char_u **srcp, const size_t src_len,
 {
   int modifiers = 0;
   int key;
-  unsigned int dlen = 0;
 
   key = find_special_key(srcp, src_len, &modifiers, keycode, false, in_string);
   if (key == 0) {
     return 0;
   }
 
+  return special_to_buf(key, modifiers, keycode, dst);
+}
+
+/// Put the character sequence for "key" with "modifiers" into "dst" and return
+/// the resulting length.
+/// When "keycode" is TRUE prefer key code, e.g. K_DEL instead of DEL.
+/// The sequence is not NUL terminated.
+/// This is how characters in a string are encoded.
+unsigned int special_to_buf(int key, int modifiers, bool keycode, char_u *dst)
+{
+  unsigned int dlen = 0;
+
   // Put the appropriate modifier in a string.
   if (modifiers != 0) {
     dst[dlen++] = K_SPECIAL;
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 86da517685..5ad9731a97 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -1465,3 +1465,40 @@ void nlua_free_typval_dict(dict_T *const d)
     d->lua_table_ref = LUA_NOREF;
   }
 }
+
+void nlua_execute_log_keystroke(int c)
+{
+  char_u buf[NUMBUFLEN];
+  size_t buf_len = special_to_buf(c, mod_mask, false, buf);
+
+  lua_State *const lstate = nlua_enter();
+
+#ifndef NDEBUG
+  int top = lua_gettop(lstate);
+#endif
+
+  // [ vim ]
+  lua_getglobal(lstate, "vim");
+
+  // [ vim, vim._log_keystroke ]
+  lua_getfield(lstate, -1, "_log_keystroke");
+  luaL_checktype(lstate, -1, LUA_TFUNCTION);
+
+  // [ vim, vim._log_keystroke, buf ]
+  lua_pushlstring(lstate, (const char *)buf, buf_len);
+
+  if (lua_pcall(lstate, 1, 0, 0)) {
+    nlua_error(
+        lstate,
+        _("Error executing vim.log_keystroke lua callback: %.*s"));
+  }
+
+  // [ vim ]
+  lua_pop(lstate, 1);
+
+#ifndef NDEBUG
+  // [ ]
+  assert(top == lua_gettop(lstate));
+#endif
+}
+
diff --git a/src/nvim/lua/vim.lua b/src/nvim/lua/vim.lua
index 820b237c4f..bfa8b91208 100644
--- a/src/nvim/lua/vim.lua
+++ b/src/nvim/lua/vim.lua
@@ -489,4 +489,60 @@ function vim.defer_fn(fn, timeout)
   return timer
 end
 
+local on_keystroke_callbacks = {}
+
+--- Register a lua {fn} with an {id} to be run after every keystroke.
+---
+--@param fn function: Function to call. It should take one argument, which is a string.
+---                   The string will contain the literal keys typed.
+---                   See |i_CTRL-V|
+---
+---                   If {fn} is nil, it removes the callback for the associated {ns_id}
+--@param ns_id number? Namespace ID. If not passed or 0, will generate and return a new
+---                    namespace ID from |nvim_create_namesapce()|
+---
+--@return number Namespace ID associated with {fn}
+---
+--@note {fn} will be automatically removed if an error occurs while calling.
+---     This is to prevent the annoying situation of every keystroke erroring
+---     while trying to remove a broken callback.
+--@note {fn} will not be cleared from |nvim_buf_clear_namespace()|
+--@note {fn} will receive the keystrokes after mappings have been evaluated
+function vim.register_keystroke_callback(fn, ns_id)
+  vim.validate {
+    fn = { fn, 'c', true},
+    ns_id = { ns_id, 'n', true }
+  }
+
+  if ns_id == nil or ns_id == 0 then
+    ns_id = vim.api.nvim_create_namespace('')
+  end
+
+  on_keystroke_callbacks[ns_id] = fn
+  return ns_id
+end
+
+--- Function that executes the keystroke callbacks.
+--@private
+function vim._log_keystroke(char)
+  local failed_ns_ids = {}
+  local failed_messages = {}
+  for k, v in pairs(on_keystroke_callbacks) do
+    local ok, err_msg = pcall(v, char)
+    if not ok then
+      vim.register_keystroke_callback(nil, k)
+
+      table.insert(failed_ns_ids, k)
+      table.insert(failed_messages, err_msg)
+    end
+  end
+
+  if failed_ns_ids[1] then
+    error(string.format(
+      "Error executing 'on_keystroke' with ns_ids of '%s'\n    With messages: %s",
+      table.concat(failed_ns_ids, ", "),
+      table.concat(failed_messages, "\n")))
+  end
+end
+
 return module
diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index 9b2697b3c2..a9e8ca9686 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -1068,6 +1068,104 @@ describe('lua stdlib', function()
     eq({5,15}, exec_lua[[ return vim.region(0,{1,5},{1,14},'v',true)[1] ]])
   end)
 
+  describe('vim.execute_on_keystroke', function()
+    it('should keep track of keystrokes', function()
+      helpers.insert([[hello world ]])
+
+      exec_lua [[
+        KeysPressed = {}
+
+        vim.register_keystroke_callback(function(buf)
+          if buf:byte() == 27 then
+            buf = "<ESC>"
+          end
+
+          table.insert(KeysPressed, buf)
+        end)
+      ]]
+
+      helpers.insert([[next 🤦 lines å ]])
+
+      -- It has escape in the keys pressed
+      eq('inext 🤦 lines å <ESC>', exec_lua [[return table.concat(KeysPressed, '')]])
+    end)
+
+    it('should allow removing trackers.', function()
+      helpers.insert([[hello world]])
+
+      exec_lua [[
+        KeysPressed = {}
+
+        return vim.register_keystroke_callback(function(buf)
+          if buf:byte() == 27 then
+            buf = "<ESC>"
+          end
+
+          table.insert(KeysPressed, buf)
+        end, vim.api.nvim_create_namespace("logger"))
+      ]]
+
+      helpers.insert([[next lines]])
+
+      exec_lua("vim.register_keystroke_callback(nil, vim.api.nvim_create_namespace('logger'))")
+
+      helpers.insert([[more lines]])
+
+      -- It has escape in the keys pressed
+      eq('inext lines<ESC>', exec_lua [[return table.concat(KeysPressed, '')]])
+    end)
+
+    it('should not call functions that error again.', function()
+      helpers.insert([[hello world]])
+
+      exec_lua [[
+        KeysPressed = {}
+
+        return vim.register_keystroke_callback(function(buf)
+          if buf:byte() == 27 then
+            buf = "<ESC>"
+          end
+
+          table.insert(KeysPressed, buf)
+
+          if buf == 'l' then
+            error("Dumb Error")
+          end
+        end)
+      ]]
+
+      helpers.insert([[next lines]])
+      helpers.insert([[more lines]])
+
+      -- Only the first letter gets added. After that we remove the callback
+      eq('inext l', exec_lua [[ return table.concat(KeysPressed, '') ]])
+    end)
+
+    it('should process mapped keys, not unmapped keys', function()
+      exec_lua [[
+        KeysPressed = {}
+
+        vim.cmd("inoremap hello world")
+
+        vim.register_keystroke_callback(function(buf)
+          if buf:byte() == 27 then
+            buf = "<ESC>"
+          end
+
+          table.insert(KeysPressed, buf)
+        end)
+      ]]
+
+      helpers.insert("hello")
+
+      local next_status = exec_lua [[
+        return table.concat(KeysPressed, '')
+      ]]
+
+      eq("iworld<ESC>", next_status)
+    end)
+  end)
+
   describe('vim.wait', function()
     before_each(function()
       exec_lua[[