feat(treesitter): playground improvements

- Render node ranges as virtual text
- Set filettype=query. The virtual text is to avoid parsing errors.
- Make sure highlights text is always in view.

2 files changed, 32 insertions, 17 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 96b1e24ba9..a15d3ec23c 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -411,6 +411,7 @@ function M.show_tree(opts)
   vim.bo[b].buflisted = false
   vim.bo[b].buftype = 'nofile'
   vim.bo[b].bufhidden = 'wipe'
+  vim.bo[b].filetype = 'query'
 
   local title = opts.title
   if not title then
@@ -425,9 +426,6 @@ function M.show_tree(opts)
 
   pg:draw(b)
 
-  vim.fn.matchadd('Comment', '\\[[0-9:-]\\+\\]')
-  vim.fn.matchadd('String', '".*"')
-
   a.nvim_buf_clear_namespace(buf, pg.ns, 0, -1)
   a.nvim_buf_set_keymap(b, 'n', '<CR>', '', {
     desc = 'Jump to the node under the cursor in the source buffer',
@@ -467,6 +465,15 @@ function M.show_tree(opts)
         end_col = math.max(0, pos.end_col),
         hl_group = 'Visual',
       })
+
+      local topline, botline = vim.fn.line('w0', win), vim.fn.line('w$', win)
+
+      -- Move the cursor if highlighted range is completely out of view
+      if pos.lnum < topline and pos.end_lnum < topline then
+        a.nvim_win_set_cursor(win, { pos.end_lnum + 1, 0 })
+      elseif pos.lnum > botline and pos.end_lnum > botline then
+        a.nvim_win_set_cursor(win, { pos.lnum + 1, 0 })
+      end
     end,
   })
 
diff --git a/runtime/lua/vim/treesitter/playground.lua b/runtime/lua/vim/treesitter/playground.lua
index be7764e6f0..001bc2d5bf 100644
--- a/runtime/lua/vim/treesitter/playground.lua
+++ b/runtime/lua/vim/treesitter/playground.lua
@@ -136,6 +136,8 @@ function TSPlayground:new(bufnr, lang)
   return t
 end
 
+local decor_ns = api.nvim_create_namespace('ts.playground')
+
 --- Write the contents of this Playground into {bufnr}.
 ---
 ---@param bufnr number Buffer number to write into.
@@ -144,22 +146,28 @@ function TSPlayground:draw(bufnr)
   vim.bo[bufnr].modifiable = true
   local lines = {} ---@type string[]
   for _, item in self:iter() do
-    lines[#lines + 1] = table.concat({
-      string.rep(' ', item.depth),
-      item.text,
-      item.lnum == item.end_lnum
-          and string.format(' [%d:%d-%d]', item.lnum + 1, item.col + 1, item.end_col)
-        or string.format(
-          ' [%d:%d-%d:%d]',
-          item.lnum + 1,
-          item.col + 1,
-          item.end_lnum + 1,
-          item.end_col
-        ),
-      self.opts.lang and string.format(' %s', item.lang) or '',
-    })
+    lines[#lines + 1] = string.rep(' ', item.depth) .. item.text
   end
   api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
+
+  api.nvim_buf_clear_namespace(bufnr, decor_ns, 0, -1)
+
+  for i, item in self:iter() do
+    local range_str
+    if item.lnum == item.end_lnum then
+      range_str = string.format('[%d:%d-%d]', item.lnum + 1, item.col + 1, item.end_col)
+    else
+      range_str =
+        string.format('[%d:%d-%d:%d]', item.lnum + 1, item.col + 1, item.end_lnum + 1, item.end_col)
+    end
+
+    local lang_str = self.opts.lang and string.format(' %s', item.lang) or ''
+
+    api.nvim_buf_set_extmark(bufnr, decor_ns, i - 1, 0, {
+      virt_text = { { range_str, 'Comment' }, { lang_str, 'Title' } },
+    })
+  end
+
   vim.bo[bufnr].modifiable = false
 end