From 68838dd44031185fa981e8cdcb9712744f4c0a9c Mon Sep 17 00:00:00 2001
From: Jannik Buhr <17450586+jmbuhr@users.noreply.github.com>
Date: Mon, 19 Dec 2022 19:03:42 +0100
Subject: [PATCH] Autocomplete (#18)

* add cmp completion source!
* autocompletion now works with one langauge at at time, but not two
  - most likely a name clash?
* update diagnostics less often by default
---
 examples/example.qmd  |   6 +-
 lua/quarto/init.lua   | 200 ++++++++++++++++++++++++++++--------------
 lua/quarto/source.lua | 163 ++++++++++++++++++++++++++++++++++
 lua/quarto/tools.lua  |   8 +-
 plugin/quarto.lua     |   6 +-
 5 files changed, 313 insertions(+), 70 deletions(-)
 create mode 100644 lua/quarto/source.lua

diff --git a/examples/example.qmd b/examples/example.qmd
index c2f436d..66101d1 100644
--- a/examples/example.qmd
+++ b/examples/example.qmd
@@ -13,6 +13,10 @@ def hello():
   print("Hello")
 ```
 
+```{python} 
+import numpy as np
+```
+
 This is how we call it:
 
 ```{python} 
@@ -22,7 +26,7 @@ hello()
 And this function is not found because we have a typo:
 
 ```{python} 
-helo()
+hello()
 ```
 
 Now, we use the function in the next code chunk
diff --git a/lua/quarto/init.lua b/lua/quarto/init.lua
index 17a4d45..3f17265 100644
--- a/lua/quarto/init.lua
+++ b/lua/quarto/init.lua
@@ -1,11 +1,13 @@
 local M = {}
 local api = vim.api
 local util = require "lspconfig.util"
+local source = require'quarto.source'
 local tools = require'quarto.tools'
 local lines = tools.lines
 local spaces = tools.spaces
+local api = vim.api
 
-local defaultConfig = {
+M.defaultConfig = {
   debug = false,
   closePreviewOnExit = true,
   lspFeatures = {
@@ -13,6 +15,10 @@ local defaultConfig = {
     languages = { 'r', 'python', 'julia' },
     diagnostics = {
       enabled = true,
+      triggers = { "BufEnter", "InsertLeave", "TextChanged" }
+    },
+    cmpSource = {
+      enabled = true,
     },
   },
   keymap = {
@@ -20,63 +26,6 @@ local defaultConfig = {
   }
 }
 
-M.config = defaultConfig
-
-
-function M.quartoPreview()
-  -- find root directory / check if it is a project
-  local buffer_path = api.nvim_buf_get_name(0)
-  local root_dir = util.root_pattern("_quarto.yml")(buffer_path)
-  local cmd
-  local mode
-  if root_dir then
-    mode = "project"
-    cmd = 'quarto preview'
-  else
-    mode = "file"
-    cmd = 'quarto preview ' .. buffer_path
-  end
-
-  local quarto_extensions = { ".qmd", ".Rmd", ".ipynb", ".md" }
-  local file_extension = buffer_path:match("^.+(%..+)$")
-  if mode == "file" and not file_extension then
-    vim.notify("Not in a file. exiting.")
-    return
-  end
-  if mode == "file" and not tools.contains(quarto_extensions, file_extension) then
-    vim.notify("Not a quarto file, ends in " .. file_extension .. " exiting.")
-    return
-  end
-
-  -- run command in embedded terminal
-  -- in a new tab and go back to the buffer
-  vim.cmd('tabedit term://' .. cmd)
-  local quartoOutputBuf = vim.api.nvim_get_current_buf()
-  vim.cmd('tabprevious')
-  api.nvim_buf_set_var(0, 'quartoOutputBuf', quartoOutputBuf)
-
-
-  -- close preview terminal on exit of the quarto buffer
-  if M.config.closePreviewOnExit then
-    api.nvim_create_autocmd({ "QuitPre", "WinClosed" }, {
-      buffer = api.nvim_get_current_buf(),
-      group = api.nvim_create_augroup("quartoPreview", {}),
-      callback = function(_, _)
-        if api.nvim_buf_is_loaded(quartoOutputBuf) then
-          api.nvim_buf_delete(quartoOutputBuf, { force = true })
-        end
-      end
-    })
-  end
-end
-
-function M.quartoClosePreview()
-  local success, quartoOutputBuf = pcall(api.nvim_buf_get_var, 0, 'quartoOutputBuf')
-  if not success then return end
-  if api.nvim_buf_is_loaded(quartoOutputBuf) then
-    api.nvim_buf_delete(quartoOutputBuf, { force = true })
-  end
-end
 
 local function get_language_content(bufnr)
   -- get and parse AST
@@ -124,7 +73,7 @@ local function get_language_content(bufnr)
   return results
 end
 
-local function update_language_buffers(qmd_bufnr)
+M.updateLanguageBuffers = function(qmd_bufnr)
   local language_content = get_language_content(qmd_bufnr)
   local bufnrs = {}
   for _, lang in ipairs(M.config.lspFeatures.languages) do
@@ -161,9 +110,118 @@ local function update_language_buffers(qmd_bufnr)
   return bufnrs
 end
 
+---Registered client and source mapping.
+M.cmp_client_source_map = {}
+
+---Setup cmp-nvim-lsp source.
+M.cmp_setup_source = function(qmdbufnr, bufnr)
+  local callback = function()
+    M.cmp_on_insert_enter(qmdbufnr, bufnr)
+  end
+  vim.api.nvim_create_autocmd('InsertEnter', {
+    buffer = qmdbufnr,
+    group = vim.api.nvim_create_augroup('cmp_quarto'..bufnr, { clear = true }),
+    callback = callback
+  })
+end
+
+---Refresh sources on InsertEnter.
+-- adds a source for the hidden language buffer bufnr
+M.cmp_on_insert_enter = function(qmdbufnr, bufnr)
+  local cmp = require('cmp')
+  local allowed_clients = {}
+
+  -- register all active clients.
+  for _, client in ipairs(vim.lsp.get_active_clients({bufnr = bufnr})) do
+    allowed_clients[client.id] = client
+    if not M.cmp_client_source_map[client.id] then
+      local s = source.new(client, qmdbufnr, bufnr, M.updateLanguageBuffers)
+      if s:is_available() then
+        P('register source for ' .. s.client.name)
+        M.cmp_client_source_map[client.id] = cmp.register_source('quarto', s)
+      end
+    end
+  end
+
+  -- register all buffer clients (early register before activation)
+  for _, client in ipairs(vim.lsp.buf_get_clients(0)) do
+    allowed_clients[client.id] = client
+    if not M.cmp_client_source_map[client.id] then
+      local s = source.new(client, qmdbufnr, bufnr, M.updateLanguageBuffers)
+      if s:is_available() then
+        M.cmp_client_source_map[client.id] = cmp.register_source('quarto', s)
+      end
+    end
+  end
+
+  -- unregister stopped/detached clients.
+  for client_id, source_id in pairs(M.cmp_client_source_map) do
+    if not allowed_clients[client_id] or allowed_clients[client_id]:is_stopped() then
+      cmp.unregister_source(source_id)
+      M.cmp_client_source_map[client_id] = nil
+    end
+  end
+end
+
+
+function M.quartoPreview()
+  -- find root directory / check if it is a project
+  local buffer_path = api.nvim_buf_get_name(0)
+  local root_dir = util.root_pattern("_quarto.yml")(buffer_path)
+  local cmd
+  local mode
+  if root_dir then
+    mode = "project"
+    cmd = 'quarto preview'
+  else
+    mode = "file"
+    cmd = 'quarto preview ' .. buffer_path
+  end
+
+  local quarto_extensions = { ".qmd", ".Rmd", ".ipynb", ".md" }
+  local file_extension = buffer_path:match("^.+(%..+)$")
+  if mode == "file" and not file_extension then
+    vim.notify("Not in a file. exiting.")
+    return
+  end
+  if mode == "file" and not tools.contains(quarto_extensions, file_extension) then
+    vim.notify("Not a quarto file, ends in " .. file_extension .. " exiting.")
+    return
+  end
+
+  -- run command in embedded terminal
+  -- in a new tab and go back to the buffer
+  vim.cmd('tabedit term://' .. cmd)
+  local quartoOutputBuf = vim.api.nvim_get_current_buf()
+  vim.cmd('tabprevious')
+  api.nvim_buf_set_var(0, 'quartoOutputBuf', quartoOutputBuf)
+
+
+  -- close preview terminal on exit of the quarto buffer
+  if M.config.closePreviewOnExit then
+    api.nvim_create_autocmd({ "QuitPre", "WinClosed" }, {
+      buffer = api.nvim_get_current_buf(),
+      group = api.nvim_create_augroup("quartoPreview", {}),
+      callback = function(_, _)
+        if api.nvim_buf_is_loaded(quartoOutputBuf) then
+          api.nvim_buf_delete(quartoOutputBuf, { force = true })
+        end
+      end
+    })
+  end
+end
+
+function M.quartoClosePreview()
+  local success, quartoOutputBuf = pcall(api.nvim_buf_get_var, 0, 'quartoOutputBuf')
+  if not success then return end
+  if api.nvim_buf_is_loaded(quartoOutputBuf) then
+    api.nvim_buf_delete(quartoOutputBuf, { force = true })
+  end
+end
+
 M.activateLspFeatures = function()
   local qmdbufnr = api.nvim_get_current_buf()
-  local bufnrs = update_language_buffers(qmdbufnr)
+  local bufnrs = M.updateLanguageBuffers(qmdbufnr)
 
   -- auto-close language files on qmd file close
   api.nvim_create_autocmd({ "QuitPre", "WinClosed" }, {
@@ -186,17 +244,31 @@ M.activateLspFeatures = function()
     M.enableDiagnostics()
   end
 
+  if M.config.lspFeatures.cmpSource.enabled then
+    for _,bufnr in ipairs(bufnrs) do
+      M.cmp_setup_source(qmdbufnr, bufnr)
+    end
+
+    api.nvim_create_autocmd({ "TextChangedI" }, {
+      buffer = 0,
+      group = api.nvim_create_augroup("quartoCmp", { clear = false }),
+      callback = function(_, _)
+        local bufnrs = M.updateLanguageBuffers(0)
+      end
+    })
+  end
+
   local key = M.config.keymap.hover
   vim.api.nvim_set_keymap('n', key, ":lua require'quarto'.quartoHover()<cr>", { silent = true })
 end
 
 M.enableDiagnostics = function()
   -- update diagnostics on changes
-  api.nvim_create_autocmd({ "CursorHold", "TextChanged" }, {
+  api.nvim_create_autocmd(M.config.lspFeatures.diagnostics.triggers, {
     buffer = 0,
     group = api.nvim_create_augroup("quartoLSPDiagnositcs", { clear = false }),
     callback = function(_, _)
-      local bufnrs = update_language_buffers(0)
+      local bufnrs = M.updateLanguageBuffers(0)
       for _, bufnr in ipairs(bufnrs) do
         local diag = vim.diagnostic.get(bufnr)
         local ns = api.nvim_create_namespace('quarto-lang-' .. bufnr)
@@ -209,7 +281,7 @@ end
 
 M.quartoHover = function()
   local qmdbufnr = api.nvim_get_current_buf()
-  local bufnrs = update_language_buffers(qmdbufnr)
+  local bufnrs = M.updateLanguageBuffers(qmdbufnr)
   for _, bufnr in ipairs(bufnrs) do
     local uri = vim.uri_from_bufnr(bufnr)
     local position_params = vim.lsp.util.make_position_params()
@@ -243,7 +315,7 @@ end
 
 -- setup
 M.setup = function(opt)
-  M.config = vim.tbl_deep_extend('force', defaultConfig, opt or {})
+  M.config = vim.tbl_deep_extend('force', M.defaultConfig, opt or {})
 end
 
 M.debug = function()
diff --git a/lua/quarto/source.lua b/lua/quarto/source.lua
new file mode 100644
index 0000000..2808e15
--- /dev/null
+++ b/lua/quarto/source.lua
@@ -0,0 +1,163 @@
+-- derived from <https://github.com/hrsh7th/cmp-nvim-lsp>
+local source = {}
+
+source.new = function(client, qmdbufnr, bufnr, updater)
+  local self = setmetatable({}, { __index = source })
+  self.client = client
+  self.bufnr = bufnr
+  self.qmdbufnr = qmdbufnr
+  self.id = qmdbufnr
+  self.request_ids = {}
+  self.updater = updater
+  return self
+end
+
+---Get debug name.
+---@return string
+source.get_debug_name = function(self)
+  return table.concat({ 'quarto', self.client.name }, ':')
+end
+
+---Return the source is available.
+---@return boolean
+source.is_available = function(self)
+  -- client is stopped.
+  if self.client.is_stopped() then
+    return false
+  end
+
+  -- don't filter clients from other buffers
+  -- client is not attached to current buffer.
+  -- if not vim.lsp.buf_get_clients(vim.api.nvim_get_current_buf())[self.client.id] then
+  --   return false
+  -- end
+
+  -- client has no completion capability.
+  if not self:_get(self.client.server_capabilities, { 'completionProvider' }) then
+    return false
+  end
+  return true;
+end
+
+---Get LSP's PositionEncodingKind.
+---@return lsp.PositionEncodingKind
+source.get_position_encoding_kind = function(self)
+  return self:_get(self.client.server_capabilities, { 'positionEncoding' }) or self.client.offset_encoding or 'utf-16'
+end
+
+---Get triggerCharacters.
+---@return string[]
+source.get_trigger_characters = function(self)
+  return self:_get(self.client.server_capabilities, { 'completionProvider', 'triggerCharacters' }) or {}
+end
+
+---Get get_keyword_pattern.
+---@param params cmp.SourceApiParams
+---@return string
+source.get_keyword_pattern = function(self, params)
+  return (params.option or {})[self.client.name] or require('cmp').get_config().completion.keyword_pattern
+end
+
+---Resolve LSP CompletionItem.
+---@param params cmp.SourceCompletionApiParams
+---@param callback function
+source.complete = function(self, params, callback)
+  local bufnrs = self.updater(self.qmdbufnr)
+  local win = vim.api.nvim_get_current_win()
+  local lsp_params = vim.lsp.util.make_position_params(win, self.client.offset_encoding)
+  lsp_params.textDocument = {
+    uri = vim.uri_from_bufnr(self.bufnr)
+  }
+  lsp_params.context = {}
+  lsp_params.context.triggerKind = params.completion_context.triggerKind
+  lsp_params.context.triggerCharacter = params.completion_context.triggerCharacter
+  self:_request('textDocument/completion', lsp_params, function(_, response)
+    callback(response)
+  end)
+end
+
+---Resolve LSP CompletionItem.
+---@param completion_item lsp.CompletionItem
+---@param callback function
+source.resolve = function(self, completion_item, callback)
+  -- client is stopped.
+  if self.client.is_stopped() then
+    return callback()
+  end
+
+  -- client has no completion capability.
+  if not self:_get(self.client.server_capabilities, { 'completionProvider', 'resolveProvider' }) then
+    return callback()
+  end
+
+  self:_request('completionItem/resolve', completion_item, function(_, response)
+    callback(response or completion_item)
+  end)
+end
+
+---Execute LSP CompletionItem.
+---@param completion_item lsp.CompletionItem
+---@param callback function
+source.execute = function(self, completion_item, callback)
+  -- client is stopped.
+  if self.client.is_stopped() then
+    return callback()
+  end
+
+  -- completion_item has no command.
+  if not completion_item.command then
+    return callback()
+  end
+
+  self:_request('workspace/executeCommand', completion_item.command, function(_, _)
+    callback()
+  end)
+end
+
+---Get object path.
+---@param root table
+---@param paths string[]
+---@return any
+source._get = function(_, root, paths)
+  local c = root
+  for _, path in ipairs(paths) do
+    c = c[path]
+    if not c then
+      return nil
+    end
+  end
+  return c
+end
+
+---Send request to nvim-lsp servers with backward compatibility.
+---@param method string
+---@param params table
+---@param callback function
+source._request = function(self, method, params, callback)
+  if self.request_ids[method] ~= nil then
+    self.client.cancel_request(self.request_ids[method])
+    self.request_ids[method] = nil
+  end
+  local _, request_id
+  _, request_id = self.client.request(method, params, function(arg1, arg2, arg3)
+    if self.request_ids[method] ~= request_id then
+      return
+    end
+    self.request_ids[method] = nil
+
+    -- Text changed, retry
+    if arg1 and arg1.code == -32801 then
+      self:_request(method, params, callback)
+      return
+    end
+
+    if method == arg2 then
+      callback(arg1, arg3) -- old signature
+    else
+      callback(arg1, arg2) -- new signature
+    end
+  end)
+  self.request_ids[method] = request_id
+end
+
+return source
diff --git a/lua/quarto/tools.lua b/lua/quarto/tools.lua
index d6141d7..1d5911d 100644
--- a/lua/quarto/tools.lua
+++ b/lua/quarto/tools.lua
@@ -1,4 +1,4 @@
-M = {}
+local M = {}
 
 M.contains = function(list, x)
   for _, v in pairs(list) do
@@ -25,5 +25,9 @@ M.spaces = function(n)
   return s
 end
 
-return M
+M.if_nil = function(val, default)
+  if val == nil then return default end
+  return val
+end
 
+return M
diff --git a/plugin/quarto.lua b/plugin/quarto.lua
index 2149d1d..798e64d 100644
--- a/plugin/quarto.lua
+++ b/plugin/quarto.lua
@@ -1,4 +1,5 @@
 local quarto = require'quarto'
+local config = require'quarto'.config
 local a = vim.api
 
 a.nvim_create_user_command('QuartoPreview', quarto.quartoPreview, {})
@@ -11,10 +12,9 @@ a.nvim_create_user_command('QuartoHover', quarto.quartoHover, {})
 a.nvim_create_autocmd({"BufEnter"}, {
   pattern = {"*.qmd"},
   callback = function ()
-    quarto = require'quarto'
-    if quarto.config.lspFeatures.enabled then
+    if config.lspFeatures.enabled then
       quarto.activateLspFeatures()
-      if quarto.config.lspFeatures.diagnostics.enabled then
+      if config.lspFeatures.diagnostics.enabled then
         quarto.enableDiagnostics()
       end
     end