progress on query editor

lua/nvim-treesitter-playground/internal.lua

@@ -2,34 +2,58 @@ local configs = require 'nvim-treesitter.configs'
 local ts_utils = require 'nvim-treesitter.ts_utils'
 local printer = require 'nvim-treesitter-playground.printer'
 local utils = require 'nvim-treesitter-playground.utils'
+local ts_query = require 'nvim-treesitter.query'
+local pl_query = require 'nvim-treesitter-playground.query'
 local api = vim.api
 
 local M = {}
 
 M._displays_by_buf = {}
 M._results_by_buf = {}
+M._queries_by_buf = {}
 
 local playground_ns = api.nvim_create_namespace('nvim-treesitter-playground')
+local query_hl_ns = api.nvim_create_namespace('nvim-treesitter-playground-query')
 
-local function clear_entry(bufnr)
-  if M._displays_by_buf[bufnr] then
-    close_buf_windows(M._displays_by_buf[bufnr])
+local function focus_buf(bufnr)
+  if not bufnr then return end
 
-    if api.nvim_buf_is_loaded(M._displays_by_buf[bufnr]) then
-      vim.cmd(string.format("bw! %d", M._displays_by_buf[bufnr]))
-    end
-  end
+  local windows = vim.fn.win_findbuf(bufnr)
 
-  M._displays_by_buf[bufnr] = nil
-  M._results_by_buf[bufnr] = nil
+  if windows[1] then
+    api.nvim_set_current_win(windows[1])
+  end
 end
 
 local function close_buf_windows(bufnr)
+  if not bufnr then return end
+
   utils.for_each_buf_window(bufnr, function(window)
     api.nvim_win_close(window, true)
   end)
 end
 
+local function close_buf(bufnr)
+  if not bufnr then return end
+
+  close_buf_windows(bufnr)
+
+  if api.nvim_buf_is_loaded(bufnr) then
+    vim.cmd(string.format("bw! %d", bufnr))
+  end
+end
+
+local function clear_entry(bufnr)
+  local display_buf = M._displays_by_buf[bufnr]
+  local query_buf = M._queries_by_buf[bufnr]
+
+  close_buf(display_buf)
+  close_buf(query_buf)
+  M._displays_by_buf[bufnr] = nil
+  M._results_by_buf[bufnr] = nil
+  M._queries_by_buf[bufnr] = nil
+end
+
 local function is_buf_visible(bufnr)
   local windows = vim.fn.win_findbuf(bufnr)
 
