rfc: use extmarks for highlighting and carets (#2099)

* Starting to make some progress on extmarks for prefixes

* Finished up extmarks for multi-icons

* Cleaned up update_prefix calls, added display highlights

* remove highlight_one_row

* Remove TODOs

* stylua fixes

* fixes for luacheck

* Fixed whitespace to appease the stylua gods

* fixed a couple nits from pr

* Got tests passing with new highlighting functionality

* Apply suggestions from code review

Co-authored-by: Fabian David Schmidt <39233597+fdschmidt93@users.noreply.github.com>

Co-authored-by: Fabian David Schmidt <39233597+fdschmidt93@users.noreply.github.com>

3 files changed, 110 insertions, 116 deletions

diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index b704bf5..616f915 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -262,43 +262,6 @@ function Picker:clear_extra_rows(results_bufnr)
   log.trace("Clearing:", worst_line)
 end
 
---- Highlight the entry corresponding to the given row
----@param results_bufnr number: the buffer number of the results buffer
----@param prompt table: table with information about the prompt buffer
----@param display string: the text corresponding to the given row
----@param row number: the number of the chosen row
-function Picker:highlight_one_row(results_bufnr, prompt, display, row)
-  if not self.sorter.highlighter then
-    return
-  end
-
-  local highlights = self.sorter:highlighter(prompt, display)
-
-  if highlights then
-    for _, hl in ipairs(highlights) do
-      local highlight, start, finish
-      if type(hl) == "table" then
-        highlight = hl.highlight or "TelescopeMatching"
-        start = hl.start
-        finish = hl.finish or hl.start
-      elseif type(hl) == "number" then
-        highlight = "TelescopeMatching"
-        start = hl
-        finish = hl
-      else
-        error "Invalid higlighter fn"
-      end
-
-      self:_increment "highlights"
-
-      vim.api.nvim_buf_add_highlight(results_bufnr, ns_telescope_matching, highlight, row, start - 1, finish)
-    end
-  end
-
-  local entry = self.manager:get_entry(self:get_index(row))
-  self.highlighter:hi_multiselect(row, self:is_multi_selected(entry))
-end
-
 --- Check if the given row number can be selected
 ---@param row number: the number of the chosen row in the results buffer
 ---@return boolean
@@ -799,7 +762,6 @@ function Picker:add_selection(row)
   local entry = self.manager:get_entry(self:get_index(row))
   self._multi:add(entry)
 
-  self:update_prefix(entry, row)
   self:get_status_updater(self.prompt_win, self.prompt_bufnr)()
   self.highlighter:hi_multiselect(row, true)
 end
@@ -810,7 +772,6 @@ function Picker:remove_selection(row)
   local entry = self.manager:get_entry(self:get_index(row))
   self._multi:drop(entry)
 
-  self:update_prefix(entry, row)
   self:get_status_updater(self.prompt_win, self.prompt_bufnr)()
   self.highlighter:hi_multiselect(row, false)
 end
@@ -839,7 +800,6 @@ function Picker:toggle_selection(row)
   end
   self._multi:toggle(entry)
 
-  self:update_prefix(entry, row)
   self:get_status_updater(self.prompt_win, self.prompt_bufnr)()
   self.highlighter:hi_multiselect(row, self._multi:is_selected(entry))
 end
@@ -987,7 +947,6 @@ function Picker:set_selection(row)
       self._selection_entry = entry
 
       if old_row >= 0 then
-        self:update_prefix(old_entry, old_row)
         self.highlighter:hi_multiselect(old_row, self:is_multi_selected(old_entry))
       end
     else
@@ -1046,25 +1005,7 @@ function Picker:update_prefix(entry, row)
     return t
   end
 
-  local line = vim.api.nvim_buf_get_lines(self.results_bufnr, row, row + 1, false)[1]
-  if not line then
-    log.trace(string.format("no line found at row %d in buffer %d", row, self.results_bufnr))
-    return
-  end
-
-  local old_caret = string.sub(line, 0, #prefix(true)) == prefix(true) and prefix(true)
-    or string.sub(line, 0, #prefix(true, true)) == prefix(true, true) and prefix(true, true)
-    or string.sub(line, 0, #prefix(false)) == prefix(false) and prefix(false)
-    or string.sub(line, 0, #prefix(false, true)) == prefix(false, true) and prefix(false, true)
-  if old_caret == false then
-    log.warn(string.format("can't identify old caret in line: %s", line))
-    return
-  end
-
-  local pre = prefix(entry == self._selection_entry, self:is_multi_selected(entry))
-  -- Only change the first couple characters, nvim_buf_set_text leaves the existing highlights
-  a.nvim_buf_set_text(self.results_bufnr, row, 0, row, #old_caret, { pre })
-  return pre
+  return prefix(entry == self._selection_entry, self:is_multi_selected(entry))
 end
 
 --- Refresh the previewer based on the current `status` of the picker
@@ -1164,8 +1105,7 @@ function Picker:entry_adder(index, entry, _, insert)
     if display_highlights then
       self.highlighter:hi_display(row, prefix, display_highlights)
     end
-    self:update_prefix(entry, row)
-    self:highlight_one_row(self.results_bufnr, self:_get_prompt(), display, row)
+    self.highlighter:hi_sorter(row, self:_get_prompt(), display)
   end
 
   if not set_ok then
@@ -1371,7 +1311,6 @@ function Picker:_do_selection(prompt)
       local old_entry, old_row = self._selection_entry, self._selection_row
       self:reset_selection() -- required to reset selection before updating prefix
       if old_row >= 0 then
-        self:update_prefix(old_entry, old_row)
         self.highlighter:hi_multiselect(old_row, self:is_multi_selected(old_entry))
       end
     end
diff --git a/lua/telescope/pickers/highlights.lua b/lua/telescope/pickers/highlights.lua
index be693a7..44f26a5 100644
--- a/lua/telescope/pickers/highlights.lua
+++ b/lua/telescope/pickers/highlights.lua
@@ -1,5 +1,4 @@
 local a = vim.api
-local log = require "telescope.log"
 local conf = require("telescope.config").values
 
 local highlights = {}
@@ -7,6 +6,13 @@ local highlights = {}
 local ns_telescope_selection = a.nvim_create_namespace "telescope_selection"
 local ns_telescope_multiselection = a.nvim_create_namespace "telescope_multiselection"
 local ns_telescope_entry = a.nvim_create_namespace "telescope_entry"
+local ns_telescope_matching = a.nvim_create_namespace "telescope_matching"
+
+-- Priorities for extmark highlights. Example: Treesitter is set to 100
+local SELECTION_MULTISELECT_PRIORITY = 100
+local DISPLAY_HIGHLIGHTS_PRIORITY = 110
+local SELECTION_HIGHLIGHTS_PRIORITY = 130
+local SORTER_HIGHLIGHTS_PRIORITY = 140
 
 local Highlighter = {}
 Highlighter.__index = Highlighter
@@ -25,19 +31,19 @@ function Highlighter:hi_display(row, prefix, display_highlights)
   end
 
   local results_bufnr = assert(self.picker.results_bufnr, "Must have a results bufnr")
+  if not a.nvim_buf_is_valid(results_bufnr) then
+    return
+  end
 
   a.nvim_buf_clear_namespace(results_bufnr, ns_telescope_entry, row, row + 1)
-  local len_prefix = #prefix
 
   for _, hl_block in ipairs(display_highlights) do
-    a.nvim_buf_add_highlight(
-      results_bufnr,
-      ns_telescope_entry,
-      hl_block[2],
-      row,
-      len_prefix + hl_block[1][1],
-      len_prefix + hl_block[1][2]
-    )
+    a.nvim_buf_set_extmark(results_bufnr, ns_telescope_entry, row, #prefix + hl_block[1][1], {
+      end_col = #prefix + hl_block[1][2],
+      hl_group = hl_block[2],
+      priority = DISPLAY_HIGHLIGHTS_PRIORITY,
+      strict = false,
+    })
   end
 end
 
@@ -61,7 +67,35 @@ function Highlighter:hi_sorter(row, prompt, display)
   end
 
   local results_bufnr = assert(self.picker.results_bufnr, "Must have a results bufnr")
-  picker:highlight_one_row(results_bufnr, prompt, display, row)
+  if not a.nvim_buf_is_valid(results_bufnr) then
+    return
+  end
+
+  local sorter_highlights = picker.sorter:highlighter(prompt, display)
+
+  if sorter_highlights then
+    for _, hl in ipairs(sorter_highlights) do
+      local highlight, start, finish
+      if type(hl) == "table" then
+        highlight = hl.highlight or "TelescopeMatching"
+        start = hl.start
+        finish = hl.finish or hl.start
+      elseif type(hl) == "number" then
+        highlight = "TelescopeMatching"
+        start = hl
+        finish = hl
+      else
+        error "Invalid highlight"
+      end
+
+      a.nvim_buf_set_extmark(results_bufnr, ns_telescope_matching, row, start - 1, {
+        end_col = finish,
+        hl_group = highlight,
+        priority = SORTER_HIGHLIGHTS_PRIORITY,
+        strict = false,
+      })
+    end
+  end
 end
 
 function Highlighter:hi_selection(row, caret)
@@ -69,52 +103,73 @@ function Highlighter:hi_selection(row, caret)
   local results_bufnr = assert(self.picker.results_bufnr, "Must have a results bufnr")
 
   a.nvim_buf_clear_namespace(results_bufnr, ns_telescope_selection, 0, -1)
-  a.nvim_buf_add_highlight(results_bufnr, ns_telescope_selection, "TelescopeSelectionCaret", row, 0, #caret)
-
-  a.nvim_buf_set_extmark(
-    results_bufnr,
-    ns_telescope_selection,
-    row,
-    #caret,
-    { end_line = row + 1, hl_eol = conf.hl_result_eol, hl_group = "TelescopeSelection" }
-  )
+
+  -- Skip if there is nothing on the actual line
+  local line = a.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1]
+  if line == nil or line == "" then
+    return
+  end
+
+  local offset = #caret
+
+  -- Highlight the caret
+  a.nvim_buf_set_extmark(results_bufnr, ns_telescope_selection, row, 0, {
+    virt_text = { { caret, "TelescopeSelectionCaret" } },
+    virt_text_pos = "overlay",
+    end_col = offset,
+    hl_group = "TelescopeSelectionCaret",
+    hl_mode = "combine",
+    priority = SELECTION_HIGHLIGHTS_PRIORITY,
+    strict = true,
+  })
+
+  -- Highlight the text after the caret
+  a.nvim_buf_set_extmark(results_bufnr, ns_telescope_selection, row, offset, {
+    end_line = row + 1,
+    hl_eol = conf.hl_result_eol,
+    hl_group = "TelescopeSelection",
+    priority = SELECTION_HIGHLIGHTS_PRIORITY,
+  })
 end
 
 function Highlighter:hi_multiselect(row, is_selected)
   local results_bufnr = assert(self.picker.results_bufnr, "Must have a results bufnr")
 
+  if not a.nvim_buf_is_valid(results_bufnr) then
+    return
+  end
+
+  a.nvim_buf_clear_namespace(results_bufnr, ns_telescope_multiselection, row, row + 1)
+
+  local line = a.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1]
+  if line == nil or line == "" then
+    return
+  end
+
   if is_selected then
-    vim.api.nvim_buf_add_highlight(results_bufnr, ns_telescope_multiselection, "TelescopeMultiSelection", row, 0, -1)
-    if self.picker.multi_icon then
-      local line = vim.api.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1]
-      local pos = line:find(self.picker.multi_icon)
-      if pos and pos <= math.max(#self.picker.selection_caret, #self.picker.entry_prefix) then
-        vim.api.nvim_buf_add_highlight(
-          results_bufnr,
-          ns_telescope_multiselection,
-          "TelescopeMultiIcon",
-          row,
-          pos - 1,
-          pos - 1 + #self.picker.multi_icon
-        )
-      end
-    end
-  else
-    local existing_marks = vim.api.nvim_buf_get_extmarks(
-      results_bufnr,
-      ns_telescope_multiselection,
-      { row, 0 },
-      { row, -1 },
-      {}
-    )
-
-    -- This is still kind of weird to me, since it seems like I'm erasing stuff
-    -- when I shouldn't... Perhaps it's about the gravity of the extmark?
-    if #existing_marks > 0 then
-      log.trace("Clearing highlight multi select row: ", row)
-
-      vim.api.nvim_buf_clear_namespace(results_bufnr, ns_telescope_multiselection, row, row + 1)
-    end
+    local multi_icon = self.picker.multi_icon
+    local offset = #multi_icon
+
+    a.nvim_buf_set_extmark(results_bufnr, ns_telescope_multiselection, row, offset, {
+      virt_text = { { multi_icon, "TelescopeMultiIcon" } },
+      virt_text_pos = "overlay",
+      end_col = offset,
+      hl_mode = "combine",
+      hl_group = "TelescopeMultiIcon",
+      priority = SELECTION_HIGHLIGHTS_PRIORITY,
+    })
+    -- highlight the caret
+    a.nvim_buf_set_extmark(results_bufnr, ns_telescope_multiselection, row, 0, {
+      end_col = #self.picker.selection_caret,
+      hl_group = "TelescopeMultiSelection",
+      priority = SELECTION_MULTISELECT_PRIORITY,
+    })
+    -- highlight the text after the multi_icon
+    a.nvim_buf_set_extmark(results_bufnr, ns_telescope_multiselection, row, offset, {
+      end_col = #line,
+      hl_group = "TelescopeMultiSelection",
+      priority = SELECTION_MULTISELECT_PRIORITY,
+    })
   end
 end
 
diff --git a/lua/tests/automated/pickers/find_files_spec.lua b/lua/tests/automated/pickers/find_files_spec.lua
index 7e1c027..98dca3f 100644
--- a/lua/tests/automated/pickers/find_files_spec.lua
+++ b/lua/tests/automated/pickers/find_files_spec.lua
@@ -26,7 +26,7 @@ describe("builtin.find_files", function()
         runner.picker('find_files', 'README.md', {
           post_typed = {
             { "> README.md", GetPrompt },
-            { "> README.md", GetBestResult },
+            { "  README.md", GetBestResult },
           },
           post_close = {
             { 'README.md', GetFile },
@@ -112,7 +112,7 @@ describe("builtin.find_files", function()
           {
             {
               "  lua/tests/fixtures/file_a.txt",
-              "> lua/tests/fixtures/file_abc.txt",
+              "  lua/tests/fixtures/file_abc.txt",
             }, GetResults
           },
         },