add Range:tsquery() to query TS within a range

lua/u/range.lua

@@ -233,39 +233,22 @@ end
 
 --- @param opts? { contains_cursor?: boolean }
 function Range.from_tsquery_caps(bufnr, query, opts)
-  if bufnr == nil or bufnr == 0 then bufnr = vim.api.nvim_get_current_buf() end
-  local lang = vim.treesitter.language.get_lang(vim.bo[bufnr].filetype)
-  if lang == nil then return end
-  local parser = vim.treesitter.get_parser(bufnr, lang)
-  if parser == nil then return end
-  local tree = parser:parse()[1]
-  if tree == nil then return end
-
   opts = opts or { contains_cursor = true }
-  local cursor = Pos.from_pos '.'
 
-  local root = tree:root()
-  local root_end = root:end_()
-  local q = vim.treesitter.query.parse(lang, query)
-  --- @type table<string, u.Range[]>
-  local ranges = {}
-  for id, match, _meta in q:iter_captures(root, bufnr, 0, root_end) do
-    local start_row0, start_col0, stop_row0, stop_col0 = match:range()
-    local range = Range.new(
-      Pos.new(bufnr, start_row0 + 1, start_col0 + 1),
-      Pos.new(bufnr, stop_row0 + 1, stop_col0 + 1),
-      'v'
-    )
-    if range.stop.lnum > vim.api.nvim_buf_line_count(bufnr) then
-      range.stop = range.stop:must_next(-1)
-    end
-    if not opts.contains_cursor or opts.contains_cursor and range:contains(cursor) then
-      local capture_name = q.captures[id]
-      if not ranges[capture_name] then ranges[capture_name] = {} end
-      table.insert(ranges[capture_name], range)
-    end
-  end
-  return ranges
+  local ranges = Range.from_buf_text(bufnr):tsquery(query)
+  if not ranges then return end
+  if not opts.contains_cursor then return ranges end
+
+  local cursor = Pos.from_pos '.'
+  return vim.tbl_map(function(cap_ranges)
+    return vim
+      .iter(cap_ranges)
+      :filter(
+        --- @param r u.Range
+        function(r) return r:contains(cursor) end
+      )
+      :totable()
+  end, ranges)
 end
 
 --- Get range information from the currently selected visual text.
@@ -440,14 +423,46 @@ function Range:set_visual_selection()
   self.stop:save_to_pos '.'
 end
 
---------------------------------------------------------------------------------
--- Range.from_* functions:
---------------------------------------------------------------------------------
-
 --------------------------------------------------------------------------------
 -- Text access/manipulation utilities:
 --------------------------------------------------------------------------------
 
+--- @param query string
+function Range:tsquery(query)
+  local bufnr = self.start.bufnr
+
+  local lang = vim.treesitter.language.get_lang(vim.bo[bufnr].filetype)
+  if lang == nil then return end
+  local parser = vim.treesitter.get_parser(bufnr, lang)
+  if parser == nil then return end
+  local tree = parser:parse()[1]
+  if tree == nil then return end
+
+  local root = tree:root()
+  local q = vim.treesitter.query.parse(lang, query)
+  --- @type table<string, u.Range[]>
+  local ranges = {}
+  for id, match, _meta in
+    q:iter_captures(root, bufnr, self.start.lnum - 1, (self.stop or self.start).lnum)
+  do
+    local start_row0, start_col0, stop_row0, stop_col0 = match:range()
+    local range = Range.new(
+      Pos.new(bufnr, start_row0 + 1, start_col0 + 1),
+      Pos.new(bufnr, stop_row0 + 1, stop_col0),
+      'v'
+    )
+    if range.stop.lnum > vim.api.nvim_buf_line_count(bufnr) then
+      range.stop = range.stop:must_next(-1)
+    end
+
+    local capture_name = q.captures[id]
+    if not ranges[capture_name] then ranges[capture_name] = {} end
+    if self:contains(range) then table.insert(ranges[capture_name], range) end
+  end
+
+  return ranges
+end
+
 function Range:length()
   if self:is_empty() then return 0 end