@@ -55,6 +79,8 @@ local function setup_buf(for_buf)
   vim.cmd(string.format([[autocmd BufWinEnter <buffer=%d> lua require'nvim-treesitter-playground.internal'.update(%d)]], 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', 'R', string.format(':lua require "nvim-treesitter-playground.internal".update(%d)<CR>', for_buf), { silent = true })
   api.nvim_buf_attach(buf, false, {
     on_detach = function() clear_entry(for_buf) end
   })
@@ -62,6 +88,28 @@ local function setup_buf(for_buf)
   return buf
 end
 
+local function setup_query_editor(bufnr)
+  if M._queries_by_buf[bufnr] then
+    return M._queries_by_buf[bufnr]
+  end
+
+  local buf = api.nvim_create_buf(false, false)
+
+  api.nvim_buf_set_option(buf, 'buftype', 'nofile')
+  api.nvim_buf_set_option(buf, 'swapfile', false)
+  api.nvim_buf_set_option(buf, 'buflisted', false)
+  api.nvim_buf_set_option(buf, 'filetype', 'query')
+
+  M._queries_by_buf[bufnr] = buf
+
+  api.nvim_buf_set_keymap(buf, 'n', 'R', string.format(':lua require "nvim-treesitter-playground.internal".update_query(%d, %d)<CR>', bufnr, buf), { silent = true })
+  api.nvim_buf_attach(buf, false, {
+    on_lines = utils.debounce(function() M.update_query(bufnr, buf) end, 1000)
+  })
+
+  return buf
+end
+
 M.highlight_playground_node = utils.debounce(function(bufnr)
   M.clear_playground_highlights(bufnr)
 
@@ -121,8 +169,47 @@ function M.highlight_node(bufnr)
   end)
 end
 
-function M.clear_highlights(bufnr)
-  api.nvim_buf_clear_namespace(bufnr, playground_ns, 0, -1)
+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 capture_by_color = {}
+  local index = 1
+
+  M.clear_highlights(query_bufnr, query_hl_ns)
+  M.clear_highlights(bufnr, query_hl_ns)
+
+  print(bufnr, query_bufnr)
+
+  for capture_match in ts_query.iter_group_results(query_bufnr, 'captures') do
+    local capture = ts_utils.get_node_text(capture_match.capture.name.node)[1]
+
+    if not capture_by_color[capture] then
+      capture_by_color[capture] = 'TSPlaygroundCapture' .. index
+      index = index + 1
+    end
+
+    ts_utils.highlight_node(capture_match.capture.def.node, query_bufnr, query_hl_ns, capture_by_color[capture])
+  end
+
+  local node_highlights = {}
+
+  for _, match in ipairs(matches) do
+    local hl_group = capture_by_color[match.tag]
+
+    if hl_group then
+      table.insert(node_highlights, { match.node, hl_group })
+    end
+  end
+
+  for _, entry in ipairs(node_highlights) do
+    ts_utils.highlight_node(entry[1], bufnr, query_hl_ns, entry[2])
+  end
+end
+
+function M.clear_highlights(bufnr, namespace)
+  local namespace = namespace or playground_ns
+
+  api.nvim_buf_clear_namespace(bufnr, namespace, 0, -1)
 end
 
 function M.clear_playground_highlights(bufnr)
@@ -133,6 +220,33 @@ function M.clear_playground_highlights(bufnr)
   end
 end
 
+function M.toggle_query_editor(bufnr)
+  local bufnr = bufnr or api.nvim_get_current_buf()
+  local display_buf = M._displays_by_buf[bufnr]
+  local current_win = api.nvim_get_current_win()
+
+  if not display_buf then
+    display_buf = M.open(bufnr)
+  end
+
+  local query_buf = setup_query_editor(bufnr)
+
+  if is_buf_visible(query_buf) then
+    close_buf_windows(query_buf)
+  else
+    M._queries_by_buf[bufnr] = query_buf
+
+    focus_buf(display_buf)
+    vim.cmd "split"
+    vim.cmd(string.format("buffer %d", query_buf))
+
+    api.nvim_win_set_option(0, 'spell', false)
+    api.nvim_win_set_option(0, 'number', true)
+
+    api.nvim_set_current_win(current_win)
+  end
+end
+
 function M.open(bufnr)
   local bufnr = bufnr or api.nvim_get_current_buf()
   local display_buf = setup_buf(bufnr)
@@ -148,6 +262,8 @@ function M.open(bufnr)
   api.nvim_win_set_option(0, 'cursorline', false)
 
   api.nvim_set_current_win(current_window)
+
+  return display_buf
 end
 
 function M.toggle(bufnr)
@@ -155,6 +271,7 @@ function M.toggle(bufnr)
   local display_buf = M._displays_by_buf[bufnr]
 
   if display_buf and is_buf_visible(display_buf) then
+    close_buf_windows(M._queries_by_buf[display_buf])
     close_buf_windows(display_buf)
   else
     M.open(bufnr)
@@ -180,8 +297,8 @@ function M.attach(bufnr, lang)
   local config = configs.get_module 'playground'
 
   api.nvim_buf_attach(bufnr, true, {
-    on_lines = vim.schedule_wrap(function() M.update(bufnr) end),
-    on_detach = function() clear_entry(bufnr) end
+    on_lines = vim.schedule_wrap(function() M.update(bufnr) end)
+    -- on_detach = function() clear_entry(bufnr) end
   })
 
   vim.cmd(string.format('augroup TreesitterPlayground_%d', bufnr))

lua/nvim-treesitter-playground/query.lua

@@ -0,0 +1,25 @@
+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 parsed_query = vim.treesitter.parse_query(lang, query)
+  local parser = parsers.get_parser(bufnr, lang)
+  local root = parser:parse():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(def, node, path)
+      table.insert(results, { node = node, tag = path })
+    end)
+  end
+
+  return results
+end
+
+return M

plugin/nvim-treesitter-playground.vim

@@ -3,5 +3,9 @@ require "nvim-treesitter-playground".init()
 EOF
 
 highlight default link TSPlaygroundFocus Visual
+highlight TSPlaygroundCapture1 ctermfg=Blue guifg=Blue
+highlight TSPlaygroundCapture2 ctermfg=LightBlue guifg=LightBlue
+highlight TSPlaygroundCapture3 ctermfg=Green guifg=Green
+highlight TSPlaygroundCapture4 ctermfg=LightGreen guifg=LightGreen
 
 command! TSPlaygroundToggle lua require "nvim-treesitter-playground.internal".toggle()

queries/query/captures.scm

@@ -0,0 +1 @@
+(capture name: (identifier) @capture.name) @capture.def