Merge pull request #22191 from lewis6991/feat/playground_imp

feat(treesitter): playground improvements

3 files changed, 39 insertions, 22 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 96b1e24ba9..4127198576 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -96,11 +96,13 @@ function M.get_parser(bufnr, lang, opts)
   if bufnr == nil or bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
-  if lang == nil then
-    lang = a.nvim_buf_get_option(bufnr, 'filetype')
-  end
 
-  if parsers[bufnr] == nil or parsers[bufnr]:lang() ~= lang then
+  if parsers[bufnr] == nil then
+    lang = lang or a.nvim_buf_get_option(bufnr, 'filetype')
+    parsers[bufnr] = M._create_parser(bufnr, lang, opts)
+  elseif lang and parsers[bufnr]:lang() ~= lang then
+    -- Only try to create a new parser if lang is provided
+    -- and it doesn't match the stored parser
     parsers[bufnr] = M._create_parser(bufnr, lang, opts)
   end
 
@@ -411,6 +413,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 +428,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 +467,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
 
diff --git a/test/functional/treesitter/language_spec.lua b/test/functional/treesitter/language_spec.lua
index df45c9b384..51d8eb62f3 100644
--- a/test/functional/treesitter/language_spec.lua
+++ b/test/functional/treesitter/language_spec.lua
@@ -82,7 +82,7 @@ describe('treesitter language API', function()
     command("set filetype=borklang")
     -- Should throw an error when filetype changes to borklang
     eq(".../language.lua:0: no parser for 'borklang' language, see :help treesitter-parsers",
-       pcall_err(exec_lua, "new_parser = vim.treesitter.get_parser(0)"))
+       pcall_err(exec_lua, "new_parser = vim.treesitter.get_parser(0, 'borklang')"))
   end)
 
   it('retrieve the tree given a range', function ()