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

Author: tjdevries

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 {

lua/telescope/make_entry.lua

@@ -685,14 +685,30 @@ function make_entry.gen_from_buffer_lines(opts)
     separator = ' │ ',
     items = {
       { width = 5 },
-      { remaining = true },
+      { remaining = true, },
     },
   }
 
   local make_display = function(entry)
+
     return displayer {
       { entry.lnum, opts.lnum_highlight_group or 'TelescopeResultsSpecialComment' },
-      entry.line
+      {
+        entry.line, function()
+          if not opts.line_highlights then return {} end
+
+          local line_hl = opts.line_highlights[entry.lnum] or {}
+          -- TODO: We could probably squash these together if the are the same...
+          --        But I don't think that it's worth it at the moment.
+          local result = {}
+
+          for col, hl in pairs(line_hl) do
+            table.insert(result, { {col, col+1}, hl })
+          end
+
+          return result
+        end
+      },
     }
   end
 

lua/telescope/pickers/entry_display.lua

@@ -59,8 +59,16 @@ entry_display.create = function(configuration)
             hl_start = hl_start + #results[j] + (#configuration.separator or 1)
           end
           local hl_end = hl_start + #str:gsub('%s*$', '')
-          table.insert(highlights, { { hl_start, hl_end }, hl })
+
+          if type(hl) == "function" then
+            for _, hl_res in ipairs(hl()) do
+              table.insert(highlights, { { hl_res[1][1] + hl_start, hl_res[1][2] + hl_start }, hl_res[2] })
+            end
+          else
+            table.insert(highlights, { { hl_start, hl_end }, hl })
+          end
         end
+
         table.insert(results, str)
       end
     end
@@ -75,6 +83,7 @@ entry_display.create = function(configuration)
         table.insert(highlights, { { hl_start, hl_end }, configuration.separator_hl })
       end
     end
+
     local final_str = table.concat(results, configuration.separator or "│")
     if configuration.hl_chars then
       for i = 1, #final_str do

scratch/buffer_highlights.lua

@@ -0,0 +1,24 @@
+local a = vim.api
+
+local ns = a.nvim_create_namespace("treesitter/highlighter")
+print(ns)
+local bufnr = 0
+
+-- P(a.nvim_buf_get_extmarks(bufnr, ns, 0, -1, { details = true }))
+
+local parser = vim.treesitter.get_parser(bufnr, "lua")
+local query = vim.treesitter.get_query("lua", "highlights")
+P(query)
+
+local root = parser:parse()[1]:root()
+print("root", root)
+
+local highlighter = vim.treesitter.highlighter.new(parser)
+local highlighter_query = highlighter:get_query("lua")
+
+for id, node, metadata in query:iter_captures(root, bufnr, 0, -1) do
+  local row1, col1, row2, col2 = node:range()
+  print(highlighter_query.hl_cache[id])
+  -- print(id, node, metadata, vim.treesitter.get_node_text(node, bufnr))
+  -- print(">>>>", row1, col1, row2, col2)
+end

scratch/short_buf.lua

@@ -0,0 +1,4 @@
+local x = {}
+print(x); print("wow");
+
+local function other() end