Repository: neovim/neovim
File: runtime/lua/vim/treesitter/dev.lua (Lua)

Code before:
---@param base_win integer
---@param lang string
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 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')

  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)
  parser:parse(true)
  parser:for_each_tree(function(tree, ltree)
    if ltree:lang() ~= lang then
      return
    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
--- @param lang? string language to open the query editor for.
--- @return boolean? `true` on success, `nil` on failure
--- @return string? error message, if applicable
function M.edit_query(lang)

Review comment (suggestion):
```suggestion
  -- <one-line comment that gives a hint about what this does>
  parser:for_each_tree(function(tree, ltree)
```

Code after:
---@param base_win integer
---@param lang string
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 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')

  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)
  -- Parse buffer including injected languages.
  parser:parse(true)
  -- Query on the trees of the language requested to highlight captures.
  parser:for_each_tree(function(tree, ltree)
    if ltree:lang() ~= lang then
      return
    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
--- @param lang? string language to open the query editor for.
--- @return boolean? `true` on success, `nil` on failure