feat: buf highlights for current buffer fuzzy find (#732)

* feat: Add buffer highlights from treesitter

* fix: Handle not having tree sitter in some buffers

* fixup

* fixup

* fixup: move back to old node

1 files changed, 48 insertions, 0 deletions

diff --git a/lua/telescope/builtin/files.lua b/lua/telescope/builtin/files.lua
index f091ad0..819d9ae 100644
--- a/lua/telescope/builtin/files.lua
+++ b/lua/telescope/builtin/files.lua
@@ -332,6 +332,7 @@ files.current_buffer_fuzzy_find = function(opts)
   -- All actions are on the current buffer
   local bufnr = vim.api.nvim_get_current_buf()
   local filename = vim.fn.expand(vim.api.nvim_buf_get_name(bufnr))
+  local filetype = vim.api.nvim_buf_get_option(bufnr, "filetype")
 
   local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
   local lines_with_numbers = {}
@@ -345,6 +346,53 @@ files.current_buffer_fuzzy_find = function(opts)
     })
   end
 
+  local ok, parser = pcall(vim.treesitter.get_parser, bufnr, filetype)
+  if ok then
+    local query = vim.treesitter.get_query(filetype, "highlights")
+
+    local root = parser:parse()[1]:root()
+
+    local highlighter = vim.treesitter.highlighter.new(parser)
+    local highlighter_query = highlighter:get_query(filetype)
+
+    local line_highlights = setmetatable({}, {
+      __index = function(t, k)
+        local obj = {}
+        rawset(t, k, obj)
+        return obj
+      end,
+    })
+    for id, node in query:iter_captures(root, bufnr, 0, -1) do
+      local hl = highlighter_query.hl_cache[id]
+      if hl then
+        local row1, col1, row2, col2 = node:range()
+
+        if row1 == row2 then
+          local row = row1 + 1
+
+          for index = col1, col2 do
+            line_highlights[row][index] = hl
+          end
+        else
+          local row = row1 + 1
+          for index = col1, #lines[row] do
+              line_highlights[row][index] = hl
+          end
+
+          while row < row2 + 1 do
+            row = row + 1
+
+            for index = 0, #lines[row] do
+              line_highlights[row][index] = hl
+            end
+          end
+        end
+      end
+    end
+
+    opts.line_highlights = line_highlights
+  end
+
   pickers.new(opts, {
     prompt_title = 'Current Buffer Fuzzy',
     finder = finders.new_table {