feat(treesitter): add a query editor (#24703)

6 files changed, 205 insertions, 24 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 82b390853c..637a33b555 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -133,6 +133,8 @@ The following new APIs and features were added.
     `vim.treesitter.language.register`.
   • The `#set!` directive now supports `injection.self` and `injection.parent` for injecting either the current node's language
     or the parent LanguageTree's language, respectively.
+  • Added `vim.treesitter.preview_query()`, for live editing of treesitter
+    queries.
 
 • |vim.ui.open()| opens URIs using the system default handler (macOS `open`,
   Windows `explorer`, Linux `xdg-open`, etc.)
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 139b986786..287985f75b 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -676,8 +676,9 @@ inspect_tree({opts})                           *vim.treesitter.inspect_tree()*
     language tree.
 
     While in the window, press "a" to toggle display of anonymous nodes, "I"
-    to toggle the display of the source language of each node, and press
-    <Enter> to jump to the node under the cursor in the source buffer.
+    to toggle the display of the source language of each node, "o" to toggle
+    the query previewer, and press <Enter> to jump to the node under the
+    cursor in the source buffer.
 
     Can also be shown with `:InspectTree`.                      *:InspectTree*
 
@@ -730,6 +731,11 @@ node_contains({node}, {range})                *vim.treesitter.node_contains()*
     Return: ~
         (boolean) True if the {node} contains the {range}
 
+preview_query()                               *vim.treesitter.preview_query()*
+    Open a window for live editing of a treesitter query.
+
+    Can also be shown with `:PreviewQuery`.                    *:PreviewQuery*
+
 start({bufnr}, {lang})                                *vim.treesitter.start()*
     Starts treesitter highlighting for a buffer
 
diff --git a/runtime/ftplugin/query.lua b/runtime/ftplugin/query.lua
index accf38c199..964c221ad4 100644
--- a/runtime/ftplugin/query.lua
+++ b/runtime/ftplugin/query.lua
@@ -1,6 +1,6 @@
 -- Neovim filetype plugin file
 -- Language:	Tree-sitter query
--- Last Change:	2022 Apr 25
+-- Last Change:	2023 Aug 23
 
 if vim.b.did_ftplugin == 1 then
   return
@@ -14,6 +14,8 @@ vim.treesitter.start()
 -- set omnifunc
 vim.bo.omnifunc = 'v:lua.vim.treesitter.query.omnifunc'
 
+vim.opt_local.iskeyword:append('.')
+
 -- 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.lua b/runtime/lua/vim/treesitter.lua
index 04420c81e3..4f84fc2e0f 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -472,8 +472,8 @@ end
 --- Open a window that displays a textual representation of the nodes in the language tree.
 ---
 --- While in the window, press "a" to toggle display of anonymous nodes, "I" to toggle the
---- display of the source language of each node, and press <Enter> to jump to the node under the
---- cursor in the source buffer.
+--- display of the source language of each node, "o" to toggle the query previewer, and press
+--- <Enter> to jump to the node under the cursor in the source buffer.
 ---
 --- Can also be shown with `:InspectTree`. *:InspectTree*
 ---
@@ -494,6 +494,13 @@ function M.inspect_tree(opts)
   require('vim.treesitter.dev').inspect_tree(opts)
 end
 
+--- Open a window for live editing of a treesitter query.
+---
+--- Can also be shown with `:PreviewQuery`. *:PreviewQuery*
+function M.preview_query()
+  require('vim.treesitter.dev').preview_query()
+end
+
 --- Returns the fold level for {lnum} in the current buffer. Can be set directly to 'foldexpr':
 --- <pre>lua
 --- vim.wo.foldexpr = 'v:lua.vim.treesitter.foldexpr()'
diff --git a/runtime/lua/vim/treesitter/dev.lua b/runtime/lua/vim/treesitter/dev.lua
index f7625eb94b..b7f2c0e473 100644
--- a/runtime/lua/vim/treesitter/dev.lua
+++ b/runtime/lua/vim/treesitter/dev.lua
@@ -124,7 +124,7 @@ function TSTreeView:new(bufnr, lang)
   end
 
   local t = {
-    ns = api.nvim_create_namespace(''),
+    ns = api.nvim_create_namespace('treesitter/dev-inspect'),
     nodes = nodes,
     named = named,
     opts = {
@@ -158,6 +158,29 @@ local function escape_quotes(text)
   return string.format('"%s"', text:sub(2, #text - 1):gsub('"', '\\"'))
 end
 
+---@param w integer
+---@return boolean closed Whether the window was closed.
+local function close_win(w)
+  if api.nvim_win_is_valid(w) then
+    api.nvim_win_close(w, true)
+    return true
+  end
+
+  return false
+end
+
+---@param w integer
+---@param b integer
+local function set_dev_properties(w, b)
+  vim.wo[w].scrolloff = 5
+  vim.wo[w].wrap = false
+  vim.wo[w].foldmethod = 'manual' -- disable folding
+  vim.bo[b].buflisted = false
+  vim.bo[b].buftype = 'nofile'
+  vim.bo[b].bufhidden = 'wipe'
+  vim.bo[b].filetype = 'query'
+end
+
 --- Write the contents of this View into {bufnr}.
 ---
 ---@param bufnr integer Buffer number to write into.
@@ -247,12 +270,9 @@ function M.inspect_tree(opts)
   local win = api.nvim_get_current_win()
   local pg = assert(TSTreeView:new(buf, opts.lang))
 
-  -- Close any existing dev window
-  if vim.b[buf].dev then
-    local w = vim.b[buf].dev
-    if api.nvim_win_is_valid(w) then
-      api.nvim_win_close(w, true)
-    end
+  -- Close any existing inspector window
+  if vim.b[buf].dev_inspect then
+    close_win(vim.b[buf].dev_inspect)
   end
 
   local w = opts.winid
@@ -268,16 +288,10 @@ function M.inspect_tree(opts)
     b = api.nvim_win_get_buf(w)
   end
 
-  vim.b[buf].dev = w
-
-  vim.wo[w].scrolloff = 5
-  vim.wo[w].wrap = false
-  vim.wo[w].foldmethod = 'manual' -- disable folding
-  vim.bo[b].buflisted = false
-  vim.bo[b].buftype = 'nofile'
-  vim.bo[b].bufhidden = 'wipe'
+  vim.b[buf].dev_inspect = w
+  vim.b[b].dev_base = win -- base window handle
   vim.b[b].disable_query_linter = true
-  vim.bo[b].filetype = 'query'
+  set_dev_properties(w, b)
 
   local title --- @type string?
   local opts_title = opts.title
@@ -306,7 +320,7 @@ function M.inspect_tree(opts)
   api.nvim_buf_set_keymap(b, 'n', 'a', '', {
     desc = 'Toggle anonymous nodes',
     callback = function()
-      local row, col = unpack(api.nvim_win_get_cursor(w))
+      local row, col = unpack(api.nvim_win_get_cursor(w)) ---@type integer, integer
       local curnode = pg:get(row)
       while curnode and not curnode.named do
         row = row - 1
@@ -336,6 +350,15 @@ function M.inspect_tree(opts)
       pg:draw(b)
     end,
   })
+  api.nvim_buf_set_keymap(b, 'n', 'o', '', {
+    desc = 'Toggle query previewer',
+    callback = function()
+      local preview_w = vim.b[buf].dev_preview
+      if not preview_w or not close_win(preview_w) then
+        M.preview_query()
+      end
+    end,
+  })
 
   local group = api.nvim_create_augroup('treesitter/dev', {})
 
@@ -436,11 +459,148 @@ function M.inspect_tree(opts)
     buffer = buf,
     once = true,
     callback = function()
-      if api.nvim_win_is_valid(w) then
-        api.nvim_win_close(w, true)
+      close_win(w)
+    end,
+  })
+end
+
+local preview_ns = api.nvim_create_namespace('treesitter/dev-preview')
+
+---@param query_win integer
+---@param base_win integer
+local function update_preview_highlights(query_win, base_win)
+  local base_buf = api.nvim_win_get_buf(base_win)
+  local query_buf = api.nvim_win_get_buf(query_win)
+  local parser = vim.treesitter.get_parser(base_buf)
+  local lang = parser:lang()
+  api.nvim_buf_clear_namespace(base_buf, preview_ns, 0, -1)
+  local query_content = table.concat(api.nvim_buf_get_lines(query_buf, 0, -1, false), '\n')
+
+  local ok_query, query = pcall(vim.treesitter.query.parse, lang, query_content)
+  if not ok_query then
+    return
+  end
+
+  local cursor_word = vim.fn.expand('<cword>') --[[@as string]]
+  -- Only highlight captures if the cursor is on a capture name
+  if cursor_word:find('^@') == nil then
+    return
+  end
+  -- Remove the '@' from the cursor word
+  cursor_word = cursor_word:sub(2)
+  local topline, botline = vim.fn.line('w0', base_win), vim.fn.line('w$', base_win)
+  for id, node in query:iter_captures(parser:trees()[1]:root(), base_buf, topline - 1, botline) do
+    local capture_name = query.captures[id]
+    if capture_name == cursor_word then
+      local lnum, col, end_lnum, end_col = node:range()
+      api.nvim_buf_set_extmark(base_buf, preview_ns, lnum, col, {
+        end_row = end_lnum,
+        end_col = end_col,
+        hl_group = 'Visual',
+        virt_text = {
+          { capture_name, 'Title' },
+        },
+      })
+    end
+  end
+end
+
+--- @private
+function M.preview_query()
+  local buf = api.nvim_get_current_buf()
+  local win = api.nvim_get_current_win()
+
+  -- Close any existing previewer window
+  if vim.b[buf].dev_preview then
+    close_win(vim.b[buf].dev_preview)
+  end
+
+  local cmd = '60vnew'
+  -- If the inspector is open, place the previewer above it.
+  local base_win = vim.b[buf].dev_base ---@type integer?
+  local base_buf = base_win and api.nvim_win_get_buf(base_win)
+  local inspect_win = base_buf and vim.b[base_buf].dev_inspect
+  if base_win and base_buf and api.nvim_win_is_valid(inspect_win) then
+    vim.api.nvim_set_current_win(inspect_win)
+    buf = base_buf
+    win = base_win
+    cmd = 'new'
+  end
+  vim.cmd(cmd)
+
+  local ok, parser = pcall(vim.treesitter.get_parser, buf)
+  if not ok then
+    return nil, 'No parser available for the given buffer'
+  end
+  local lang = parser:lang()
+
+  local query_win = api.nvim_get_current_win()
+  local query_buf = api.nvim_win_get_buf(query_win)
+
+  vim.b[buf].dev_preview = query_win
+  vim.bo[query_buf].omnifunc = 'v:lua.vim.treesitter.query.omnifunc'
+  set_dev_properties(query_win, query_buf)
+
+  -- Note that omnifunc guesses the language based on the containing folder,
+  -- so we add the parser's language to the buffer's name so that omnifunc
+  -- can infer the language later.
+  api.nvim_buf_set_name(query_buf, string.format('%s/query_previewer.scm', lang))
+
+  local group = api.nvim_create_augroup('treesitter/dev-preview', {})
+  api.nvim_create_autocmd({ 'TextChanged', 'InsertLeave' }, {
+    group = group,
+    buffer = query_buf,
+    desc = 'Update query previewer diagnostics when the query changes',
+    callback = function()
+      vim.treesitter.query.lint(query_buf, { langs = lang, clear = false })
+    end,
+  })
+  api.nvim_create_autocmd({ 'TextChanged', 'InsertLeave', 'CursorMoved', 'BufEnter' }, {
+    group = group,
+    buffer = query_buf,
+    desc = 'Update query previewer highlights when the cursor moves',
+    callback = function()
+      update_preview_highlights(query_win, win)
+    end,
+  })
+  api.nvim_create_autocmd('BufLeave', {
+    group = group,
+    buffer = query_buf,
+    desc = 'Clear the query previewer highlights when leaving the previewer',
+    callback = function()
+      api.nvim_buf_clear_namespace(buf, preview_ns, 0, -1)
+    end,
+  })
+  api.nvim_create_autocmd('BufLeave', {
+    group = group,
+    buffer = buf,
+    desc = 'Clear the query previewer highlights when leaving the source buffer',
+    callback = function()
+      if not api.nvim_buf_is_loaded(query_buf) then
+        return true
       end
+
+      api.nvim_buf_clear_namespace(query_buf, preview_ns, 0, -1)
+    end,
+  })
+  api.nvim_create_autocmd('BufHidden', {
+    group = group,
+    buffer = buf,
+    desc = 'Close the previewer window when the source buffer is hidden',
+    once = true,
+    callback = function()
+      close_win(query_win)
     end,
   })
+
+  api.nvim_buf_set_lines(query_buf, 0, -1, false, {
+    ';; Write your query here. Use @captures to highlight matches in the source buffer.',
+    ';; Completion for grammar nodes is available (see :h compl-omni)',
+    '',
+    '',
+  })
+  vim.cmd('normal! G')
+  vim.cmd.startinsert()
 end
 
 return M
diff --git a/runtime/plugin/nvim.lua b/runtime/plugin/nvim.lua
index 0a33826b82..2ddccfcff6 100644
--- a/runtime/plugin/nvim.lua
+++ b/runtime/plugin/nvim.lua
@@ -18,3 +18,7 @@ vim.api.nvim_create_user_command('InspectTree', function(cmd)
     vim.treesitter.inspect_tree()
   end
 end, { desc = 'Inspect treesitter language tree for buffer', count = true })
+
+vim.api.nvim_create_user_command('PreviewQuery', function()
+  vim.treesitter.preview_query()
+end, { desc = 'Preview treesitter query' })