feat(lua): pass keys before mapping to vim.on_key() callback (#28098)

Keys before mapping (i.e. typed keys) are passed as the second argument.

runtime/doc/lua.txt

@@ -1640,12 +1640,14 @@ vim.on_key({fn}, {ns_id})                                       *vim.on_key()*
     Note: ~
       • {fn} will be removed on error.
       • {fn} will not be cleared by |nvim_buf_clear_namespace()|
-      • {fn} will receive the keys after mappings have been evaluated
 
     Parameters: ~
-      • {fn}     (`fun(key: string)?`) Function invoked on every key press.
-                 |i_CTRL-V| Passing in nil when {ns_id} is specified removes
-                 the callback associated with namespace {ns_id}.
+      • {fn}     (`fun(key: string, typed: string)?`) Function invoked on
+                 every key press. |i_CTRL-V| {key} is the key after mappings
+                 have been applied, and {typed} is the key(s) before mappings
+                 are applied, which may be empty if {key} is produced by
+                 non-typed keys. When {fn} is nil and {ns_id} is specified,
+                 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.
 

runtime/doc/news.txt

@@ -354,6 +354,9 @@ The following changes to existing APIs or features add new behavior.
 
 • |vim.region()| can use a string accepted by |getpos()| as position.
 
+• |vim.on_key()| callbacks receive a second argument for keys typed before
+  mappings are applied.
+
 • |vim.diagnostic.config()| now accepts a function for the virtual_text.prefix
   option, which allows for rendering e.g., diagnostic severities differently.
 

runtime/lua/vim/_editor.lua

@@ -654,11 +654,14 @@ local on_key_cbs = {} --- @type table<integer,function>
 ---
 ---@note {fn} will be removed on error.
 ---@note {fn} will not be cleared by |nvim_buf_clear_namespace()|
----@note {fn} will receive the keys after mappings have been evaluated
 ---
----@param fn fun(key: string)? Function invoked on every key press. |i_CTRL-V|
----                   Passing in nil when {ns_id} is specified removes the
----                   callback associated with namespace {ns_id}.
+---@param fn fun(key: string, typed: string)?
+---                   Function invoked on every key press. |i_CTRL-V|
+---                   {key} is the key after mappings have been applied, and
+---                   {typed} is the key(s) before mappings are applied, which
+---                   may be empty if {key} is produced by non-typed keys.
+---                   When {fn} is nil and {ns_id} is specified, 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.
 ---
@@ -684,11 +687,11 @@ end
 
 --- Executes the on_key callbacks.
 ---@private
-function vim._on_key(char)
+function vim._on_key(buf, typed_buf)
   local failed_ns_ids = {}
   local failed_messages = {}
   for k, v in pairs(on_key_cbs) do
-    local ok, err_msg = pcall(v, char)
+    local ok, err_msg = pcall(v, buf, typed_buf)
     if not ok then
       vim.on_key(nil, k)
       table.insert(failed_ns_ids, k)

src/klib/kvec.h

@@ -153,6 +153,12 @@
     type init_array[INIT_SIZE]; \
   }
 
+#define KVI_INITIAL_VALUE(v) { \
+  .size = 0, \
+  .capacity = ARRAY_SIZE((v).init_array), \
+  .items = (v).init_array \
+}
+
 /// Initialize vector with preallocated array
 ///
 /// @param[out]  v  Vector to initialize.
@@ -218,6 +224,17 @@ static inline void *_memcpy_free(void *const restrict dest, void *const restrict
     } \
   } while (0)
 
+#define kvi_concat_len(v, data, len) \
+  if (len > 0) { \
+    kvi_ensure_more_space(v, len); \
+    assert((v).items); \
+    memcpy((v).items + (v).size, data, sizeof((v).items[0]) * len); \
+    (v).size = (v).size + len; \
+  }
+
+#define kvi_concat(v, str) kvi_concat_len(v, str, strlen(str))
+#define kvi_splice(v1, v0) kvi_concat_len(v1, (v0).items, (v0).size)
+
 /// Get location where to store new element to a vector with preallocated array
 ///
 /// @param[in,out]  v  Vector to push to.

