fix: update multi_icon with select/drop/toggle_all actions (#1682)

* fix: `multi_icon` with `select/drop/toggle_all`

* typos

* fix: add check for no caret found

* fix: add check for no line found

* fix: check `max_results` in `Picker:can_select_row`

- also switch order of highlighting in `select/drop/toggle_all` actions

* fix: make `max_results` check a strict inequality

* [docgen] Update doc/telescope.txt
skip-checks: true

* fix: update `prompt_status` on `select/drop/toggle_all` actions

Co-authored-by: Github Actions <actions@github>

doc/telescope.txt

@@ -792,6 +792,7 @@ builtin.fd()                                                    *builtin.fd()*
     This is an alias for the `find_files` picker
 
 
+
 builtin.treesitter()                                    *builtin.treesitter()*
     Lists function names, variables, and other symbols from treesitter queries
     - Default keymaps:

lua/telescope/actions/init.lua

@@ -113,10 +113,15 @@ function actions.select_all(prompt_bufnr)
     if not current_picker._multi:is_selected(entry) then
       current_picker._multi:add(entry)
       if current_picker:can_select_row(row) then
+        local caret = current_picker:update_prefix(entry, row)
+        if current_picker._selection_entry == entry and current_picker._selection_row == row then
+          current_picker.highlighter:hi_selection(row, caret:match "(.*%S)")
+        end
         current_picker.highlighter:hi_multiselect(row, current_picker._multi:is_selected(entry))
       end
     end
   end)
+  current_picker:get_status_updater(current_picker.prompt_win, current_picker.prompt_bufnr)()
 end
 
 --- Drop all entries from the current multi selection.
@@ -126,9 +131,14 @@ function actions.drop_all(prompt_bufnr)
   action_utils.map_entries(prompt_bufnr, function(entry, _, row)
     current_picker._multi:drop(entry)
     if current_picker:can_select_row(row) then
+      local caret = current_picker:update_prefix(entry, row)
+      if current_picker._selection_entry == entry and current_picker._selection_row == row then
+        current_picker.highlighter:hi_selection(row, caret:match "(.*%S)")
+      end
       current_picker.highlighter:hi_multiselect(row, current_picker._multi:is_selected(entry))
     end
   end)
+  current_picker:get_status_updater(current_picker.prompt_win, current_picker.prompt_bufnr)()
 end
 
 --- Toggle multi selection for all entries.
@@ -139,9 +149,14 @@ function actions.toggle_all(prompt_bufnr)
   action_utils.map_entries(prompt_bufnr, function(entry, _, row)
     current_picker._multi:toggle(entry)
     if current_picker:can_select_row(row) then
+      local caret = current_picker:update_prefix(entry, row)
+      if current_picker._selection_entry == entry and current_picker._selection_row == row then
+        current_picker.highlighter:hi_selection(row, caret:match "(.*%S)")
+      end
       current_picker.highlighter:hi_multiselect(row, current_picker._multi:is_selected(entry))
     end
   end)
+  current_picker:get_status_updater(current_picker.prompt_win, current_picker.prompt_bufnr)()
 end
 
 function actions.preview_scrolling_up(prompt_bufnr)

lua/telescope/pickers.lua

@@ -291,7 +291,7 @@ end
 ---@return boolean
 function Picker:can_select_row(row)
   if self.sorting_strategy == "ascending" then
-    return row <= self.manager:num_results()
+    return row <= self.manager:num_results() and row < self.max_results
   else
     return row >= 0 and row <= self.max_results and row >= self.max_results - self.manager:num_results()
   end
@@ -766,6 +766,7 @@ 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
@@ -776,6 +777,7 @@ 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
@@ -801,6 +803,7 @@ function Picker:toggle_selection(row)
   local entry = self.manager:get_entry(self:get_index(row))
   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
