From 4e77964569ef47a70f9bb76c668dcfea2d089d5a Mon Sep 17 00:00:00 2001
From: Sergey Tarasov <dzirtusss@gmail.com>
Date: Wed, 17 Jul 2024 01:12:17 +0300
Subject: [PATCH] feat: Adds syntax highlighting in the telescope picker (#386)

* Add telescope highlighting

* Fix no TS language parser edge case

* Fix language detection for non-matching languages

* Remove nvim-treesitter dependency, streamline logic

* Change to show_columns option

* Update README
---
 README.md                            |  2 +
 lua/telescope/_extensions/aerial.lua | 87 +++++++++++++++++++++++++---
 2 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index d456ce7..9ee80cd 100644
--- a/README.md
+++ b/README.md
@@ -613,6 +613,8 @@ require("telescope").setup({
         json = true, -- You can set the option for specific filetypes
         yaml = true,
       },
+      -- Available modes: symbols, lines, both
+      show_columns = "both",
     },
   },
 })
diff --git a/lua/telescope/_extensions/aerial.lua b/lua/telescope/_extensions/aerial.lua
index c0634fa..da4655e 100644
--- a/lua/telescope/_extensions/aerial.lua
+++ b/lua/telescope/_extensions/aerial.lua
@@ -5,7 +5,8 @@ local pickers = require("telescope.pickers")
 local telescope = require("telescope")
 
 local ext_config = {
-  show_lines = true,
+  -- show_lines = true, -- deprecated in favor of show_columns
+  show_columns = "both", -- { "symbols", "lines", "both" }
   show_nesting = {
     ["_"] = false,
     json = true,
@@ -26,6 +27,19 @@ local function aerial_picker(opts)
   local filename = vim.api.nvim_buf_get_name(0)
   local filetype = vim.bo[bufnr].filetype
   local show_nesting = ext_config.show_nesting[filetype]
+
+  local show_columns = opts.show_columns or conf.show_columns
+  local show_lines = opts.show_lines or conf.show_lines -- show_lines is deprecated
+  if show_columns == nil then
+    if show_lines == true then
+      show_columns = "both"
+    elseif show_lines == false then
+      show_columns = "symbols"
+    else
+      show_columns = ext_config.show_columns
+    end
+  end
+
   if show_nesting == nil then
     show_nesting = ext_config.show_nesting["_"]
   end
@@ -39,7 +53,7 @@ local function aerial_picker(opts)
   end
 
   local layout
-  if ext_config.show_lines then
+  if show_columns == "both" then
     layout = {
       { width = 4 },
       { width = 30 },
@@ -57,8 +71,55 @@ local function aerial_picker(opts)
       items = layout,
     })
 
+  local function collect_buf_highlights()
+    local parser = vim.treesitter.get_parser(bufnr)
+    if not parser then
+      return {}
+    end
+
+    local lang = parser:lang()
+    local root = parser:trees()[1]:root() -- get root of already parsed cached tree
+
+    local highlights = {}
+    local query = vim.treesitter.query.get(lang, "highlights")
+    if query then
+      for _, captures, _ in query:iter_matches(root, bufnr, 0, -1) do
+        for id, node in pairs(captures) do
+          local start_row, start_col, _, end_col = node:range()
+          highlights[start_row] = highlights[start_row] or {}
+          table.insert(highlights[start_row], { start_col, end_col, query.captures[id] })
+        end
+      end
+    end
+    return highlights
+  end
+
+  local buf_highlights = {}
+  if show_columns == "lines" or show_columns == "both" then
+    buf_highlights = collect_buf_highlights() -- collect buffer highlights only if needed
+  end
+
+  local function highlights_for_row(row, offset)
+    offset = offset or 0
+    local row_highlights = buf_highlights[row] or {}
+    local highlights = {}
+    for _, value in ipairs(row_highlights) do
+      local start_col, end_col, hl_type = unpack(value)
+      hl_type = hl_type:match("^[^.]+") -- strip subtypes after dot
+      table.insert(highlights, { { start_col + offset, end_col + offset }, hl_type })
+    end
+    return highlights
+  end
+
   local function make_display(entry)
     local item = entry.value
+    local row = item.lnum - 1
+
+    if show_columns == "lines" then
+      local text = vim.api.nvim_buf_get_lines(bufnr, row, row + 1, false)[1] or ""
+      return text, highlights_for_row(row)
+    end
+
     local icon = config.get_icon(bufnr, item.kind)
     local icon_hl = highlight.get_highlight(item, true, false) or "NONE"
     local name_hl = highlight.get_highlight(item, false, false) or "NONE"
@@ -66,12 +127,21 @@ local function aerial_picker(opts)
       { icon, icon_hl },
       { entry.name, name_hl },
     }
-    if ext_config.show_lines then
-      local text = vim.api.nvim_buf_get_lines(bufnr, item.lnum - 1, item.lnum, false)[1] or ""
-      text = vim.trim(text)
-      table.insert(columns, text)
+
+    local highlights = {}
+    if show_columns == "both" then
+      local text = vim.api.nvim_buf_get_lines(bufnr, row, row + 1, false)[1] or ""
+      table.insert(columns, vim.trim(text))
+
+      local leading_spaces = text:match("^%s*")
+      local offset = layout[1].width + layout[2].width - #leading_spaces + #icon
+      if #entry.name > layout[2].width then
+        offset = offset + 2 -- '...' symbol
+      end
+      highlights = highlights_for_row(row, offset)
     end
-    return displayer(columns)
+
+    return displayer(columns), highlights
   end
 
   local function make_entry(item)
@@ -114,7 +184,8 @@ local function aerial_picker(opts)
   end
 
   -- Reverse the symbols so they have the same top-to-bottom order as in the file
-  if conf.sorting_strategy == "descending" then
+  local sorting_strategy = opts.sorting_strategy or conf.sorting_strategy
+  if sorting_strategy == "descending" then
     util.tbl_reverse(results)
     default_selection_index = #results - (default_selection_index - 1)
   end