src/nvim/getchar.c

@@ -94,6 +94,12 @@ static buffheader_T readbuf1 = { { NULL, { NUL } }, NULL, 0, 0 };
 /// Second read ahead buffer. Used for redo.
 static buffheader_T readbuf2 = { { NULL, { NUL } }, NULL, 0, 0 };
 
+/// Buffer used to store typed characters for vim.on_key().
+static kvec_withinit_t(char, MAXMAPLEN) on_key_buf = KVI_INITIAL_VALUE(on_key_buf);
+
+/// Number of following bytes that should not be stored for vim.on_key().
+static size_t no_on_key_len = 0;
+
 static int typeahead_char = 0;  ///< typeahead char that's not flushed
 
 /// When block_redo is true the redo buffer will not be changed.
@@ -994,14 +1000,19 @@ int ins_typebuf(char *str, int noremap, int offset, bool nottyped, bool silent)
 /// Uses cmd_silent, KeyTyped and KeyNoremap to restore the flags belonging to
 /// the char.
 ///
+/// @param no_on_key don't store these bytes for vim.on_key()
+///
 /// @return the length of what was inserted
-int ins_char_typebuf(int c, int modifiers)
+int ins_char_typebuf(int c, int modifiers, bool no_on_key)
 {
   char buf[MB_MAXBYTES * 3 + 4];
   unsigned len = special_to_buf(c, modifiers, true, buf);
   assert(len < sizeof(buf));
   buf[len] = NUL;
   ins_typebuf(buf, KeyNoremap, 0, !KeyTyped, cmd_silent);
+  if (KeyTyped && no_on_key) {
+    no_on_key_len += len;
+  }
   return (int)len;
 }
 
@@ -1162,12 +1173,22 @@ static void gotchars(const uint8_t *chars, size_t len)
       updatescript(buf[i]);
     }
 
-    if (reg_recording != 0) {
     buf[buflen] = NUL;
+
+    if (reg_recording != 0) {
       add_buff(&recordbuff, (char *)buf, (ptrdiff_t)buflen);
       // remember how many chars were last recorded
       last_recorded_len += buflen;
     }
+
+    if (buflen > no_on_key_len) {
+      vim_unescape_ks((char *)buf + no_on_key_len);
+      kvi_concat(on_key_buf, (char *)buf + no_on_key_len);
+      no_on_key_len = 0;
+    } else {
+      no_on_key_len -= buflen;
+    }
+
     buflen = 0;
   }
 
@@ -1185,6 +1206,7 @@ static void gotchars(const uint8_t *chars, size_t len)
 void gotchars_ignore(void)
 {
   uint8_t nop_buf[3] = { K_SPECIAL, KS_EXTRA, KE_IGNORE };
+  no_on_key_len += 3;
   gotchars(nop_buf, 3);
 }
 
@@ -1655,9 +1677,13 @@ int vgetc(void)
       if (!no_mapping && KeyTyped && mod_mask == MOD_MASK_ALT && !(State & MODE_TERMINAL)
           && !is_mouse_key(c)) {
         mod_mask = 0;
-        int len = ins_char_typebuf(c, 0);
-        ins_char_typebuf(ESC, 0);
-        ungetchars(len + 3);  // K_SPECIAL KS_MODIFIER MOD_MASK_ALT takes 3 more bytes
+        int len = ins_char_typebuf(c, 0, false);
+        ins_char_typebuf(ESC, 0, false);
+        int old_len = len + 3;  // K_SPECIAL KS_MODIFIER MOD_MASK_ALT takes 3 more bytes
+        ungetchars(old_len);
+        if (on_key_buf.size >= (size_t)old_len) {
+          on_key_buf.size -= (size_t)old_len;
+        }
         continue;
       }
 
