wip

lua/ssr/search.lua

@@ -113,7 +113,7 @@ local function build_sexpr(node, source)
 end
 
 ---@class ssr.Searcher
----@field query vim.treesitter.Query
+---@field lang_queries table<string, vim.treesitter.Query>
 ---@field captures table<string, integer>
 ---@field rough_regex string
 local Searcher = {}
@@ -128,13 +128,13 @@ function Searcher.new(lang, pattern, parse_context)
     return
   end
   local sexpr, captures, rough_regex = build_sexpr(node, source)
-  local parse_query = ts.query.parse or ts.parse_query
-  local query = parse_query(lang, sexpr)
+  local query = ts.query.parse(lang, sexpr)
   return setmetatable({
     lang = lang,
-    query = query,
+    lang_queries = { lang = query },
     captures = captures,
     rough_regex = rough_regex,
+    pattern = pattern,
   }, { __index = Searcher })
 end
 
@@ -149,10 +149,22 @@ function Searcher:search(file)
   ---@type ssr.Match[]
   local matches = {}
   file.tree:for_each_tree(function(tree, lang_tree) -- must called :parse(true)
-    if lang_tree:lang() ~= self.lang then
+    local query = self.lang_queries[lang_tree:lang()]
+    if query == vim.NIL then
       return
     end
-    for _, nodes in self.query:iter_matches(tree:root(), file.source, 0, -1) do
+    if not query then
+      local ok = pcall(function()
+        local node, source = require("ssr.parse_context").empty(lang_tree:lang()):parse(self.pattern)
+        query = ts.query.parse(lang_tree:lang(), build_sexpr(node, source))
+        self.lang_queries[lang_tree:lang()] = query
+      end)
+      if not ok then
+        self.lang_queries[lang_tree:lang()] = vim.NIL
+        return
+      end
+    end
+    for _, nodes in query:iter_matches(tree:root(), file.source, 0, -1) do
       local range = Range.from_node(nodes[#nodes])
       local captures = {}
       for var, idx in pairs(self.captures) do