3 files changed, 76 insertions, 3 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 {
diff --git a/lua/telescope/make_entry.lua b/lua/telescope/make_entry.lua
index e4b3d4e..607b8ea 100644
--- a/lua/telescope/make_entry.lua
+++ b/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
 
diff --git a/lua/telescope/pickers/entry_display.lua b/lua/telescope/pickers/entry_display.lua
index 991c9bc..f64f693 100644
--- a/lua/telescope/pickers/entry_display.lua
+++ b/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