fix(playground): show highlight groups (#71)

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

@@ -1,65 +1,29 @@
+local utils = require "nvim-treesitter-playground.utils"
 local highlighter = require "vim.treesitter.highlighter"
-local ts_utils = require "nvim-treesitter.ts_utils"
 
 local M = {}
 
 function M.get_treesitter_hl()
-  local buf = vim.api.nvim_get_current_buf()
+  local bufnr = vim.api.nvim_get_current_buf()
   local row, col = unpack(vim.api.nvim_win_get_cursor(0))
   row = row - 1
 
-  local self = highlighter.active[buf]
-
-  if not self then
-    return {}
-  end
-
-  local matches = {}
-
-  self.tree:for_each_tree(function(tstree, tree)
-    if not tstree then
-      return
+  local results = utils.get_hl_groups_at_position(bufnr, row, col)
+  local highlights = {}
+  for _, hl in pairs(results) do
+    local line = "* **@" .. hl.capture .. "**"
+    if hl.specific then
+      line = line .. " -> **" .. hl.specific .. "**"
     end
-
-    local root = tstree:root()
-    local root_start_row, _, root_end_row, _ = root:range()
-
-    -- Only worry about trees within the line range
-    if root_start_row > row or root_end_row < row then
-      return
+    if hl.general then
+      line = line .. " -> **" .. hl.general .. "**"
     end
-
-    local query = self:get_query(tree:lang())
-
-    -- Some injected languages may not have highlight queries.
-    if not query:query() then
-      return
+    if hl.priority then
+      line = line .. "(" .. hl.priority .. ")"
     end
-
-    local iter = query:query():iter_captures(root, self.bufnr, row, row + 1)
-
-    for capture, node, metadata in iter do
-      if ts_utils.is_in_node_range(node, row, col) then
-        local c = query._query.captures[capture] -- name of the capture in the query
-        if c ~= nil then
-          local general_hl, is_vim_hl = query:_get_hl_from_capture(capture)
-          local local_hl = not is_vim_hl and (tree:lang() .. general_hl)
-          local line = "* **@" .. c .. "**"
-          if local_hl then
-            line = line .. " -> **" .. local_hl .. "**"
-          end
-          if general_hl then
-            line = line .. " -> **" .. general_hl .. "**"
-          end
-          if metadata.priority then
-            line = line .. " *(priority " .. metadata.priority .. ")*"
-          end
-          table.insert(matches, line)
-        end
-      end
-    end
-  end, true)
-  return matches
+    table.insert(highlights, line)
+  end
+  return highlights
 end
 
 function M.get_syntax_hl()

lua/nvim-treesitter-playground/printer.lua

@@ -3,22 +3,17 @@ local utils = require "nvim-treesitter-playground.utils"
 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 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 start_row, start_col, _, _ = node:range()
+  local hlgroups = utils.get_hl_groups_at_position(bufnr, start_row, start_col)
   local groups = {}
 
-  if #extmarks > 0 then
-    for _, ext in ipairs(extmarks) do
-      table.insert(groups, ext[4].hl_group)
+  if #hlgroups > 0 then
+    for _, hl in pairs(hlgroups) do
+      table.insert(groups, hl.general)
     end
   end
 

lua/nvim-treesitter-playground/utils.lua

@@ -1,4 +1,6 @@
 local api = vim.api
+local ts_utils = require "nvim-treesitter.ts_utils"
+local highlighter = require "vim.treesitter.highlighter"
 
 local M = {}
 
@@ -26,6 +28,56 @@ function M.debounce(fn, debounce_time)
   end
 end
 
+function M.get_hl_groups_at_position(bufnr, row, col)
+  local buf_highlighter = highlighter.active[bufnr]
+
+  if not buf_highlighter then
+    return {}
+  end
+
+  local matches = {}
+
+  buf_highlighter.tree:for_each_tree(function(tstree, tree)
+    if not tstree then
+      return
+    end
+
+    local root = tstree:root()
+    local root_start_row, _, root_end_row, _ = root:range()
+
+    -- Only worry about trees within the line range
+    if root_start_row > row or root_end_row < row then
+      return
+    end
+
+    local query = buf_highlighter:get_query(tree:lang())
+
+    -- Some injected languages may not have highlight queries.
+    if not query:query() then
+      return
+    end
+
+    local iter = query:query():iter_captures(root, buf_highlighter.bufnr, row, row + 1)
+
+    for capture, node, metadata in iter do
+      local hl = query.hl_cache[capture]
+
+      if hl and ts_utils.is_in_node_range(node, row, col) then
+        local c = query._query.captures[capture] -- name of the capture in the query
+        if c ~= nil then
+          local general_hl, is_vim_hl = query:_get_hl_from_capture(capture)
+          local local_hl = not is_vim_hl and (tree:lang() .. general_hl)
+          table.insert(
+            matches,
+            { capture = c, specific = local_hl, general = general_hl, priority = metadata.priority }
+          )
+        end
+      end
+    end
+  end, true)
+  return matches
+end
+
 function M.for_each_buf_window(bufnr, fn)
   if not api.nvim_buf_is_loaded(bufnr) then
     return