refactor(tui): refactor Kitty keyboard query implementation

Refactor our implementation of querying for Kitty keyboard protocol
support:

- Remove usage of the "extkeys" term. This is not standard or really
  used elsewhere. Use "key encoding" instead
- Replace usages of "CSIu" with "Kitty". "Kitty keyboard protocol" is
  vastly more common than "CSIu" now
- Replace the countdown response counter with a simple boolean flag. We
  don't actually need a countdown counter because we request the primary
  device attributes along with the Kitty keyboard query, so we will
  always receive a "terminating event", making a countdown/timer
  unnecessary
- Move the CSI response handling into a dedicated function
- Bypass Unibilium for sending key encoding escape sequences. These
  sequences are not part of terminfo and do not have any parameters, so
  there's no reason to go through Unibilium

3 files changed, 114 insertions, 75 deletions

diff --git a/src/nvim/tui/input.c b/src/nvim/tui/input.c
index de4afa68cf..03f22745bf 100644
--- a/src/nvim/tui/input.c
+++ b/src/nvim/tui/input.c
@@ -122,7 +122,7 @@ void tinput_init(TermInput *input, Loop *loop)
   input->loop = loop;
   input->paste = 0;
   input->in_fd = STDIN_FILENO;
-  input->extkeys_type = kExtkeysNone;
+  input->key_encoding = kKeyEncodingLegacy;
   input->ttimeout = (bool)p_ttimeout;
   input->ttimeoutlen = p_ttm;
   input->key_buffer = rbuffer_new(KEY_BUFFER_SIZE);
@@ -444,38 +444,7 @@ static void tk_getkeys(TermInput *input, bool force)
     } else if (key.type == TERMKEY_TYPE_MOUSE) {
       forward_mouse_event(input, &key);
     } else if (key.type == TERMKEY_TYPE_UNKNOWN_CSI) {
-      // There is no specified limit on the number of parameters a CSI sequence can contain, so just
-      // allocate enough space for a large upper bound
-      long args[16];
-      size_t nargs = 16;
-      unsigned long cmd;
-      if (termkey_interpret_csi(input->tk, &key, args, &nargs, &cmd) == TERMKEY_RES_KEY) {
-        uint8_t intermediate = (cmd >> 16) & 0xFF;
-        uint8_t initial = (cmd >> 8) & 0xFF;
-        uint8_t command = cmd & 0xFF;
-
-        // Currently unused
-        (void)intermediate;
-
-        if (input->waiting_for_csiu_response > 0) {
-          if (initial == '?' && command == 'u') {
-            // The first (and only) argument contains the current progressive
-            // enhancement flags. Only enable CSI u mode if the first bit
-            // (disambiguate escape codes) is not already set
-            if (nargs > 0 && (args[0] & 0x1) == 0) {
-              input->extkeys_type = kExtkeysCSIu;
-            } else {
-              input->extkeys_type = kExtkeysNone;
-            }
-          } else if (initial == '?' && command == 'c') {
-            // Received Primary Device Attributes response
-            input->waiting_for_csiu_response = 0;
-            tui_enable_extkeys(input->tui_data);
-          } else {
-            input->waiting_for_csiu_response--;
-          }
-        }
-      }
+      handle_unknown_csi(input, &key);
     } else if (key.type == TERMKEY_TYPE_OSC || key.type == TERMKEY_TYPE_DCS) {
       handle_term_response(input, &key);
     } else if (key.type == TERMKEY_TYPE_MODEREPORT) {
@@ -578,6 +547,7 @@ static HandleState handle_bracketed_paste(TermInput *input)
   return kNotApplicable;
 }
 
+/// Handle an OSC or DCS response sequence from the terminal.
 static void handle_term_response(TermInput *input, const TermKeyKey *key)
   FUNC_ATTR_NONNULL_ALL
 {
@@ -612,6 +582,7 @@ static void handle_term_response(TermInput *input, const TermKeyKey *key)
   }
 }
 
+/// Handle a mode report (DECRPM) sequence from the terminal.
 static void handle_modereport(TermInput *input, const TermKeyKey *key)
   FUNC_ATTR_NONNULL_ALL
 {
@@ -622,6 +593,59 @@ static void handle_modereport(TermInput *input, const TermKeyKey *key)
   tui_dec_report_mode(input->tui_data, (TerminalDecMode)mode, value);
 }
 
