feat: multiselect icon (#1572)

* feat: add `multi_icon` option to pickers and corresponding highlight

* feat: allow `multi_icon` to be any length

* fix: adjust `selection_caret` highlighting

5 files changed, 64 insertions, 7 deletions

diff --git a/doc/telescope.txt b/doc/telescope.txt
index 65d6b43..0c1f33d 100644
--- a/doc/telescope.txt
+++ b/doc/telescope.txt
@@ -184,6 +184,15 @@ telescope.setup({opts})                                    *telescope.setup()*
 
         Default: '  '
 
+                                            *telescope.defaults.multi_icon*
+    multi_icon: ~
+        Symbol to add in front of a multi-selected result entry.
+        Replaces final character of |telescope.defaults.selection_caret| and
+        |telescope.defaults.entry_prefix| as appropriate.
+        To have no icon, set to the empty string.
+
+        Default: '+'
+
                                           *telescope.defaults.initial_mode*
     initial_mode: ~
         Determines in which mode telescope starts. Valid Keys:
diff --git a/lua/telescope/config.lua b/lua/telescope/config.lua
index 7bd07dc..043ac8a 100644
--- a/lua/telescope/config.lua
+++ b/lua/telescope/config.lua
@@ -246,6 +246,18 @@ append(
 )
 
 append(
+  "multi_icon",
+  "+",
+  [[
+  Symbol to add in front of a multi-selected result entry.
+  Replaces final character of |telescope.defaults.selection_caret| and
+  |telescope.defaults.entry_prefix| as appropriate.
+  To have no icon, set to the empty string.
+
+  Default: '+']]
+)
+
+append(
   "initial_mode",
   "insert",
   [[
diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index 66e6370..9385a49 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -27,6 +27,9 @@ local MultiSelect = require "telescope.pickers.multi"
 
 local get_default = utils.get_default
 
+local truncate = require("plenary.strings").truncate
+local strdisplaywidth = require("plenary.strings").strdisplaywidth
+
 local ns_telescope_matching = a.nvim_create_namespace "telescope_matching"
 local ns_telescope_prompt = a.nvim_create_namespace "telescope_prompt"
 local ns_telescope_prompt_prefix = a.nvim_create_namespace "telescope_prompt_prefix"
@@ -68,6 +71,8 @@ function Picker:new(opts)
     prompt_prefix = get_default(opts.prompt_prefix, config.values.prompt_prefix),
     selection_caret = get_default(opts.selection_caret, config.values.selection_caret),
     entry_prefix = get_default(opts.entry_prefix, config.values.entry_prefix),
+    multi_icon = get_default(opts.multi_icon, config.values.multi_icon),
+
     initial_mode = get_default(opts.initial_mode, config.values.initial_mode),
     debounce = get_default(tonumber(opts.debounce), nil),
 
@@ -923,29 +928,45 @@ 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]
+
       --Check if that row still has the caret
-      if string.sub(line, 0, #self.selection_caret) == self.selection_caret then
+      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,
-          #self.selection_caret,
-          { self.entry_prefix }
+          #old_caret,
+          { prefix(false, old_multiselected) }
         )
-        self.highlighter:hi_multiselect(row_old_selection, self:is_multi_selected(self._selection_entry))
+        self.highlighter:hi_multiselect(row_old_selection, old_multiselected)
       end
     end
 
-    local caret = self.selection_caret
+    local caret = prefix(true, self:is_multi_selected(entry))
 
     local display, display_highlights = entry_display.resolve(self, entry)
     display = caret .. display
@@ -958,8 +979,8 @@ function Picker:set_selection(row)
     end
     a.nvim_buf_set_lines(results_bufnr, row, row + 1, false, { display })
 
-    -- don't highlight the ' ' at the end of caret
-    self.highlighter:hi_selection(row, caret:sub(1, -2))
+    -- 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)
 
diff --git a/lua/telescope/pickers/highlights.lua b/lua/telescope/pickers/highlights.lua
index 6d89b78..3ac6510 100644
--- a/lua/telescope/pickers/highlights.lua
+++ b/lua/telescope/pickers/highlights.lua
@@ -84,6 +84,20 @@ function Highlighter:hi_multiselect(row, is_selected)
 
   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,
diff --git a/plugin/telescope.vim b/plugin/telescope.vim
index 5bf3cf0..1aa033f 100644
--- a/plugin/telescope.vim
+++ b/plugin/telescope.vim
@@ -12,6 +12,7 @@ let g:loaded_telescope = 1
 highlight default link TelescopeSelection Visual
 highlight default link TelescopeSelectionCaret TelescopeSelection
 highlight default link TelescopeMultiSelection Type
+highlight default link TelescopeMultiIcon Identifier
 
 " "Normal" in the floating windows created by telescope.
 highlight default link TelescopeNormal Normal