fix(treesitter): ":EditQuery [lang]" with injected languages #34914

Problem:
`:EditQuery` command accepts a language argument, but it doesn't
highlight properly for injected languages.

Solution:
- Fully parse with the root language and then filter the query on the
  child trees that are of the language requested.
- Also support completion (`EditQuery <tab>`).

runtime/doc/treesitter.txt

@@ -1323,6 +1323,9 @@ edit({lang})                                     *vim.treesitter.query.edit()*
 
     Can also be shown with `:EditQuery`.                          *:EditQuery*
 
+    `:EditQuery <tab>` completes the treesitter parser names in the runtime
+    path.
+
     If you move the cursor to a capture name ("@foo"), text matching the
     capture is highlighted in the source buffer. The query editor is a scratch
     buffer, use `:write` to save it. You can find example queries at

runtime/lua/vim/_defaults.lua

@@ -23,7 +23,13 @@ do
 
   vim.api.nvim_create_user_command('EditQuery', function(cmd)
     vim.treesitter.query.edit(cmd.fargs[1])
-  end, { desc = 'Edit treesitter query', nargs = '?' })
+  end, {
+    desc = 'Edit treesitter query',
+    nargs = '?',
+    complete = function()
+      return vim.treesitter.language._complete()
+    end,
+  })
 
   vim.api.nvim_create_user_command('Open', function(cmd)
     vim.ui.open(assert(cmd.fargs[1]))

runtime/lua/vim/treesitter/dev.lua

@@ -565,7 +565,8 @@ local edit_ns = api.nvim_create_namespace('nvim.treesitter.dev_edit')
 local function update_editor_highlights(query_win, base_win, lang)
   local base_buf = api.nvim_win_get_buf(base_win)
   local query_buf = api.nvim_win_get_buf(query_win)
-  local parser = assert(vim.treesitter.get_parser(base_buf, lang, { error = false }))
+  local root_lang = vim.treesitter.language.get_lang(vim.bo[base_buf].filetype)
+  local parser = assert(vim.treesitter.get_parser(base_buf, root_lang, { error = false }))
   api.nvim_buf_clear_namespace(base_buf, edit_ns, 0, -1)
   local query_content = table.concat(api.nvim_buf_get_lines(query_buf, 0, -1, false), '\n')
 
@@ -581,21 +582,30 @@ local function update_editor_highlights(query_win, base_win, lang)
   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)
+  -- Parse buffer including injected languages.
-  for id, node in query:iter_captures(parser:trees()[1]:root(), base_buf, topline - 1, botline) do
+  parser:parse(true)
-    local capture_name = query.captures[id]
+  -- Query on the trees of the language requested to highlight captures.
-    if capture_name == cursor_word then
+  parser:for_each_tree(function(tree, ltree)
-      local lnum, col, end_lnum, end_col = node:range()
+    if ltree:lang() ~= lang then
-      api.nvim_buf_set_extmark(base_buf, edit_ns, lnum, col, {
+      return
-        end_row = end_lnum,
-        end_col = end_col,
-        hl_group = 'Visual',
-        virt_text = {
-          { capture_name, 'Title' },
-        },
-      })
     end
-  end
+    local root = tree:root()
+    local topline, botline = vim.fn.line('w0', base_win), vim.fn.line('w$', base_win)
+    for id, node in query:iter_captures(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, edit_ns, lnum, col, {
+          end_row = end_lnum,
+          end_col = end_col,
+          hl_group = 'Visual',
+          virt_text = {
+            { capture_name, 'Title' },
+          },
+        })
+      end
+    end
+  end)
 end
 
 --- @nodoc

runtime/lua/vim/treesitter/language.lua

@@ -186,4 +186,15 @@ function M.inspect(lang)
   return vim._ts_inspect_language(lang)
 end
 
+--- Returns available treesitter languages.
+function M._complete()
+  local parsers = vim.api.nvim_get_runtime_file('parser/*', true)
+  local parser_names_set = {} ---@type table<string, boolean>
+  for _, parser in ipairs(parsers) do
+    local parser_name = vim.fn.fnamemodify(parser, ':t:r')
+    parser_names_set[parser_name] = true
+  end
+  return vim.tbl_keys(parser_names_set)
+end
+
 return M

runtime/lua/vim/treesitter/query.lua

@@ -1165,6 +1165,8 @@ end
 ---
 --- Can also be shown with `:EditQuery`. [:EditQuery]()
 ---
+--- `:EditQuery <tab>` completes the treesitter parser names in the runtime path.
+---
 --- If you move the cursor to a capture name ("@foo"), text matching the capture is highlighted in
 --- the source buffer. The query editor is a scratch buffer, use `:write` to save it. You can find
 --- example queries at `$VIMRUNTIME/queries/`.