@@ -932,47 +935,24 @@ function Picker:set_selection(row)
   --        Not sure.
   local set_ok, set_errmsg = pcall(function()
     local prompt = self:_get_prompt()
-    local prefix = function(sel, multi)
-      local t
-      if sel then
-        t = self.selection_caret
-      else
-        t = self.entry_prefix
-      end
-      if multi and type(self.multi_icon) == "string" then
-        t = truncate(t, strdisplaywidth(t) - strdisplaywidth(self.multi_icon), "") .. self.multi_icon
-      end
-      return t
-    end
 
-    -- This block handles removing the caret from beginning of previous selection (if still visible)
     -- Check if previous selection is still visible
     if self._selection_entry and self.manager:find_entry(self._selection_entry) then
-      -- Find the (possibly new) row of the old selection
-      local row_old_selection = self:get_row(self.manager:find_entry(self._selection_entry))
-      local old_multiselected = self:is_multi_selected(self._selection_entry)
-      local line = a.nvim_buf_get_lines(results_bufnr, row_old_selection, row_old_selection + 1, false)[1]
+      -- Find old selection, and update prefix and highlights
+      local old_entry = self._selection_entry
+      local old_row = self:get_row(self.manager:find_entry(old_entry))
 
-      --Check if that row still has the caret
-      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)
-      if old_caret then
-        -- Only change the first couple characters, nvim_buf_set_text leaves the existing highlights
-        a.nvim_buf_set_text(
-          results_bufnr,
-          row_old_selection,
-          0,
-          row_old_selection,
-          #old_caret,
-          { prefix(false, old_multiselected) }
-        )
-        self.highlighter:hi_multiselect(row_old_selection, old_multiselected)
-      end
+      self._selection_entry = entry
+
+      self:update_prefix(old_entry, old_row)
+      self.highlighter:hi_multiselect(old_row, self:is_multi_selected(old_entry))
+    else
+      self._selection_entry = entry
     end
 
-    local caret = prefix(true, self:is_multi_selected(entry))
+    local caret = self:update_prefix(entry, row)
 
-    local display, display_highlights = entry_display.resolve(self, entry)
+    local display, _ = entry_display.resolve(self, entry)
     display = caret .. display
 
     -- TODO: You should go back and redraw the highlights for this line from the sorter.
@@ -981,11 +961,9 @@ function Picker:set_selection(row)
       log.debug "Invalid buf somehow..."
       return
     end
-    a.nvim_buf_set_lines(results_bufnr, row, row + 1, false, { display })
 
     -- don't highlight any whitespace at the end of caret
     self.highlighter:hi_selection(row, caret:match "(.*%S)")
-    self.highlighter:hi_display(row, caret, display_highlights)
     self.highlighter:hi_sorter(row, prompt, display)
 
     self.highlighter:hi_multiselect(row, self:is_multi_selected(entry))
@@ -1009,6 +987,42 @@ function Picker:set_selection(row)
   vim.api.nvim_win_set_cursor(self.results_win, { row + 1, 0 })
 end
 
+--- Update prefix for entry on a given row
+function Picker:update_prefix(entry, row)
+  local prefix = function(sel, multi)
+    local t
+    if sel then
+      t = self.selection_caret
+    else
+      t = self.entry_prefix
+    end
+    if multi and type(self.multi_icon) == "string" then
+      t = truncate(t, strdisplaywidth(t) - strdisplaywidth(self.multi_icon), "") .. self.multi_icon
+    end
+    return t
+  end
+
+  local line = vim.api.nvim_buf_get_lines(self.results_bufnr, row, row + 1, false)[1]
+  if not line then
+    log.warn(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
+end
+
 --- Refresh the previewer based on the current `status` of the picker
 function Picker:refresh_previewer()
   local status = state.get_status(self.prompt_bufnr)
@@ -1110,6 +1124,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)
   end
 

lua/telescope/pickers/highlights.lua

@@ -5,7 +5,7 @@ local conf = require("telescope.config").values
 local highlights = {}
 
 local ns_telescope_selection = a.nvim_create_namespace "telescope_selection"
-local ns_telescope_multiselection = a.nvim_create_namespace "telescope_mulitselection"
+local ns_telescope_multiselection = a.nvim_create_namespace "telescope_multiselection"
 local ns_telescope_entry = a.nvim_create_namespace "telescope_entry"
 
 local Highlighter = {}
@@ -19,7 +19,7 @@ end
 
 function Highlighter:hi_display(row, prefix, display_highlights)
   -- This is the bug that made my highlight fixes not work.
-  -- We will leave the solutino commented, so the test fails.
+  -- We will leave the solution commented, so the test fails.
   if not display_highlights or vim.tbl_isempty(display_highlights) then
     return
   end
@@ -109,9 +109,9 @@ function Highlighter:hi_multiselect(row, is_selected)
     )
 
     -- This is still kind of weird to me, since it seems like I'm erasing stuff
-    -- when i shouldn't... perhaps it's a bout the gravity of the extmark?
+    -- when I shouldn't... Perhaps it's about the gravity of the extmark?
     if #existing_marks > 0 then
-      log.trace("Clearning highlight multi select row: ", row)
+      log.trace("Clearing highlight multi select row: ", row)
 
       vim.api.nvim_buf_clear_namespace(results_bufnr, ns_telescope_multiselection, row, row + 1)
     end