@@ -1673,7 +1699,9 @@ int vgetc(void)
   may_garbage_collect = false;
 
   // Execute Lua on_key callbacks.
-  nlua_execute_on_key(c);
+  nlua_execute_on_key(c, on_key_buf.items, on_key_buf.size);
+  kvi_destroy(on_key_buf);
+  kvi_init(on_key_buf);
 
   // Need to process the character before we know it's safe to do something
   // else.

src/nvim/lua/executor.c

@@ -2064,7 +2064,7 @@ char *nlua_register_table_as_callable(const typval_T *const arg)
   return name;
 }
 
-void nlua_execute_on_key(int c)
+void nlua_execute_on_key(int c, char *typed_buf, size_t typed_len)
 {
   char buf[MB_MAXBYTES * 3 + 4];
   size_t buf_len = special_to_buf(c, mod_mask, false, buf);
@@ -2085,9 +2085,12 @@ void nlua_execute_on_key(int c)
   // [ vim, vim._on_key, buf ]
   lua_pushlstring(lstate, buf, buf_len);
 
+  // [ vim, vim._on_key, buf, typed_buf ]
+  lua_pushlstring(lstate, typed_buf, typed_len);
+
   int save_got_int = got_int;
   got_int = false;  // avoid interrupts when the key typed is Ctrl-C
-  if (nlua_pcall(lstate, 1, 0)) {
+  if (nlua_pcall(lstate, 2, 0)) {
     nlua_error(lstate,
                _("Error executing  vim.on_key Lua callback: %.*s"));
   }

src/nvim/message.c

@@ -1265,7 +1265,7 @@ void wait_return(int redraw)
     } else if (vim_strchr("\r\n ", c) == NULL && c != Ctrl_C) {
       // Put the character back in the typeahead buffer.  Don't use the
       // stuff buffer, because lmaps wouldn't work.
-      ins_char_typebuf(vgetc_char, vgetc_mod_mask);
+      ins_char_typebuf(vgetc_char, vgetc_mod_mask, true);
       do_redraw = true;             // need a redraw even though there is
                                     // typeahead
     }
@@ -3431,7 +3431,7 @@ int do_dialog(int type, const char *title, const char *message, const char *butt
       }
       if (c == ':' && ex_cmd) {
         retval = dfltbutton;
-        ins_char_typebuf(':', 0);
+        ins_char_typebuf(':', 0, false);
         break;
       }
 

src/nvim/normal.c

@@ -1078,7 +1078,7 @@ static int normal_execute(VimState *state, int key)
     // When "restart_edit" is set fake a "d"elete command, Insert mode will restart automatically.
     // Insert the typed character in the typeahead buffer, so that it can
     // be mapped in Insert mode.  Required for ":lmap" to work.
-    int len = ins_char_typebuf(vgetc_char, vgetc_mod_mask);
+    int len = ins_char_typebuf(vgetc_char, vgetc_mod_mask, true);
 
     // When recording and gotchars() was called the character will be
     // recorded again, remove the previous recording.

src/nvim/terminal.c

@@ -1631,7 +1631,7 @@ end:
     return false;
   }
 
-  int len = ins_char_typebuf(vgetc_char, vgetc_mod_mask);
+  int len = ins_char_typebuf(vgetc_char, vgetc_mod_mask, true);
   if (KeyTyped) {
     ungetchars(len);
   }

test/functional/editor/meta_key_spec.lua

@@ -141,4 +141,29 @@ describe('meta-keys #8226 #13042', function()
       // This is some text: bar
       // This is some text: baz]])
   end)