+/// Handle a CSI sequence from the terminal that is unrecognized by libtermkey.
+static void handle_unknown_csi(TermInput *input, const TermKeyKey *key)
+  FUNC_ATTR_NONNULL_ALL
+{
+  // There is no specified limit on the number of parameters a CSI sequence can
+  // contain, so just allocate enough space for a large upper bound
+  long args[16];
+  size_t nargs = 16;
+  unsigned long cmd;
+  if (termkey_interpret_csi(input->tk, key, args, &nargs, &cmd) != TERMKEY_RES_KEY) {
+    return;
+  }
+
+  uint8_t intermediate = (cmd >> 16) & 0xFF;
+  uint8_t initial = (cmd >> 8) & 0xFF;
+  uint8_t command = cmd & 0xFF;
+
+  // Currently unused
+  (void)intermediate;
+
+  switch (command) {
+  case 'u':
+    switch (initial) {
+    case '?':
+      // Kitty keyboard protocol query response.
+      if (input->waiting_for_kkp_response) {
+        input->waiting_for_kkp_response = false;
+        input->key_encoding = kKeyEncodingKitty;
+        tui_set_key_encoding(input->tui_data);
+      }
+
+      break;
+    }
+    break;
+  case 'c':
+    switch (initial) {
+    case '?':
+      // Primary Device Attributes response
+      if (input->waiting_for_kkp_response) {
+        input->waiting_for_kkp_response = false;
+
+        // Enable the fallback key encoding (if any)
+        tui_set_key_encoding(input->tui_data);
+      }
+
+      break;
+    }
+    break;
+  default:
+    break;
+  }
+}
+
 static void handle_raw_buffer(TermInput *input, bool force)
 {
   HandleState is_paste = kNotApplicable;
diff --git a/src/nvim/tui/input.h b/src/nvim/tui/input.h
index 01514269be..2d72d1978c 100644
--- a/src/nvim/tui/input.h
+++ b/src/nvim/tui/input.h
@@ -14,18 +14,22 @@
 #include "nvim/types.h"
 
 typedef enum {
-  kExtkeysNone,
-  kExtkeysCSIu,
-  kExtkeysXterm,
-} ExtkeysType;
+  kKeyEncodingLegacy,  ///< Legacy key encoding
+  kKeyEncodingKitty,   ///< Kitty keyboard protocol encoding
+  kKeyEncodingXterm,   ///< Xterm's modifyOtherKeys encoding (XTMODKEYS)
+} KeyEncoding;
 
