feat(terminal): add support for kitty keyboard protocol

This commit adds basic support for the kitty keyboard protocol to
Neovim's builtin terminal. For now only the first mode ("Disambiguate
escape codes") is supported.

8 files changed, 301 insertions, 47 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index f897220374..33ffeae2bb 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -341,6 +341,9 @@ TERMINAL
 • |jobstart()| gained the "term" flag.
 • The |terminal| will send theme update notifications when 'background' is
   changed and DEC mode 2031 is enabled.
+• The |terminal| has experimental support for the Kitty keyboard protocol
+  (sometimes called "CSI u" key encoding). Only the "Disambiguate escape
+  codes" mode is currently supported.
 
 TREESITTER
 
diff --git a/src/nvim/getchar.c b/src/nvim/getchar.c
index 6cf4556a9f..6ec84ff543 100644
--- a/src/nvim/getchar.c
+++ b/src/nvim/getchar.c
@@ -1517,12 +1517,10 @@ int merge_modifiers(int c_arg, int *modifiers)
   int c = c_arg;
 
   if (*modifiers & MOD_MASK_CTRL) {
-    if ((c >= '`' && c <= 0x7f) || (c >= '@' && c <= '_')) {
-      if (!(State & MODE_TERMINAL) || !(c == 'I' || c == 'J' || c == 'M' || c == '[')) {
-        c &= 0x1f;
-        if (c == NUL) {
-          c = K_ZERO;
-        }
+    if (c >= '@' && c <= 0x7f) {
+      c &= 0x1f;
+      if (c == NUL) {
+        c = K_ZERO;
       }
     } else if (c == '6') {
       // CTRL-6 is equivalent to CTRL-^
@@ -2058,6 +2056,12 @@ static bool at_ins_compl_key(void)
 /// @return  the length of the replaced bytes, 0 if nothing changed, -1 for error.
 static int check_simplify_modifier(int max_offset)
 {
+  // We want full modifiers in Terminal mode so that the key can be correctly
+  // encoded
+  if (State & MODE_TERMINAL) {
+    return 0;
+  }
+
   for (int offset = 0; offset < max_offset; offset++) {
     if (offset + 3 >= typebuf.tb_len) {
       break;
diff --git a/src/nvim/terminal.c b/src/nvim/terminal.c
index 2ad5ac49ca..197a225209 100644
--- a/src/nvim/terminal.c
+++ b/src/nvim/terminal.c
@@ -783,7 +783,12 @@ static int terminal_execute(VimState *state, int key)
 {
   TerminalState *s = (TerminalState *)state;
 
-  switch (key) {
+  // Check for certain control keys like Ctrl-C and Ctrl-\. We still send the
+  // unmerged key and modifiers to the terminal.
+  int tmp_mod_mask = mod_mask;
+  int mod_key = merge_modifiers(key, &tmp_mod_mask);
+
+  switch (mod_key) {
   case K_LEFTMOUSE:
   case K_LEFTDRAG:
   case K_LEFTRELEASE:
@@ -841,13 +846,13 @@ static int terminal_execute(VimState *state, int key)
     FALLTHROUGH;
 
   default:
-    if (key == Ctrl_C) {
+    if (mod_key == Ctrl_C) {
       // terminal_enter() always sets `mapped_ctrl_c` to avoid `got_int`. 8eeda7169aa4
       // But `got_int` may be set elsewhere, e.g. by interrupt() or an autocommand,
       // so ensure that it is cleared.
       got_int = false;
     }
-    if (key == Ctrl_BSL && !s->got_bsl) {
+    if (mod_key == Ctrl_BSL && !s->got_bsl) {
       s->got_bsl = true;
       break;
     }
@@ -1016,7 +1021,7 @@ static void terminal_send_key(Terminal *term, int c)
 
   VTermKey key = convert_key(&c, &mod);
 
-  if (key) {
+  if (key != VTERM_KEY_NONE) {
     vterm_keyboard_key(term->vt, key, mod);
   } else if (!IS_SPECIAL(c)) {
     vterm_keyboard_unichar(term->vt, (uint32_t)c, mod);
diff --git a/src/nvim/vterm/keyboard.c b/src/nvim/vterm/keyboard.c
index 696b09157e..dd088ac40e 100644
--- a/src/nvim/vterm/keyboard.c
+++ b/src/nvim/vterm/keyboard.c
@@ -10,54 +10,77 @@
 # include "vterm/keyboard.c.generated.h"
 #endif
 
+static VTermKeyEncodingFlags vterm_state_get_key_encoding_flags(const VTermState *state)
+{
+  int screen = state->mode.alt_screen ? BUFIDX_ALTSCREEN : BUFIDX_PRIMARY;
+  const struct VTermKeyEncodingStack *stack = &state->key_encoding_stacks[screen];
+  assert(stack->size > 0);
+  return stack->items[stack->size - 1];
+}
+
 void vterm_keyboard_unichar(VTerm *vt, uint32_t c, VTermModifier mod)
 {
-  // The shift modifier is never important for Unicode characters apart from Space
-  if (c != ' ') {
-    mod &= (unsigned)~VTERM_MOD_SHIFT;
+  bool passthru = false;
+  if (c == ' ') {
+    // Space is passed through only when there are no modifiers (including shift)
+    passthru = mod == VTERM_MOD_NONE;
+  } else {
+    // Otherwise pass through when there are no modifiers (ignoring shift)
+    passthru = (mod & (unsigned)~VTERM_MOD_SHIFT) == 0;
   }
 
-  if (mod == 0) {
-    // Normal text - ignore just shift
+  if (passthru) {
     char str[6];
     int seqlen = fill_utf8((int)c, str);
     vterm_push_output_bytes(vt, str, (size_t)seqlen);
     return;
   }
 
-  int needs_CSIu;
-  switch (c) {
-  // Special Ctrl- letters that can't be represented elsewise
-  case 'i':
-  case 'j':
-  case 'm':
-  case '[':
-    needs_CSIu = 1;
-    break;
-  // Ctrl-\ ] ^ _ don't need CSUu
-  case '\\':
-  case ']':
-  case '^':
-  case '_':
-    needs_CSIu = 0;
-    break;
-  // Shift-space needs CSIu
-  case ' ':
-    needs_CSIu = !!(mod & VTERM_MOD_SHIFT);
-    break;
-  // All other characters needs CSIu except for letters a-z
-  default:
-    needs_CSIu = (c < 'a' || c > 'z');
-  }
+  VTermKeyEncodingFlags flags = vterm_state_get_key_encoding_flags(vt->state);
+  if (flags.disambiguate) {
+    // Always use unshifted codepoint
+    if (c >= 'A' && c <= 'Z') {
+      c += 'a' - 'A';
+      mod |= VTERM_MOD_SHIFT;
+    }
 
-  // ALT we can just prefix with ESC; anything else requires CSI u
-  if (needs_CSIu && (mod & (unsigned)~VTERM_MOD_ALT)) {
     vterm_push_output_sprintf_ctrl(vt, C1_CSI, "%d;%du", c, mod + 1);
     return;
   }
 
   if (mod & VTERM_MOD_CTRL) {
-    c &= 0x1f;
+    // Handle special cases. These are taken from kitty, but seem mostly
+    // consistent across terminals.
+    switch (c) {
+    case '2':
+    case ' ':
+      // Ctrl+2 is NUL to match Ctrl+@ (which is Shift+2 on US keyboards)
+      // Ctrl+Space is also NUL for some reason
+      c = 0x00;
+      break;
+    case '3':
+    case '4':
+    case '5':
+    case '6':
+    case '7':
+      // Ctrl+3 through Ctrl+7 are sequential starting from 0x1b. Importantly,
+      // this means that Ctrl+6 emits 0x1e (the same as Ctrl+^ on US keyboards)
+      c = 0x1b + c - '3';
+      break;
+    case '8':
+      // Ctrl+8 is DEL
+      c = 0x7f;
+      break;
+    case '/':
+      // Ctrl+/ is equivalent to Ctrl+_ for historic reasons
+      c = 0x1f;
+      break;
+    default:
+      if (c >= '@' && c <= 0x7f) {
+        c &= 0x1f;
+      }
+      break;
+    }
   }
 
   vterm_push_output_sprintf(vt, "%s%c", mod & VTERM_MOD_ALT ? ESC_S : "", c);
@@ -75,7 +98,7 @@ typedef struct {
     KEYCODE_CSINUM,
     KEYCODE_KEYPAD,
   } type;
-  char literal;
+  int literal;
   int csinum;
 } keycodes_s;
 
@@ -137,12 +160,35 @@ static keycodes_s keycodes_kp[] = {
   { KEYCODE_KEYPAD, '=',  'X' },  // KP_EQUAL
 };
 
+static keycodes_s keycodes_kp_csiu[] = {
+  { KEYCODE_KEYPAD, 57399, 'p' },  // KP_0
+  { KEYCODE_KEYPAD, 57400, 'q' },  // KP_1
+  { KEYCODE_KEYPAD, 57401, 'r' },  // KP_2
+  { KEYCODE_KEYPAD, 57402, 's' },  // KP_3
+  { KEYCODE_KEYPAD, 57403, 't' },  // KP_4
+  { KEYCODE_KEYPAD, 57404, 'u' },  // KP_5
+  { KEYCODE_KEYPAD, 57405, 'v' },  // KP_6
+  { KEYCODE_KEYPAD, 57406, 'w' },  // KP_7
+  { KEYCODE_KEYPAD, 57407, 'x' },  // KP_8
+  { KEYCODE_KEYPAD, 57408, 'y' },  // KP_9
+  { KEYCODE_KEYPAD, 57411, 'j' },  // KP_MULT
+  { KEYCODE_KEYPAD, 57413, 'k' },  // KP_PLUS
+  { KEYCODE_KEYPAD, 57416, 'l' },  // KP_COMMA
+  { KEYCODE_KEYPAD, 57412, 'm' },  // KP_MINUS
+  { KEYCODE_KEYPAD, 57409, 'n' },  // KP_PERIOD
+  { KEYCODE_KEYPAD, 57410, 'o' },  // KP_DIVIDE
+  { KEYCODE_KEYPAD, 57414, 'M' },  // KP_ENTER
+  { KEYCODE_KEYPAD, 57415, 'X' },  // KP_EQUAL
+};
+
 void vterm_keyboard_key(VTerm *vt, VTermKey key, VTermModifier mod)
 {
   if (key == VTERM_KEY_NONE) {
     return;
   }
 
+  VTermKeyEncodingFlags flags = vterm_state_get_key_encoding_flags(vt->state);
+
   keycodes_s k;
   if (key < VTERM_KEY_FUNCTION_0) {
     if (key >= sizeof(keycodes)/sizeof(keycodes[0])) {
@@ -158,7 +204,12 @@ void vterm_keyboard_key(VTerm *vt, VTermKey key, VTermModifier mod)
     if ((key - VTERM_KEY_KP_0) >= sizeof(keycodes_kp)/sizeof(keycodes_kp[0])) {
       return;
     }
-    k = keycodes_kp[key - VTERM_KEY_KP_0];
+
+    if (flags.disambiguate) {
+      k = keycodes_kp_csiu[key - VTERM_KEY_KP_0];
+    } else {
+      k = keycodes_kp[key - VTERM_KEY_KP_0];
+    }
   }
 
   switch (k.type) {
@@ -167,7 +218,9 @@ void vterm_keyboard_key(VTerm *vt, VTermKey key, VTermModifier mod)
 
   case KEYCODE_TAB:
     // Shift-Tab is CSI Z but plain Tab is 0x09
-    if (mod == VTERM_MOD_SHIFT) {
+    if (flags.disambiguate) {
+      goto case_LITERAL;
+    } else if (mod == VTERM_MOD_SHIFT) {
       vterm_push_output_sprintf_ctrl(vt, C1_CSI, "Z");
     } else if (mod & VTERM_MOD_SHIFT) {
       vterm_push_output_sprintf_ctrl(vt, C1_CSI, "1;%dZ", mod + 1);
@@ -187,7 +240,20 @@ void vterm_keyboard_key(VTerm *vt, VTermKey key, VTermModifier mod)
 
   case KEYCODE_LITERAL:
                         case_LITERAL:
-    if (mod & (VTERM_MOD_SHIFT|VTERM_MOD_CTRL)) {
+    if (flags.disambiguate) {
+      switch (key) {
+      case VTERM_KEY_TAB:
+      case VTERM_KEY_ENTER:
+      case VTERM_KEY_BACKSPACE:
+        // If there are no mods then leave these as-is
+        flags.disambiguate = mod != VTERM_MOD_NONE;
+        break;
+      default:
+        break;
+      }
+    }
+
+    if (flags.disambiguate) {
       vterm_push_output_sprintf_ctrl(vt, C1_CSI, "%d;%du", k.literal, mod + 1);
     } else {
       vterm_push_output_sprintf(vt, mod & VTERM_MOD_ALT ? ESC_S "%c" : "%c", k.literal);
@@ -229,7 +295,7 @@ void vterm_keyboard_key(VTerm *vt, VTermKey key, VTermModifier mod)
 
   case KEYCODE_KEYPAD:
     if (vt->state->mode.keypad) {
-      k.literal = (char)k.csinum;
+      k.literal = k.csinum;
       goto case_SS3;
     } else {
       goto case_LITERAL;
diff --git a/src/nvim/vterm/state.c b/src/nvim/vterm/state.c
index 4ad07377de..0e43107347 100644
--- a/src/nvim/vterm/state.c
+++ b/src/nvim/vterm/state.c
@@ -1,3 +1,4 @@
+#include <assert.h>
 #include <stdio.h>
 #include <string.h>
 
@@ -116,6 +117,15 @@ static VTermState *vterm_state_new(VTerm *vt)
     (*state->encoding_utf8.enc->init)(state->encoding_utf8.enc, state->encoding_utf8.data);
   }
 
+  for (size_t i = 0; i < ARRAY_SIZE(state->key_encoding_stacks); i++) {
+    struct VTermKeyEncodingStack *stack = &state->key_encoding_stacks[i];
+    for (size_t j = 0; j < ARRAY_SIZE(stack->items); j++) {
+      memset(&stack->items[j], 0, sizeof(stack->items[j]));
+    }
+
+    stack->size = 1;
+  }
+
   return state;
 }
 
@@ -916,6 +926,115 @@ static void request_version_string(VTermState *state)
                                 VTERM_VERSION_MAJOR, VTERM_VERSION_MINOR);
 }
 
+static void request_key_encoding_flags(VTermState *state)
+{
+  int screen = state->mode.alt_screen ? BUFIDX_ALTSCREEN : BUFIDX_PRIMARY;
+  struct VTermKeyEncodingStack *stack = &state->key_encoding_stacks[screen];
+
+  int reply = 0;
+
+  assert(stack->size > 0);
+  VTermKeyEncodingFlags flags = stack->items[stack->size - 1];
+
+  if (flags.disambiguate) {
+    reply |= KEY_ENCODING_DISAMBIGUATE;
+  }
+
+  if (flags.report_events) {
+    reply |= KEY_ENCODING_REPORT_EVENTS;
+  }
+
+  if (flags.report_alternate) {
+    reply |= KEY_ENCODING_REPORT_ALTERNATE;
+  }
+
+  if (flags.report_all_keys) {
+    reply |= KEY_ENCODING_REPORT_ALL_KEYS;
+  }
+
+  if (flags.report_associated) {
+    reply |= KEY_ENCODING_REPORT_ASSOCIATED;
+  }
+
+  vterm_push_output_sprintf_ctrl(state->vt, C1_CSI, "?%du", reply);
+}
+
+static void set_key_encoding_flags(VTermState *state, int arg, int mode)
+{
+  // When mode is 3, bits set in arg reset the corresponding mode
+  bool set = mode != 3;
+
+  // When mode is 1, unset bits are reset
+  bool reset_unset = mode == 1;
+
+  struct VTermKeyEncodingFlags flags = { 0 };
+  if (arg & KEY_ENCODING_DISAMBIGUATE) {
+    flags.disambiguate = set;
+  } else if (reset_unset) {
+    flags.disambiguate = false;
+  }
+
+  if (arg & KEY_ENCODING_REPORT_EVENTS) {
+    flags.report_events = set;
+  } else if (reset_unset) {
+    flags.report_events = false;
+  }
+
+  if (arg & KEY_ENCODING_REPORT_ALTERNATE) {
+    flags.report_alternate = set;
+  } else if (reset_unset) {
+    flags.report_alternate = false;
+  }
+  if (arg & KEY_ENCODING_REPORT_ALL_KEYS) {
+    flags.report_all_keys = set;
+  } else if (reset_unset) {
+    flags.report_all_keys = false;
+  }
+
+  if (arg & KEY_ENCODING_REPORT_ASSOCIATED) {
+    flags.report_associated = set;
+  } else if (reset_unset) {
+    flags.report_associated = false;
+  }
+
+  int screen = state->mode.alt_screen ? BUFIDX_ALTSCREEN : BUFIDX_PRIMARY;
+  struct VTermKeyEncodingStack *stack = &state->key_encoding_stacks[screen];
+  assert(stack->size > 0);
+  stack->items[stack->size - 1] = flags;
+}
+
+static void push_key_encoding_flags(VTermState *state, int arg)
+{
+  int screen = state->mode.alt_screen ? BUFIDX_ALTSCREEN : BUFIDX_PRIMARY;
+  struct VTermKeyEncodingStack *stack = &state->key_encoding_stacks[screen];
+  assert(stack->size <= ARRAY_SIZE(stack->items));
+
+  if (stack->size == ARRAY_SIZE(stack->items)) {
+    // Evict oldest entry when stack is full
+    for (size_t i = 0; i < ARRAY_SIZE(stack->items) - 1; i++) {
+      stack->items[i] = stack->items[i + 1];
+    }
+  } else {
+    stack->size++;
+  }
+
+  set_key_encoding_flags(state, arg, 1);
+}
+
+static void pop_key_encoding_flags(VTermState *state, int arg)
+{
+  int screen = state->mode.alt_screen ? BUFIDX_ALTSCREEN : BUFIDX_PRIMARY;
+  struct VTermKeyEncodingStack *stack = &state->key_encoding_stacks[screen];
+  if (arg >= stack->size) {
+    stack->size = 1;
+
+    // If a pop request is received that empties the stack, all flags are reset.
+    memset(&stack->items[0], 0, sizeof(stack->items[0]));
+  } else if (arg > 0) {
+    stack->size -= arg;
+  }
+}
+
 static int on_csi(const char *leader, const long args[], int argcount, const char *intermed,
                   char command, void *user)
 {
@@ -932,6 +1051,8 @@ static int on_csi(const char *leader, const long args[], int argcount, const cha
     switch (leader[0]) {
     case '?':
     case '>':
+    case '<':
+    case '=':
       leader_byte = (int)leader[0];
       break;
     default:
@@ -1542,6 +1663,23 @@ static int on_csi(const char *leader, const long args[], int argcount, const cha
 
     break;
 
+  case LEADER('?', 0x75):  // Kitty query
+    request_key_encoding_flags(state);
+    break;
+
+  case LEADER('>', 0x75):  // Kitty push flags
+    push_key_encoding_flags(state, CSI_ARG_OR(args[0], 0));
+    break;
+
+  case LEADER('<', 0x75):  // Kitty pop flags
+    pop_key_encoding_flags(state, CSI_ARG_OR(args[0], 1));
+    break;
+
+  case LEADER('=', 0x75):  // Kitty set flags
+    val = argcount < 2 || CSI_ARG_IS_MISSING(args[1]) ? 1 : CSI_ARG(args[1]);
+    set_key_encoding_flags(state, CSI_ARG_OR(args[0], 0), val);
+    break;
+
   case INTERMED('\'', 0x7D):  // DECIC
     count = CSI_ARG_COUNT(args[0]);
 
diff --git a/src/nvim/vterm/vterm_internal_defs.h b/src/nvim/vterm/vterm_internal_defs.h
index d4d59867bf..19e809490f 100644
--- a/src/nvim/vterm/vterm_internal_defs.h
+++ b/src/nvim/vterm/vterm_internal_defs.h
@@ -21,7 +21,14 @@
 #define BUFIDX_PRIMARY   0
 #define BUFIDX_ALTSCREEN 1
 
+#define KEY_ENCODING_DISAMBIGUATE 0x1
+#define KEY_ENCODING_REPORT_EVENTS 0x2
+#define KEY_ENCODING_REPORT_ALTERNATE 0x4
+#define KEY_ENCODING_REPORT_ALL_KEYS 0x8
+#define KEY_ENCODING_REPORT_ASSOCIATED 0x10
+
 typedef struct VTermEncoding VTermEncoding;
+typedef struct VTermKeyEncodingFlags VTermKeyEncodingFlags;
 
 typedef struct {
   VTermEncoding *enc;
@@ -46,6 +53,21 @@ struct VTermPen {
   unsigned baseline:2;
 };
 
+// https://sw.kovidgoyal.net/kitty/keyboard-protocol/#progressive-enhancement
+struct VTermKeyEncodingFlags {
+  bool disambiguate:1;
+  bool report_events:1;
+  bool report_alternate:1;
+  bool report_all_keys:1;
+  bool report_associated:1;
+};
+
+struct VTermKeyEncodingStack {
+  VTermKeyEncodingFlags items[16];
+  uint8_t size;  ///< Number of items in the stack. This is at least 1 and at
+                 ///< most the length of the "items" array.
+};
+
 struct VTermState {
   VTerm *vt;
 
@@ -171,6 +193,9 @@ struct VTermState {
     char *buffer;
     size_t buflen;
   } selection;
+
+  // Maintain two stacks, one for primary screen and one for altscreen
+  struct VTermKeyEncodingStack key_encoding_stacks[2];
 };
 
 struct VTerm {
diff --git a/test/functional/terminal/buffer_spec.lua b/test/functional/terminal/buffer_spec.lua
index 66b75a4ea2..a524e49ef4 100644
--- a/test/functional/terminal/buffer_spec.lua
+++ b/test/functional/terminal/buffer_spec.lua
@@ -629,6 +629,14 @@ describe('terminal input', function()
   -- TODO(bfredl): getcharstr() erases the distinction between <C-I> and <Tab>.
   -- If it was enhanced or replaced this could get folded into the test above.
   it('can send TAB/C-I and ESC/C-[ separately', function()
+    if
+      skip(
+        is_os('win'),
+        "The escape sequence to enable kitty keyboard mode doesn't work on Windows"
+      )
+    then
+      return
+    end
     clear()
     local screen = tt.setup_child_nvim({
       '-u',
diff --git a/test/unit/vterm_spec.lua b/test/unit/vterm_spec.lua
index 6ff3c18d2a..9f70187fad 100644
--- a/test/unit/vterm_spec.lua
+++ b/test/unit/vterm_spec.lua
@@ -2324,6 +2324,9 @@ putglyph 1f3f4,200d,2620,fe0f 2 0,4]])
     local vt = init()
     local state = wantstate(vt)
 
+    -- Disambiguate escape codes
+    push('\x1b[>1u', vt)
+
     -- Unmodified ASCII
     inchar(41, vt)
     expect('output 29')
@@ -2478,6 +2481,8 @@ putglyph 1f3f4,200d,2620,fe0f 2 0,4]])
     expect_output('\x1b[I')
     vterm.vterm_state_focus_out(state)
     expect_output('\x1b[O')
+
+    push('\x1b[<u', vt)
   end)
 
   itp('26state_query', function()