From 2afb7fc9113d1135a52800a2f1c03f747c132dad Mon Sep 17 00:00:00 2001
From: Luis Calle <53507599+TheLeoP@users.noreply.github.com>
Date: Wed, 31 Dec 2025 17:39:33 -0500
Subject: [PATCH] perf: improve general treesitter performance and markdown
 particular performance (#423)

* perf(treesitter): query only the necessary trees instead of all of them

* perf(markdown): remove expensive latex query
---
 after/queries/markdown/matchup.scm        |  8 +++---
 after/queries/markdown_inline/matchup.scm | 31 ++++++++-------------
 lua/treesitter-matchup/internal.lua       | 34 ++++++++++++-----------
 3 files changed, 34 insertions(+), 39 deletions(-)

diff --git a/after/queries/markdown/matchup.scm b/after/queries/markdown/matchup.scm
index 54af626..c3bac2a 100644
--- a/after/queries/markdown/matchup.scm
+++ b/after/queries/markdown/matchup.scm
@@ -1,6 +1,6 @@
-;; ====================
-;; fenced code block ```
-;; ====================
+; ====================
+; fenced code block ```
+; ====================
 (fenced_code_block
   (fenced_code_block_delimiter) @open.code_fence
-  (fenced_code_block_delimiter) @close.code_fence) @scope.code_fence
\ No newline at end of file
+  (fenced_code_block_delimiter) @close.code_fence) @scope.code_fence
diff --git a/after/queries/markdown_inline/matchup.scm b/after/queries/markdown_inline/matchup.scm
index 7973405..76691e0 100644
--- a/after/queries/markdown_inline/matchup.scm
+++ b/after/queries/markdown_inline/matchup.scm
@@ -1,36 +1,29 @@
-;; ===========
-;; inline code `
-;; ===========
+; ===========
+; inline code `
+; ===========
 (code_span
   (code_span_delimiter) @open.code
   (code_span_delimiter) @close.code) @scope.code
 
-;; ===========================
-;; emphasis (italic, *...* / _..._)
-;; ===========================
+; ===========================
+; emphasis (italic, *...* / _..._)
+; ===========================
 (emphasis
   (emphasis_delimiter) @open.emphasis
   (emphasis_delimiter) @close.emphasis) @scope.emphasis
 
-;; ===========================
-;; strong emphasis (bold, **...** / __...__)
-;; ===========================
+; ===========================
+; strong emphasis (bold, **...** / __...__)
+; ===========================
 (strong_emphasis
   (emphasis_delimiter) @open.strong
   (emphasis_delimiter) @open.strong
   (emphasis_delimiter) @close.strong
   (emphasis_delimiter) @close.strong) @scope.strong
 
-;; ===========================
-;; strikethrough (~~...~~)
-;; ===========================
+; ===========================
+; strikethrough (~~...~~)
+; ===========================
 (strikethrough
   (emphasis_delimiter) @open.strikethrough
   (emphasis_delimiter) @close.strikethrough) @scope.strikethrough
-
-;; ===========================
-;; LaTeX math ($...$ / $$...$$)
-;; ===========================
-(latex_block
-  (latex_span_delimiter) @open.math
-  (latex_span_delimiter) @close.math) @scope.math
diff --git a/lua/treesitter-matchup/internal.lua b/lua/treesitter-matchup/internal.lua
index aa860f2..551418c 100644
--- a/lua/treesitter-matchup/internal.lua
+++ b/lua/treesitter-matchup/internal.lua
@@ -78,8 +78,7 @@ end
 ---@param lang string
 ---@return matchup.treesitter.Match[]
 local get_memoized_matches = memoize(function(bufnr, root, lang)
-  local query_name = 'matchup'
-  local query = ts.query.get(lang, query_name)
+  local query = ts.query.get(lang, 'matchup')
 
   if not query then
     return {}
@@ -137,28 +136,31 @@ end, buf_root_lang_hash)
 ---@param bufnr integer
 ---@return matchup.treesitter.Match[]
 M.get_matches = function(bufnr)
-  local parser = ts.get_parser(bufnr)
+  local lang_tree = ts.get_parser(bufnr)
   local matches = {} ---@type matchup.treesitter.Match[]
 
-  if parser then
+  if lang_tree then
     -- NOTE: assummes that we are always parsing the current window. May cause
     -- issues if that's not always the case
-    local win = api.nvim_get_current_win()
-    local cur_row = unpack(api.nvim_win_get_cursor(win))
+    local cursor = vim.api.nvim_win_get_cursor(0)
     local stopline = vim.g.matchup_treesitter_stopline ---@type integer
-    local start_row = math.max(cur_row - stopline, 0)
-    local end_row = math.min(cur_row + stopline, api.nvim_buf_line_count(bufnr))
+    local start_row = math.max(cursor[1] - stopline, 0)
+    local end_row = math.min(cursor[1] + stopline, api.nvim_buf_line_count(bufnr))
 
-    parser:parse({start_row, end_row})
-    parser:for_each_tree(function(tree, lang_tree)
-      if not tree or lang_tree:lang() == 'comment' then
-        return
+    lang_tree:parse({ start_row, end_row })
+    ---@type vim.treesitter.LanguageTree?
+    local nested_lang_tree = lang_tree:language_for_range({ cursor[1]-1, cursor[2], cursor[1]-1, cursor[2] })
+    while vim.tbl_isempty(matches) and nested_lang_tree ~= nil do
+      local lang = nested_lang_tree:lang()
+      if lang ~= 'comment' then
+        for _, tree in ipairs(nested_lang_tree:trees()) do
+          local group_results = get_memoized_matches(bufnr, tree:root(), lang)
+          vim.list_extend(matches, group_results)
+        end
       end
 
-      local lang = lang_tree:lang()
-      local group_results = get_memoized_matches(bufnr, tree:root(), lang)
-      vim.list_extend(matches, group_results)
-    end)
+      nested_lang_tree = nested_lang_tree:parent()
+    end
   end
 
   return matches