-typedef struct term_input {
+typedef struct {
   int in_fd;
   // Phases: -1=all 0=disabled 1=first-chunk 2=continue 3=last-chunk
   int8_t paste;
   bool ttimeout;
-  int8_t waiting_for_csiu_response;
-  ExtkeysType extkeys_type;
+
+  bool waiting_for_kkp_response;  ///< True if we are expecting to receive a response to a query for
+                                  ///< Kitty keyboard protocol support
+
+  KeyEncoding key_encoding;       ///< The key encoding used by the terminal emulator
+
   OptInt ttimeoutlen;
   TermKey *tk;
   TermKey_Terminfo_Getstr_Hook *tk_ti_hook_fn;  ///< libtermkey terminfo hook
diff --git a/src/nvim/tui/tui.c b/src/nvim/tui/tui.c
index a64f9af454..cdc8f769a1 100644
--- a/src/nvim/tui/tui.c
+++ b/src/nvim/tui/tui.c
@@ -134,8 +134,6 @@ struct TUIData {
     int save_title, restore_title;
     int set_underline_style;
     int set_underline_color;
-    int enable_extended_keys, disable_extended_keys;
-    int get_extkeys;
     int sync;
   } unibi_ext;
   char *space_buf;
@@ -183,29 +181,34 @@ void tui_start(TUIData **tui_p, int *width, int *height, char **term)
   *term = tui->term;
 }
 
-void tui_enable_extkeys(TUIData *tui)
+void tui_set_key_encoding(TUIData *tui)
+  FUNC_ATTR_NONNULL_ALL
 {
-  TermInput input = tui->input;
-  unibi_term *ut = tui->ut;
-
-  switch (input.extkeys_type) {
-  case kExtkeysCSIu:
-    tui->unibi_ext.enable_extended_keys = (int)unibi_add_ext_str(ut, "ext.enable_extended_keys",
-                                                                 "\x1b[>1u");
-    tui->unibi_ext.disable_extended_keys = (int)unibi_add_ext_str(ut, "ext.disable_extended_keys",
-                                                                  "\x1b[<1u");
+  switch (tui->input.key_encoding) {
+  case kKeyEncodingKitty:
+    out(tui, S_LEN("\x1b[>1u"));
     break;
-  case kExtkeysXterm:
-    tui->unibi_ext.enable_extended_keys = (int)unibi_add_ext_str(ut, "ext.enable_extended_keys",
-                                                                 "\x1b[>4;2m");
-    tui->unibi_ext.disable_extended_keys = (int)unibi_add_ext_str(ut, "ext.disable_extended_keys",
-                                                                  "\x1b[>4;0m");
+  case kKeyEncodingXterm:
+    out(tui, S_LEN("\x1b[>4;2m"));
     break;
-  default:
+  case kKeyEncodingLegacy:
     break;
   }
+}
 
-  unibi_out_ext(tui, tui->unibi_ext.enable_extended_keys);
+static void tui_reset_key_encoding(TUIData *tui)
+  FUNC_ATTR_NONNULL_ALL
+{
+  switch (tui->input.key_encoding) {
+  case kKeyEncodingKitty:
+    out(tui, S_LEN("\x1b[<1u"));
+    break;
+  case kKeyEncodingXterm:
+    out(tui, S_LEN("\x1b[>4;0m"));
+    break;
+  case kKeyEncodingLegacy:
+    break;
+  }
 }
 
 /// Request the terminal's DEC mode (DECRQM).
@@ -244,6 +247,21 @@ void tui_dec_report_mode(TUIData *tui, TerminalDecMode mode, TerminalModeState s
   }
 }
 
+/// Query the terminal emulator to see if it supports Kitty's keyboard protocol.
+///
+/// Write CSI ? u followed by a primary device attributes request (CSI c). If
+/// a primary device attributes response is received without first receiving an
+/// answer to the progressive enhancement query (CSI u), then the terminal does
+/// not support the Kitty keyboard protocol.
+///
+/// See https://sw.kovidgoyal.net/kitty/keyboard-protocol/#detection-of-support-for-this-protocol
+static void tui_query_kitty_keyboard(TUIData *tui)
+  FUNC_ATTR_NONNULL_ALL
+{
+  tui->input.waiting_for_kkp_response = true;
+  out(tui, S_LEN("\x1b[?u\x1b[c"));
+}
+
 static void terminfo_start(TUIData *tui)
 {
   tui->scroll_region_is_full_screen = true;
@@ -277,9 +295,6 @@ static void terminfo_start(TUIData *tui)
   tui->unibi_ext.set_cursor_style = -1;
   tui->unibi_ext.reset_cursor_style = -1;
   tui->unibi_ext.set_underline_color = -1;
-  tui->unibi_ext.enable_extended_keys = -1;
-  tui->unibi_ext.disable_extended_keys = -1;
-  tui->unibi_ext.get_extkeys = -1;
   tui->unibi_ext.sync = -1;
   tui->out_fd = STDOUT_FILENO;
   tui->out_isatty = os_isatty(tui->out_fd);
@@ -358,9 +373,8 @@ static void terminfo_start(TUIData *tui)
   // mode 2026
   tui_dec_request_mode(tui, kDecModeSynchronizedOutput);
 
-  // Query the terminal to see if it supports CSI u
-  tui->input.waiting_for_csiu_response = 5;
-  unibi_out_ext(tui, tui->unibi_ext.get_extkeys);
+  // Query the terminal to see if it supports Kitty's keyboard protocol
+  tui_query_kitty_keyboard(tui);
 
   int ret;
   uv_loop_init(&tui->write_loop);
@@ -403,8 +417,10 @@ static void terminfo_stop(TUIData *tui)
   // Reset cursor to normal before exiting alternate screen.
   unibi_out(tui, unibi_cursor_normal);
   unibi_out(tui, unibi_keypad_local);
-  // Disable extended keys before exiting alternate screen.
-  unibi_out_ext(tui, tui->unibi_ext.disable_extended_keys);
+
+  // Reset the key encoding
+  tui_reset_key_encoding(tui);
+
   // May restore old title before exiting alternate screen.
   tui_set_title(tui, (String)STRING_INIT);
   if (ui_client_exit_status == 0) {
@@ -1949,12 +1965,6 @@ static void patch_terminfo_bugs(TUIData *tui, const char *term, const char *colo
 #define XTERM_SETAB_16 \
   "\x1b[%?%p1%{8}%<%t4%p1%d%e%p1%{16}%<%t10%p1%{8}%-%d%e39%;m"
 
-  // Query the terminal to see if it supports CSI u key encoding by writing CSI
-  // ? u followed by a request for the primary device attributes (CSI c)
-  // See https://sw.kovidgoyal.net/kitty/keyboard-protocol/#detection-of-support-for-this-protocol
-  tui->unibi_ext.get_extkeys = (int)unibi_add_ext_str(ut, "ext.get_extkeys",
-                                                      "\x1b[?u\x1b[c");
-
   // Terminals with 256-colour SGR support despite what terminfo says.
   if (unibi_get_num(ut, unibi_max_colors) < 256) {
     // See http://fedoraproject.org/wiki/Features/256_Color_Terminals
@@ -2247,8 +2257,9 @@ static void augment_terminfo(TUIData *tui, const char *term, int vte_version, in
   }
 
   if (!kitty && (vte_version == 0 || vte_version >= 5400)) {
-    // Fallback to Xterm's modifyOtherKeys if terminal does not support CSI u
-    tui->input.extkeys_type = kExtkeysXterm;
+    // Fallback to Xterm's modifyOtherKeys if terminal does not support the
+    // Kitty keyboard protocol
+    tui->input.key_encoding = kKeyEncodingXterm;
   }
 }