From 668f16bac779ac52d7bd9452e6001a7a6d1e9965 Mon Sep 17 00:00:00 2001
From: Christian Clason <c.clason@uni-graz.at>
Date: Sun, 30 Apr 2023 11:01:54 +0200
Subject: [PATCH] feat(treesitter): upstream query omnifunc from playground
 (#23394)

and set by default in `ftplugin/query.lua`
---
 runtime/doc/news.txt                         |  6 +-
 runtime/doc/treesitter.txt                   |  9 ++-
 runtime/ftplugin/query.lua                   |  3 +
 runtime/lua/vim/treesitter/_query_linter.lua | 60 +++++++++++++++++++-
 runtime/lua/vim/treesitter/query.lua         | 12 +++-
 5 files changed, 84 insertions(+), 6 deletions(-)

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index c343525a09..f33cffa22e 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -37,14 +37,18 @@ The following new APIs or features were added.
 
 • |vim.iter()| provides a generic iterator interface for tables and Lua
 iterators |luaref-in|.
+
 • Added |vim.keycode()| for translating keycodes in a string.
 
 • Added automatic linting of treesitter query files (see |ft-query-plugin|).
   Automatic linting can be turned off via >lua
-	vim.g.query_lint_on = {}
+      vim.g.query_lint_on = {}
 <
 • Enabled treesitter highlighting for treesitter query files by default.
 
+• Added |vim.treesitter.query.omnifunc()| for treesitter query files (set by
+  default).
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 32c97ce3ad..94690f0b7f 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -853,7 +853,7 @@ lint({buf}, {opts})                              *vim.treesitter.query.lint()*
 
     The found diagnostics are reported using |diagnostic-api|. By default, the
     parser used for verification is determined by the containing folder of the
-    query file, e.g., if the path is `**/lua/highlights.scm` , the parser for the `lua` language will be used.
+    query file, e.g., if the path ends in `/lua/highlights.scm` , the parser for the `lua` language will be used.
 
     Parameters: ~
       • {buf}   (integer) Buffer handle
@@ -875,6 +875,13 @@ list_predicates()                     *vim.treesitter.query.list_predicates()*
     Return: ~
         string[] List of supported predicates.
 
+omnifunc({findstart}, {base})                *vim.treesitter.query.omnifunc()*
+    Omnifunc for completing node names and predicates in treesitter queries.
+
+    Use via >lua
+      vim.bo.omnifunc = 'v:lua.vim.treesitter.query.omnifunc'
+<
+
 parse({lang}, {query})                          *vim.treesitter.query.parse()*
     Parse {query} as a string. (If the query is in a file, the caller should
     read the contents into a string before calling).
diff --git a/runtime/ftplugin/query.lua b/runtime/ftplugin/query.lua
index 842d338fd9..accf38c199 100644
--- a/runtime/ftplugin/query.lua
+++ b/runtime/ftplugin/query.lua
@@ -11,6 +11,9 @@ end
 -- use treesitter over syntax
 vim.treesitter.start()
 
+-- set omnifunc
+vim.bo.omnifunc = 'v:lua.vim.treesitter.query.omnifunc'
+
 -- query linter
 local buf = vim.api.nvim_get_current_buf()
 local query_lint_on = vim.g.query_lint_on or { 'BufEnter', 'BufWrite' }
diff --git a/runtime/lua/vim/treesitter/_query_linter.lua b/runtime/lua/vim/treesitter/_query_linter.lua
index 62f28d3097..ecdee5fc95 100644
--- a/runtime/lua/vim/treesitter/_query_linter.lua
+++ b/runtime/lua/vim/treesitter/_query_linter.lua
@@ -1,4 +1,6 @@
-local namespace = vim.api.nvim_create_namespace('vim.treesitter.query_linter')
+local api = vim.api
+
+local namespace = api.nvim_create_namespace('vim.treesitter.query_linter')
 -- those node names exist for every language
 local BUILT_IN_NODE_NAMES = { '_', 'ERROR' }
 
@@ -49,7 +51,7 @@ end
 --- @param buf integer
 --- @return string?
 local function guess_query_lang(buf)
-  local filename = vim.api.nvim_buf_get_name(buf)
+  local filename = api.nvim_buf_get_name(buf)
   if filename ~= '' then
     local ok, query_lang = pcall(vim.fn.fnamemodify, filename, ':p:h:t')
     if ok then
@@ -256,7 +258,7 @@ end
 --- @param opts QueryLinterOpts|QueryLinterNormalizedOpts|nil Options for linting
 function M.lint(buf, opts)
   if buf == 0 then
-    buf = vim.api.nvim_get_current_buf()
+    buf = api.nvim_get_current_buf()
   end
 
   local diagnostics = {}
@@ -299,4 +301,56 @@ function M.clear(buf)
   vim.diagnostic.reset(namespace, buf)
 end
 
+--- @private
+--- @param findstart integer
+--- @param base string
+function M.omnifunc(findstart, base)
+  if findstart == 1 then
+    local result =
+      api.nvim_get_current_line():sub(1, api.nvim_win_get_cursor(0)[2]):find('["#%-%w]*$')
+    return result - 1
+  end
+
+  local buf = api.nvim_get_current_buf()
+  local query_lang = guess_query_lang(buf)
+
+  local ok, parser_info = pcall(vim.treesitter.language.inspect, query_lang)
+  if not ok then
+    return -2
+  end
+
+  local items = {}
+  for _, f in pairs(parser_info.fields) do
+    if f:find(base, 1, true) then
+      table.insert(items, f .. ':')
+    end
+  end
+  for _, p in pairs(vim.treesitter.query.list_predicates()) do
+    local text = '#' .. p
+    local found = text:find(base, 1, true)
+    if found and found <= 2 then -- with or without '#'
+      table.insert(items, text)
+    end
+    text = '#not-' .. p
+    found = text:find(base, 1, true)
+    if found and found <= 2 then -- with or without '#'
+      table.insert(items, text)
+    end
+  end
+  for _, p in pairs(vim.treesitter.query.list_directives()) do
+    local text = '#' .. p
+    local found = text:find(base, 1, true)
+    if found and found <= 2 then -- with or without '#'
+      table.insert(items, text)
+    end
+  end
+  for _, s in pairs(parser_info.symbols) do
+    local text = s[2] and s[1] or '"' .. s[1]:gsub([[\]], [[\\]]) .. '"'
+    if text:find(base, 1, true) then
+      table.insert(items, text)
+    end
+  end
+  return { words = items, refresh = 'always' }
+end
+
 return M
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 492bfd1ffb..93841bb31e 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -728,7 +728,7 @@ end
 ---
 --- The found diagnostics are reported using |diagnostic-api|.
 --- By default, the parser used for verification is determined by the containing folder
---- of the query file, e.g., if the path is `**/lua/highlights.scm`, the parser for the
+--- of the query file, e.g., if the path ends in `/lua/highlights.scm`, the parser for the
 --- `lua` language will be used.
 ---@param buf (integer) Buffer handle
 ---@param opts (QueryLinterOpts|nil) Optional keyword arguments:
@@ -743,4 +743,14 @@ function M.lint(buf, opts)
   end
 end
 
+--- Omnifunc for completing node names and predicates in treesitter queries.
+---
+--- Use via
+--- <pre>lua
+---   vim.bo.omnifunc = 'v:lua.vim.treesitter.query.omnifunc'
+--- </pre>
+function M.omnifunc(findstart, base)
+  return require('vim.treesitter._query_linter').omnifunc(findstart, base)
+end
+
 return M