2 files changed, 121 insertions, 46 deletions
diff --git a/runtime/lua/vim/termcap.lua b/runtime/lua/vim/termcap.lua
new file mode 100644
index 0000000000..0eefc5eee4
--- /dev/null
+++ b/runtime/lua/vim/termcap.lua
@@ -0,0 +1,60 @@
+local M = {}
+
+--- Query the host terminal emulator for terminfo capabilities.
+---
+--- This function sends the XTGETTCAP DCS sequence to the host terminal emulator asking the terminal
+--- to send us its terminal capabilities. These are strings that are normally taken from a terminfo
+--- file, however an up to date terminfo database is not always available (particularly on remote
+--- machines), and many terminals continue to misidentify themselves or do not provide their own
+--- terminfo file, making the terminfo database unreliable.
+---
+--- Querying the terminal guarantees that we get a truthful answer, but only if the host terminal
+--- emulator supports the XTGETTCAP sequence.
+---
+--- @param caps string|table A terminal capability or list of capabilities to query
+--- @param cb function(cap:string, seq:string) Function to call when a response is received
+function M.query(caps, cb)
+  vim.validate({
+    caps = { caps, { 'string', 'table' } },
+    cb = { cb, 'f' },
+  })
+
+  if type(caps) ~= 'table' then
+    caps = { caps }
+  end
+
+  local count = #caps
+
+  vim.api.nvim_create_autocmd('TermResponse', {
+    callback = function(args)
+      local resp = args.data ---@type string
+      local k, v = resp:match('^\027P1%+r(%x+)=(%x+)$')
+      if k and v then
+        local cap = vim.text.hexdecode(k)
+        local seq =
+          vim.text.hexdecode(v):gsub('\\E', '\027'):gsub('%%p%d', ''):gsub('\\(%d+)', string.char)
+
+        -- TODO: When libtermkey is patched to accept BEL as an OSC terminator, this workaround can
+        -- be removed
+        seq = seq:gsub('\007$', '\027\\')
+
+        cb(cap, seq)
+
+        count = count - 1
+        if count == 0 then
+          return true
+        end
+      end
+    end,
+  })
+
+  local encoded = {} ---@type string[]
+  for i = 1, #caps do
+    encoded[i] = vim.text.hexencode(caps[i])
+  end
+
+  local query = string.format('\027P+q%s\027\\', table.concat(encoded, ';'))
+  io.stdout:write(query)
+end
+
+return M
diff --git a/runtime/lua/vim/ui/clipboard/osc52.lua b/runtime/lua/vim/ui/clipboard/osc52.lua
index 035a6abb86..f1d454010f 100644
--- a/runtime/lua/vim/ui/clipboard/osc52.lua
+++ b/runtime/lua/vim/ui/clipboard/osc52.lua
@@ -1,60 +1,75 @@
 local M = {}
 
-function M.copy(lines)
-  local s = table.concat(lines, '\n')
-  io.stdout:write(string.format('\027]52;;%s\027\\', vim.base64.encode(s)))
+--- Return the OSC 52 escape sequence
+---
+--- @param clipboard string The clipboard to read from or write to
+--- @param contents string The Base64 encoded contents to write to the clipboard, or '?' to read
+---                        from the clipboard
+local function osc52(clipboard, contents)
+  return string.format('\027]52;%s;%s\027\\', clipboard, contents)
 end
 
-function M.paste()
-  local contents = nil
-  local id = vim.api.nvim_create_autocmd('TermResponse', {
-    callback = function(args)
-      local resp = args.data ---@type string
-      local encoded = resp:match('\027%]52;%w?;([A-Za-z0-9+/=]*)')
-      if encoded then
-        contents = vim.base64.decode(encoded)
-        return true
-      end
-    end,
-  })
-
-  io.stdout:write('\027]52;;?\027\\')
-
-  local ok, res
-
-  -- Wait 1s first for terminals that respond quickly
-  ok, res = vim.wait(1000, function()
-    return contents ~= nil
-  end)
-
-  if res == -1 then
-    -- If no response was received after 1s, print a message and keep waiting
-    vim.api.nvim_echo(
-      { { 'Waiting for OSC 52 response from the terminal. Press Ctrl-C to interrupt...' } },
-      false,
-      {}
-    )
-    ok, res = vim.wait(9000, function()
+function M.copy(reg)
+  local clipboard = reg == '+' and 'c' or 's'
+  return function(lines)
+    local s = table.concat(lines, '\n')
+    io.stdout:write(osc52(clipboard, vim.base64.encode(s)))
+  end
+end
+
+function M.paste(reg)
+  local clipboard = reg == '+' and 'c' or 's'
+  return function()
+    local contents = nil
+    local id = vim.api.nvim_create_autocmd('TermResponse', {
+      callback = function(args)
+        local resp = args.data ---@type string
+        local encoded = resp:match('\027%]52;%w?;([A-Za-z0-9+/=]*)')
+        if encoded then
+          contents = vim.base64.decode(encoded)
+          return true
+        end
+      end,
+    })
+
+    io.stdout:write(osc52(clipboard, '?'))
+
+    local ok, res
+
+    -- Wait 1s first for terminals that respond quickly
+    ok, res = vim.wait(1000, function()
       return contents ~= nil
     end)
-  end
 
-  if not ok then
-    vim.api.nvim_del_autocmd(id)
     if res == -1 then
-      vim.notify(
-        'Timed out waiting for a clipboard response from the terminal',
-        vim.log.levels.WARN
+      -- If no response was received after 1s, print a message and keep waiting
+      vim.api.nvim_echo(
+        { { 'Waiting for OSC 52 response from the terminal. Press Ctrl-C to interrupt...' } },
+        false,
+        {}
       )
-    elseif res == -2 then
-      -- Clear message area
-      vim.api.nvim_echo({ { '' } }, false, {})
+      ok, res = vim.wait(9000, function()
+        return contents ~= nil
+      end)
+    end
+
+    if not ok then
+      vim.api.nvim_del_autocmd(id)
+      if res == -1 then
+        vim.notify(
+          'Timed out waiting for a clipboard response from the terminal',
+          vim.log.levels.WARN
+        )
+      elseif res == -2 then
+        -- Clear message area
+        vim.api.nvim_echo({ { '' } }, false, {})
+      end
+      return 0
     end
-    return 0
-  end
 
-  -- If we get here, contents should be non-nil
-  return vim.split(assert(contents), '\n')
+    -- If we get here, contents should be non-nil
+    return vim.split(assert(contents), '\n')
+  end
 end
 
 return M