+
+  it('ALT/META with vim.on_key()', function()
+    feed('ifoo<CR>bar<CR>baz<Esc>gg0')
+
+    exec_lua [[
+      keys = {}
+      typed = {}
+
+      vim.on_key(function(buf, typed_buf)
+        table.insert(keys, vim.fn.keytrans(buf))
+        table.insert(typed, vim.fn.keytrans(typed_buf))
+      end)
+    ]]
+
+    -- <M-"> is reinterpreted as <Esc>"
+    feed('qrviw"ayc$FOO.<M-">apq')
+    expect([[
+      FOO.foo
+      bar
+      baz]])
+
+    -- vim.on_key() callback should only receive <Esc>"
+    eq('qrviw"ayc$FOO.<Esc>"apq', exec_lua [[return table.concat(keys, '')]])
+    eq('qrviw"ayc$FOO.<Esc>"apq', exec_lua [[return table.concat(typed, '')]])
+  end)
 end)

test/functional/lua/vim_spec.lua

@@ -3012,13 +3012,18 @@ describe('lua stdlib', function()
 
       exec_lua [[
         keys = {}
+        typed = {}
 
-        vim.on_key(function(buf)
+        vim.on_key(function(buf, typed_buf)
           if buf:byte() == 27 then
             buf = "<ESC>"
           end
+          if typed_buf:byte() == 27 then
+            typed_buf = "<ESC>"
+          end
 
           table.insert(keys, buf)
+          table.insert(typed, typed_buf)
         end)
       ]]
 
@@ -3026,20 +3031,41 @@ describe('lua stdlib', function()
 
       -- It has escape in the keys pressed
       eq('inext 🤦 lines å …<ESC>', exec_lua [[return table.concat(keys, '')]])
+      eq('inext 🤦 lines å …<ESC>', exec_lua [[return table.concat(typed, '')]])
     end)
 
     it('tracks input with modifiers', function()
       exec_lua [[
         keys = {}
+        typed = {}
 
-        vim.on_key(function(buf)
+        vim.on_key(function(buf, typed_buf)
           table.insert(keys, vim.fn.keytrans(buf))
+          table.insert(typed, vim.fn.keytrans(typed_buf))
         end)
       ]]
 
       feed([[i<C-V><C-;><C-V><C-…><Esc>]])
 
       eq('i<C-V><C-;><C-V><C-…><Esc>', exec_lua [[return table.concat(keys, '')]])
+      eq('i<C-V><C-;><C-V><C-…><Esc>', exec_lua [[return table.concat(typed, '')]])
+    end)
+
+    it('works with character find and Select mode', function()
+      insert('12345')
+
+      exec_lua [[
+        typed = {}
+
+        vim.cmd('snoremap # @')
+
+        vim.on_key(function(buf, typed_buf)
+          table.insert(typed, vim.fn.keytrans(typed_buf))
+        end)
+      ]]
+
+      feed('F3gHβγδεζ<Esc>gH…<Esc>gH#$%^')
+      eq('F3gHβγδεζ<Esc>gH…<Esc>gH#$%^', exec_lua [[return table.concat(typed, '')]])
     end)
 
     it('allows removing on_key listeners', function()
@@ -3101,23 +3127,29 @@ describe('lua stdlib', function()
       eq('inext l', exec_lua [[ return table.concat(keys, '') ]])
     end)
 
-    it('processes mapped keys, not unmapped keys', function()
+    it('argument 1 is keys after mapping, argument 2 is typed keys', function()
       exec_lua [[
         keys = {}
+        typed = {}
 
         vim.cmd("inoremap hello world")
 
-        vim.on_key(function(buf)
+        vim.on_key(function(buf, typed_buf)
           if buf:byte() == 27 then
             buf = "<ESC>"
           end
+          if typed_buf:byte() == 27 then
+            typed_buf = "<ESC>"
+          end
 
           table.insert(keys, buf)
+          table.insert(typed, typed_buf)
         end)
       ]]
       insert('hello')
 
       eq('iworld<ESC>', exec_lua [[return table.concat(keys, '')]])
+      eq('ihello<ESC>', exec_lua [[return table.concat(typed, '')]])
     end)
 
     it('can call vim.fn functions on Ctrl-C #17273', function()