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>

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