refactor(printer): integrate injected languages

Co-authored-by: Jędrzej Boczar <yendreij@gmail.com>

README.md

@@ -43,7 +43,13 @@ The tree can be toggled using the command `:TSPlaygroundToggle`.
 ### Keybindings
 
 - `R`: Refreshes the playground view when focused or reloads the query when the query editor is focused.
-- `o`: Toggles the query editor when the playground is focused
+- `o`: Toggles the query editor when the playground is focused.
+- `a`: Toggles visibility of anonymous nodes.
+- `i`: Toggles visibility of highlight groups.
+- `I`: Toggles visibility of the language the node belongs to.
+- `t`: Toggles visibility of injected languages.
+- `f`: Focuses the language tree under the cursor in the playground. The query editor will now be using the focused language.
+- `F`: Unfocuses the currently focused language.
 - `<cr>`: Go to current node in code buffer
 
 ## Query Linter
@@ -61,6 +67,3 @@ require "nvim-treesitter.configs".setup {
 ```
 
 ![image](https://user-images.githubusercontent.com/7189118/101246661-06089a00-3715-11eb-9c57-6d6439defbf8.png)
-
-# Roadmap
-  - [ ] Add interactive query highlighting

lua/nvim-treesitter-playground.lua

@@ -1,10 +1,6 @@
 local parsers = require 'nvim-treesitter.parsers'
 local M = {}
 
-vim.cmd [[
-  command! TSHighlightCapturesUnderCursor :lua require'nvim-treesitter-playground.hl-info'.show_hl_captures()<cr>
-]]
-
 function M.init()
   require "nvim-treesitter".define_modules {
     playground = {
@@ -21,6 +17,10 @@ function M.init()
       end,
     },
   }
+
+  vim.cmd [[
+    command! TSHighlightCapturesUnderCursor :lua require'nvim-treesitter-playground.hl-info'.show_hl_captures()<cr>
+  ]]
 end
 
 return M

lua/nvim-treesitter-playground/hl-info.lua

@@ -18,18 +18,22 @@ function M.show_hl_captures()
   local parser = parsers.get_parser(bufnr, lang)
   if not parser then return function() end end
 
-  local root = parser:parse()[1]:root()
-  if not root then return end
-  local start_row, _, end_row, _ = root:range()
+  parser = parser:language_for_range({row, col, row, col})
 
   local matches = {}
-  local query = queries.get_query(lang, 'highlights')
-  for _, match in query:iter_matches(root, bufnr, start_row, end_row) do
-    for id, node in pairs(match) do
-      if ts_utils.is_in_node_range(node, row, col) then
-        local c = query.captures[id] -- name of the capture in the query
-        if c ~= nil then
-          table.insert(matches, '@'..c..' -> '..(hlmap[c] or 'nil'))
+  local query = queries.get_query(parser:lang(), 'highlights')
+
+  for _, tree in ipairs(parser:trees()) do
+    local root = tree:root()
+    local start_row, _, end_row, _ = root:range()
+
+    for _, match in query:iter_matches(root, bufnr, start_row, end_row) do
+      for id, node in pairs(match) do
+        if ts_utils.is_in_node_range(node, row, col) then
+          local c = query.captures[id] -- name of the capture in the query
+          if c ~= nil then
+            table.insert(matches, '@'..c..' -> '..(hlmap[c] or 'nil'))
+          end
         end
       end
     end

lua/nvim-treesitter-playground/internal.lua

@@ -1,3 +1,4 @@
+local parsers = require 'nvim-treesitter.parsers'
 local configs = require 'nvim-treesitter.configs'
 local ts_utils = require 'nvim-treesitter.ts_utils'
 local printer = require 'nvim-treesitter-playground.printer'
@@ -23,7 +24,13 @@ M._entries = setmetatable({}, {
     local entry = rawget(tbl, key)
 
     if not entry then
-      entry = {}
+      entry = {
+        include_anonymous_nodes = false,
+        suppress_injected_languages = false,
+        include_language = false,
+        include_hl_groups = false,
+        focused_language_tree = nil
+      }
       rawset(tbl, key, entry)
     end
 
@@ -35,9 +42,16 @@ local playground_ns = api.nvim_create_namespace('nvim-treesitter-playground')
 local query_hl_ns = api.nvim_create_namespace('nvim-treesitter-playground-query')
 
 local function get_node_at_cursor()
-  local success, node_at_point = pcall(function() return ts_utils.get_node_at_cursor() end)
+  local lnum, col = unpack(vim.api.nvim_win_get_cursor(0))
+  local root_lang_tree = parsers.get_parser()
+  local owning_lang_tree = root_lang_tree:language_for_range({lnum - 1, col, lnum - 1, col})
+  local result
 
-  return success and node_at_point or nil
+  for _, tree in ipairs(owning_lang_tree:trees()) do
+    result = tree:root():named_descendant_for_range(lnum - 1, col, lnum - 1, col)
+
+    if result then return result end
+  end
 end
 
 local function focus_buf(bufnr)
@@ -88,6 +102,38 @@ local function get_update_time()
   return config and config.updatetime or 25
 end
 
+local function make_entry_toggle(property, update_fn)
+  update_fn = update_fn or function(entry)
+    entry[property] = not entry[property]
+  end
+
+  return function(bufnr)
+    bufnr = bufnr or api.nvim_get_current_buf()
+    update_fn(M._entries[bufnr])
+
+    local current_cursor = vim.api.nvim_win_get_cursor(0)
+    local node_at_cursor = M.get_current_node(bufnr)
+
+    M.update(bufnr)
+
+    -- Restore the cursor to the same node or at least the previous cursor position.
+    local cursor_pos = current_cursor
+    local node_entries = M._entries[bufnr].results
+
+    if node_at_cursor then
+      for lnum, node_entry in ipairs(node_entries) do
+        if node_entry.node:id() == node_at_cursor:id() then
+          cursor_pos = {lnum, cursor_pos[2]}
+        end
+      end
+    end
+
+    -- This could be out of bounds
+    -- TODO(steelsojka): set to end if out of bounds
+    pcall(vim.api.nvim_win_set_cursor, 0, cursor_pos)
+  end
+end
+
 local function setup_buf(for_buf)
   if M._entries[for_buf].display_bufnr then
     return M._entries[for_buf].display_bufnr
@@ -108,7 +154,12 @@ local function setup_buf(for_buf)
   vim.cmd 'augroup END'
 
   api.nvim_buf_set_keymap(buf, 'n', 'o', string.format(':lua require "nvim-treesitter-playground.internal".toggle_query_editor(%d)<CR>', for_buf), { silent = true })
-  api.nvim_buf_set_keymap(buf, 'n', 'i', string.format(':lua require "nvim-treesitter-playground.internal".toggle_highlights(%d)<CR>', for_buf), { silent = true })
+  api.nvim_buf_set_keymap(buf, 'n', 'i', string.format(':lua require "nvim-treesitter-playground.internal".toggle_hl_groups(%d)<CR>', for_buf), { silent = true })
+  api.nvim_buf_set_keymap(buf, 'n', 't', string.format(':lua require "nvim-treesitter-playground.internal".toggle_injected_languages(%d)<CR>', for_buf), { silent = true })
+  api.nvim_buf_set_keymap(buf, 'n', 'a', string.format(':lua require "nvim-treesitter-playground.internal".toggle_anonymous_nodes(%d)<CR>', for_buf), { silent = true })
+  api.nvim_buf_set_keymap(buf, 'n', 'I', string.format(':lua require "nvim-treesitter-playground.internal".toggle_language_display(%d)<CR>', for_buf), { silent = true })
+  api.nvim_buf_set_keymap(buf, 'n', 'f', string.format(':lua require "nvim-treesitter-playground.internal".focus_language(%d)<CR>', for_buf), { silent = true })
+  api.nvim_buf_set_keymap(buf, 'n', 'F', string.format(':lua require "nvim-treesitter-playground.internal".unfocus_language(%d)<CR>', for_buf), { silent = true })
   api.nvim_buf_set_keymap(buf, 'n', 'R', string.format(':lua require "nvim-treesitter-playground.internal".update(%d)<CR>', for_buf), { silent = true })
   api.nvim_buf_set_keymap(buf, 'n', '<cr>', string.format(':lua require "nvim-treesitter-playground.internal".goto_node(%d)<CR>', for_buf), { silent = true })
   api.nvim_buf_attach(buf, false, {
@@ -118,6 +169,23 @@ local function setup_buf(for_buf)
   return buf
 end
 
+local function resolve_lang_tree(bufnr)
+  local entry = M._entries[bufnr]
+
+  if entry.focused_language_tree then
+    local root_lang_tree = parsers.get_parser(bufnr)
+    local found
+
+    root_lang_tree:for_each_child(function(lang_tree)
+      if not found and lang_tree == entry.focused_language_tree then
+        found = lang_tree
+      end
+    end)
+
+    if found then return found end
+  end
+end
+
 local function setup_query_editor(bufnr)
   if M._entries[bufnr].query_bufnr then
     return M._entries[bufnr].query_bufnr
@@ -183,7 +251,20 @@ function M.read_saved_query(bufnr)
       :then_(function(data) return fs_close(fd)
         :then_(function() return vim.split(data, '\n') end) end)
     end)
-    :catch(function(err) return {} end)
+    :catch(function() return {} end)
+end
+
+function M.focus_language(bufnr)
+  local node_entry = M.get_current_entry(bufnr)
+
+  if not node_entry then return end
+
+  M.update(bufnr, node_entry.language_tree)
+end
+
+function M.unfocus_language(bufnr)
+  M._entries[bufnr].focused_language_tree = nil
+  M.update(bufnr)
 end
 
 function M.highlight_playground_nodes(bufnr, nodes)
@@ -192,24 +273,25 @@ function M.highlight_playground_nodes(bufnr, nodes)
   local display_buf = entry.display_bufnr
   local lines = {}
   local count = 0
+  local node_map = utils.to_lookup_table(nodes, function(node) return node:id() end)
 
   if not results or not display_buf then return end
 
-  for i, node in ipairs(results.nodes) do
-    if vim.tbl_contains(nodes, node) then
-      table.insert(lines, i)
+  for line, result in ipairs(results) do
+    if node_map[result.node:id()] then
+      table.insert(lines, line)
       count = count + 1
+    end
 
-      if count >= #nodes then
-        break
-      end
+    if count >= #nodes then
+      break
     end
   end
 
   for _, lnum in ipairs(lines) do
-    local lines = api.nvim_buf_get_lines(display_buf, lnum - 1, lnum, false)
+    local buf_lines = api.nvim_buf_get_lines(display_buf, lnum - 1, lnum, false)
 
-    if lines[1] then
+    if buf_lines[1] then
       vim.api.nvim_buf_add_highlight(display_buf, playground_ns, 'TSPlaygroundFocus', lnum - 1, 0, -1)
     end
   end
@@ -239,13 +321,17 @@ end
 
 M._highlight_playground_node_debounced = utils.debounce(M.highlight_playground_node_from_buffer, get_update_time)
 
-function M.get_current_node(bufnr)
+function M.get_current_entry(bufnr)
   local row, _ = unpack(api.nvim_win_get_cursor(0))
   local results = M._entries[bufnr].results
 
-  if not results then return end
+  return results and results[row]
+end
+
+function M.get_current_node(bufnr)
+  local entry = M.get_current_entry(bufnr)
 
-  return results.nodes[row]
+  return entry and entry.node
 end
 
 function M.highlight_node(bufnr)
@@ -289,7 +375,7 @@ end
 
 function M.update_query(bufnr, query_bufnr)
   local query = table.concat(api.nvim_buf_get_lines(query_bufnr, 0, -1, false), '\n')
-  local matches = pl_query.parse(bufnr, query)
+  local matches = pl_query.parse(bufnr, query, M._entries[bufnr].focused_language_tree)
   local capture_by_color = {}
   local index = 1
 
@@ -379,7 +465,7 @@ end
 function M.clear_highlights(bufnr, namespace)
   if not bufnr then return end
 
-  local namespace = namespace or playground_ns
+  namespace = namespace or playground_ns
 
   api.nvim_buf_clear_namespace(bufnr, namespace, 0, -1)
 end
@@ -389,7 +475,8 @@ function M.clear_playground_highlights(bufnr)
 end
 
 function M.toggle_query_editor(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
+  bufnr = bufnr or api.nvim_get_current_buf()
+
   local display_buf = M._entries[bufnr].display_bufnr
   local current_win = api.nvim_get_current_win()
 
@@ -416,7 +503,8 @@ function M.toggle_query_editor(bufnr)
 end
 
 function M.open(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
+  bufnr = bufnr or api.nvim_get_current_buf()
+
   local display_buf = setup_buf(bufnr)
   local current_window = api.nvim_get_current_win()
 
@@ -435,7 +523,8 @@ function M.open(bufnr)
 end
 
 function M.toggle(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
+  bufnr = bufnr or api.nvim_get_current_buf()
+
   local display_buf = M._entries[bufnr].display_bufnr
 
   if display_buf and is_buf_visible(display_buf) then
@@ -446,38 +535,55 @@ function M.toggle(bufnr)
   end
 end
 
-local print_virt_hl = false
-
-function M.toggle_highlights(bufnr)
-  print_virt_hl = not print_virt_hl
-  local bufnr = bufnr or api.nvim_get_current_buf()
-  local display_buf = M._entries[bufnr].display_bufnr
+M.toggle_anonymous_nodes = make_entry_toggle("include_anonymous_nodes")
+M.toggle_injected_languages = make_entry_toggle("suppress_injected_languages")
+M.toggle_hl_groups = make_entry_toggle("include_hl_groups")
+M.toggle_language_display = make_entry_toggle("include_language")
 
-  if print_virt_hl then
-    printer.print_hl_groups(bufnr, display_buf)
-  else
-    printer.remove_hl_groups(display_buf)
-  end
-end
+function M.update(bufnr, lang_tree)
+  bufnr = bufnr or api.nvim_get_current_buf()
+  lang_tree = lang_tree or resolve_lang_tree(bufnr)
 
-function M.update(bufnr)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-  local display_buf = M._entries[bufnr].display_bufnr
+  local entry = M._entries[bufnr]
+  local display_buf = entry.display_bufnr
 
   -- Don't bother updating if the playground isn't shown
   if not display_buf or not is_buf_visible(display_buf) then return end
 
-  local results = printer.print(bufnr)
+  entry.focused_language_tree = lang_tree
+
+  local results = printer.process(bufnr, lang_tree, {
+    include_anonymous_nodes = entry.include_anonymous_nodes,
+    suppress_injected_languages = entry.suppress_injected_languages,
+    include_hl_groups = entry.include_hl_groups
+  })
 
   M._entries[bufnr].results = results
 
-  api.nvim_buf_set_lines(display_buf, 0, -1, false, results.lines)
-  if print_virt_hl then
-    printer.print_hl_groups(bufnr, display_buf)
+  api.nvim_buf_set_lines(display_buf, 0, -1, false, printer.print_entries(results))
+
+  if entry.query_bufnr then
+    M.update_query(bufnr, entry.query_bufnr)
   end
+
+  if entry.include_language then
+    printer.print_language(display_buf, results)
+  else
+    printer.remove_language(display_buf)
+  end
+
+  if entry.include_hl_groups then
+    printer.print_hl_groups(display_buf, results)
+  else
+    printer.remove_hl_groups(display_buf)
+  end
+end
+
+function M.get_entries()
+  return M._entries
 end
 
-function M.attach(bufnr, lang)
+function M.attach(bufnr)
   api.nvim_buf_attach(bufnr, true, {
     on_lines = vim.schedule_wrap(function() M.update(bufnr) end)
   })

lua/nvim-treesitter-playground/printer.lua

@@ -2,92 +2,191 @@ local parsers = require 'nvim-treesitter.parsers'
 local api = vim.api
 
 local M = {}
+local treesitter_namespace = api.nvim_get_namespaces()["treesitter/highlighter"]
+local virt_text_id = api.nvim_create_namespace('TSPlaygroundHlGroups')
+local lang_virt_text_id = api.nvim_create_namespace('TSPlaygroundLangGroups')
 
-local function print_tree(root, results, indent)
-  local results = results or { lines = {}, nodes = {} }
-  local indentation = indent or ""
+local function get_extmarks(bufnr, start, end_)
+  return api.nvim_buf_get_extmarks(bufnr, treesitter_namespace, start, end_, { details = true })
+end
+
+local function get_hl_group_for_node(bufnr, node)
+  local start_row, start_col, end_row, end_col = node:range()
+  local extmarks = get_extmarks(bufnr, {start_row, start_col}, {end_row, end_col})
+  local groups = {}
+
+  if #extmarks > 0 then
+    for _, ext in ipairs(extmarks) do
+      table.insert(groups, ext[4].hl_group)
+    end
+  end
+
+  return groups
+end
+
+local function flatten_node(root, results, level, language_tree, options)
+  level = level or 0
+  results = results or {}
 
   for node, field in root:iter_children() do
-    if node:named() then
-      local line
-      if field then
-        line = string.format("%s%s: %s [%d, %d] - [%d, %d]",
-          indentation,
-          field,
-          node:type(),
-          node:range())
-      else
-        line = string.format("%s%s [%d, %d] - [%d, %d]",
-          indentation,
-          node:type(),
-          node:range())
+    if node:named() or options.include_anonymous_nodes then
+      local node_entry = {
+        level = level,
+        node = node,
+        field = field,
+        language_tree = language_tree,
+        hl_groups = options.include_hl_groups
+          and options.bufnr
+          and get_hl_group_for_node(options.bufnr, node)
+          or {}
+      }
+
+      table.insert(results, node_entry)
+
+      flatten_node(node, results, level + 1, language_tree, options)
+    end
+  end
+
+  return results
+end
+
+local function node_contains(node, range)
+  local start_row, start_col, end_row, end_col = node:range()
+  local start_fits = start_row < range[1] or (start_row == range[1] and start_col <= range[2])
+  local end_fits = end_row > range[3] or (end_row == range[3] and end_col >= range[4])
+
+  return start_fits and end_fits
+end
+
+local function flatten_lang_tree(lang_tree, results, options)
+  results = results or {}
+
+  for _, tree in ipairs(lang_tree:trees()) do
+    local root = tree:root()
+    local head_entry = nil
+    local head_entry_index = nil
+
+    for i, node_entry in ipairs(results) do
+      local is_contained = node_contains(node_entry.node, {root:range()})
+
+      if is_contained then
+        if not head_entry then
+          head_entry = node_entry
+          head_entry_index = i
+        else
+          if node_entry.level >= head_entry.level then
+            head_entry = node_entry
+            head_entry_index = i
+          else
+            -- If entry contains the root tree but is less specific, then we
+            -- can exit the loop
+            break
+          end
+        end
       end
+    end
+
+    local insert_index = head_entry_index and head_entry_index or #results
+    local level = head_entry and head_entry.level + 1 or nil
 
-      table.insert(results.lines, line)
-      table.insert(results.nodes, node)
+    local flattened_root = flatten_node(root, nil, level, lang_tree, options)
+    local i = insert_index + 1
 
-      print_tree(node, results, indentation .. "  ")
+    -- Insert new items into the table at the correct positions
+    for _, entry in ipairs(flattened_root) do
+      table.insert(results, i, entry)
+      i = i + 1
+    end
+  end
+
+  if not options.suppress_injected_languages then
+    for _, child in pairs(lang_tree:children()) do
+      flatten_lang_tree(child, results, options)
     end
   end
 
   return results
 end
 
-function M.print(bufnr, lang)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-  local parser = parsers.get_parser(bufnr, lang)
+function M.process(bufnr, lang_tree, options)
+  bufnr = bufnr or api.nvim_get_current_buf()
+  options = options or {}
+  lang_tree = lang_tree or parsers.get_parser(bufnr)
+  options.bufnr = options.bufnr or bufnr
 
-  if not parser then return end
+  if not lang_tree then return {} end
 
-  return print_tree(parser:parse()[1]:root())
+  return flatten_lang_tree(lang_tree, nil, options)
 end
 
-local treesitter_namespace = api.nvim_get_namespaces().treesitter_hl
-local virt_text_id = api.nvim_create_namespace('TSPlaygroundHlGroups')
+function M.print_entry(node_entry)
+  local line
+  local indent = string.rep("  ", node_entry.level)
+  local node = node_entry.node
+  local field = node_entry.field
+  local node_name = node:type()
 
-local function get_extmarks(bufnr, start, end_)
-  return api.nvim_buf_get_extmarks(bufnr, treesitter_namespace, start, end_, { details = true })
+  if not node:named() then
+    node_name = string.format([["%s"]], node_name)
+    node_name = string.gsub(node_name, "\n", "\\n")
+  end
+
+  if field then
+    line = string.format("%s%s: %s [%d, %d] - [%d, %d]",
+      indent,
+      field,
+      node_name,
+      node:range())
+  else
+    line = string.format("%s%s [%d, %d] - [%d, %d]",
+      indent,
+      node_name,
+      node:range())
+  end
+
+  return line
 end
 
-function M.print_hl_groups(bufnr, display_bufnr, lang)
-  local bufnr = bufnr or api.nvim_get_current_buf()
-  local parser = parsers.get_parser(bufnr, lang)
-
-  if not parser then return end
-
-  local i = 0
-  local function iter(root)
-    for node, _ in root:iter_children() do
-      if node:named() then
-        local start_row, start_col, end_row, end_col = node:range()
-        local extmarks = get_extmarks(bufnr, {start_row, start_col}, {end_row, end_col})
-
-        if #extmarks > 0 then
-          local groups = {}
-
-          for idx, ext in ipairs(extmarks) do
-            local infos = ext[4]
-            local hl_group = infos.hl_group
-            local str = hl_group.." / "
-            if idx == #extmarks then
-              str = str:sub(0, -3)
-            end
-            table.insert(groups, {str, hl_group})
-          end
-          api.nvim_buf_set_virtual_text(display_bufnr, virt_text_id, i, groups, {})
-        end
+function M.print_entries(node_entries)
+  local results = {}
+
+  for _, entry in ipairs(node_entries) do
+    table.insert(results, M.print_entry(entry))
+  end
+
+  return results
+end
+
+function M.print_hl_groups(bufnr, node_entries)
+  for i, node_entry in ipairs(node_entries) do
+    local groups = {}
+
+    for j, hl_group in ipairs(node_entry.hl_groups) do
+      local str = hl_group .. " / "
 
-        i = i + 1
-        iter(node)
+      if j == #hl_group then
+        str = string.sub(str, 0, -3)
       end
+
+      table.insert(groups, {str, hl_group})
     end
+
+    api.nvim_buf_set_virtual_text(bufnr, virt_text_id, i, groups, {})
   end
+end
 
-  iter(parser:parse()[1]:root())
+function M.print_language(bufnr, node_entries)
+  for i, node_entry in ipairs(node_entries) do
+    api.nvim_buf_set_virtual_text(bufnr, lang_virt_text_id, i - 1, {{node_entry.language_tree:lang()}}, {})
+  end
 end
 
 function M.remove_hl_groups(bufnr)
   api.nvim_buf_clear_namespace(bufnr, virt_text_id, 0, -1)
 end
 
+function M.remove_language(bufnr)
+  api.nvim_buf_clear_namespace(bufnr, lang_virt_text_id, 0, -1)
+end
+
 return M

lua/nvim-treesitter-playground/query.lua

@@ -1,25 +1,29 @@
 local ts_query = require 'nvim-treesitter.query'
 local parsers = require 'nvim-treesitter.parsers'
 local locals = require 'nvim-treesitter.locals'
-local api = vim.api
 
 local M = {}
 
-function M.parse(bufnr, query)
-  local lang = api.nvim_buf_get_option(bufnr, 'ft')
-  local success, parsed_query = pcall(function() return vim.treesitter.parse_query(lang, query) end)
+function M.parse(bufnr, query, lang_tree)
+  lang_tree = lang_tree or parsers.get_parser(bufnr)
+
+  local success, parsed_query = pcall(function()
+    return vim.treesitter.parse_query(lang_tree:lang(), query)
+  end)
 
   if not success then return {} end
 
-  local parser = parsers.get_parser(bufnr, lang)
-  local root = parser:parse()[1]:root()
-  local start_row, _, end_row, _ = root:range()
   local results = {}
 
-  for match in ts_query.iter_prepared_matches(parsed_query, root, bufnr, start_row, end_row) do
-    locals.recurse_local_nodes(match, function(_, node, path)
-      table.insert(results, { node = node, tag = path })
-    end)
+  for _, tree in ipairs(lang_tree:trees()) do
+    local root = tree:root()
+    local start_row, _, end_row, _ = root:range()
+
+    for match in ts_query.iter_prepared_matches(parsed_query, root, bufnr, start_row, end_row) do
+      locals.recurse_local_nodes(match, function(_, node, path)
+        table.insert(results, { node = node, tag = path })
+      end)
+    end
   end
 
   return results

lua/nvim-treesitter-playground/query_linter.lua

@@ -39,35 +39,41 @@ local function table_contains(predicate, table)
 end
 
 local function query_lang_from_playground_buf(buf)
-  for lang_buf, entry in pairs(playground_module._entries or {}) do
+  for lang_buf, entry in pairs(playground_module.get_entries() or {}) do
     if entry.query_bufnr == buf then
+      if entry.focused_language_tree then
+        return entry.focused_language_tree:lang()
+      end
+
       return parsers.get_buf_lang(lang_buf)
     end
   end
 end
 
-function M.lint(buf)
-  buf = buf or api.nvim_get_current_buf()
-  M.clear_virtual_text(buf)
-  M.lints[buf] = {}
+function M.lint(query_buf)
+  query_buf = query_buf or api.nvim_get_current_buf()
+  M.clear_virtual_text(query_buf)
+  M.lints[query_buf] = {}
 
-  local filename = api.nvim_buf_get_name(buf)
+  local filename = api.nvim_buf_get_name(query_buf)
   local ok, query_lang = pcall(vim.fn.fnamemodify, filename, ":p:h:t")
   query_lang = filename ~= "" and query_lang
-  local query_lang = ok and query_lang
+  query_lang = ok and query_lang
   if not query_lang then
-    query_lang = query_lang_from_playground_buf(buf)
+    query_lang = query_lang_from_playground_buf(query_buf)
   end
+
   local ok, parser_info = pcall(vim.treesitter.inspect_language, query_lang)
-  local parser_info = ok and parser_info
 
-  local matches = queries.get_matches(buf, "query-linter-queries")
+  parser_info = ok and parser_info
+
+  local matches = queries.get_matches(query_buf, "query-linter-queries")
 
   for _, m in pairs(matches) do
     local error_node = utils.get_at_path(m, "error.node")
 
     if error_node then
-      lint_node(error_node, buf, "Syntax Error")
+      lint_node(error_node, query_buf, "Syntax Error")
     end
 
     local toplevel_node = utils.get_at_path(m, "toplevel-query.node")
@@ -75,7 +81,7 @@ function M.lint(buf)
       local query_text = table.concat(ts_utils.get_node_text(toplevel_node), "\n")
       local ok, err = pcall(vim.treesitter.parse_query, query_lang, query_text)
       if not ok then
-        lint_node(toplevel_node, buf, "Invalid Query", err)
+        lint_node(toplevel_node, query_buf, "Invalid Query", err)
       end
     end
 
@@ -102,7 +108,7 @@ function M.lint(buf)
           )
 
         if not found then
-          lint_node(node, buf, "Invalid Node Type")
+          lint_node(node, query_buf, "Invalid Node Type")
         end
       end
 
@@ -112,12 +118,12 @@ function M.lint(buf)
         local field_name = ts_utils.get_node_text(field_node)[1]
         local found = vim.tbl_contains(parser_info.fields, field_name)
         if not found then
-          lint_node(field_node, buf, "Invalid Field")
+          lint_node(field_node, query_buf, "Invalid Field")
         end
       end
     end
   end
-  return M.lints[buf]
+  return M.lints[query_buf]
 end
 
 function M.clear_virtual_text(buf)

lua/nvim-treesitter-playground/utils.lua

@@ -24,4 +24,20 @@ function M.for_each_buf_window(bufnr, fn)
   end
 end
 
+function M.to_lookup_table(list, key_mapper)
+  local result = {}
+
+  for i, v in ipairs(list) do
+    local key = v
+
+    if key_mapper then
+      key = key_mapper(v, i)
+    end
+
+    result[key] = v
+  end
+
+  return result
+end
+
 return M

plugin/nvim-treesitter-playground.vim

@@ -4,5 +4,6 @@ EOF
 
 highlight default link TSPlaygroundFocus Visual
 highlight default link TSQueryLinterError Error
+highlight default link TSPlaygroundLang String
 
 command! TSPlaygroundToggle lua require "nvim-treesitter-playground.internal".toggle()

syntax/tsplayground.vim

@@ -6,10 +6,12 @@ syn match nodeType "[a-zA-Z_]\+"
 syn match nodeNumber "\d\+"
 syn match nodeOp "[,\-\)]\+"
 syn match nodeTag "\k\+:"
+syn match nodeAnonymous "\".\+\""
 
 hi def link nodeType Identifier
 hi def link nodeNumber Number
 hi def link nodeOp Operator
 hi def link nodeTag Tag
+hi def link nodeAnonymous String
 
 let b:current_syntax = 'tsplayground'