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.

6 files changed, 60 insertions, 12 deletions

diff --git a/src/klib/kvec.h b/src/klib/kvec.h
index a32b35a14c..1b9e6fd9f8 100644
--- a/src/klib/kvec.h
+++ b/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.
diff --git a/src/nvim/getchar.c b/src/nvim/getchar.c
index 0a848c4676..665f60a49e 100644
--- a/src/nvim/getchar.c
+++ b/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]);
     }
 
+    buf[buflen] = NUL;
+
     if (reg_recording != 0) {
-      buf[buflen] = NUL;
       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.
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 9e4b698b69..cfc68dc08f 100644
--- a/src/nvim/lua/executor.c
+++ b/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"));
   }
diff --git a/src/nvim/message.c b/src/nvim/message.c
index 5a47908eb6..68a8b8e88b 100644
--- a/src/nvim/message.c
+++ b/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;
       }
 
diff --git a/src/nvim/normal.c b/src/nvim/normal.c
index bbba8069c7..c7eb5c5793 100644
--- a/src/nvim/normal.c
+++ b/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.
diff --git a/src/nvim/terminal.c b/src/nvim/terminal.c
index 9033e00f3d..2b05a8047e 100644
--- a/src/nvim/terminal.c
+++ b/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);
   }