From 6120afa159d1dd3ba112ee4360b4ab4562a9b266 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Thu, 27 Oct 2022 14:06:26 +0200
Subject: [PATCH] feat: use vim.treesitter.get_captures_at_pos to detect
 comments

---
 lua/todo-comments/highlight.lua | 41 ++++++++++-----------------------
 1 file changed, 12 insertions(+), 29 deletions(-)

diff --git a/lua/todo-comments/highlight.lua b/lua/todo-comments/highlight.lua
index ffeae13..0a53835 100644
--- a/lua/todo-comments/highlight.lua
+++ b/lua/todo-comments/highlight.lua
@@ -39,35 +39,14 @@ function M.match(str, patterns)
 end
 
 -- This method returns nil if this buf doesn't have a treesitter parser
--- @return true or false otherwise
-function M.is_comment(buf, line)
-  local highlighter = require("vim.treesitter.highlighter")
-  local hl = highlighter.active[buf]
-
-  if not hl then
-    return
-  end
-
-  local is_comment = false
-  hl.tree:for_each_tree(function(tree, lang_tree)
-    if is_comment then
-      return
-    end
-
-    local query = hl:get_query(lang_tree:lang())
-    if not (query and query:query()) then
-      return
+--- @return boolean? true or false otherwise
+function M.is_comment(buf, row, col)
+  local captures = vim.treesitter.get_captures_at_pos(buf, row, col)
+  for _, c in ipairs(captures) do
+    if c.capture == "comment" then
+      return true
     end
-
-    local iter = query:query():iter_captures(tree:root(), buf, line, line + 1)
-
-    for capture, _ in iter do
-      if query._query.captures[capture] == "comment" then
-        is_comment = true
-      end
-    end
-  end)
-  return is_comment
+  end
 end
 
 local function add_highlight(buffer, ns, hl, line, from, to)
@@ -103,7 +82,11 @@ function M.highlight(buf, first, last, _event)
     local lnum = first + l - 1
 
     if ok and start then
-      if Config.options.highlight.comments_only and not M.is_quickfix(buf) and M.is_comment(buf, lnum) == false then
+      if
+        Config.options.highlight.comments_only
+        and not M.is_quickfix(buf)
+        and M.is_comment(buf, lnum, start - 1) == false
+      then
         kw = nil
       end
     end