10 files changed, 271 insertions, 73 deletions

diff --git a/doc/telescope.txt b/doc/telescope.txt
index c80b093..c9a999c 100644
--- a/doc/telescope.txt
+++ b/doc/telescope.txt
@@ -118,6 +118,102 @@ builtin.live_grep()                                      *builtin.live_grep()*
 
 
 ================================================================================
+                                                             *telescope.actions*
+
+Actions functions that are useful for people creating their own mappings.
+
+actions.add_selection({prompt_bufnr})                *actions.add_selection()*
+
+    Add current entry to multi select
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_selection_better({prompt_bufnr})*actions.move_selection_better()*
+
+    Move the selection to the entry that has a better score
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_selection_next({prompt_bufnr})    *actions.move_selection_next()*
+
+    Move the selection to the next entry
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_selection_previous({prompt_bufnr})*actions.move_selection_previous()*
+
+    Move the selection to the previous entry
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_selection_worse({prompt_bufnr})  *actions.move_selection_worse()*
+
+    Move the selection to the entry that has a worse score
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_to_bottom({prompt_bufnr})              *actions.move_to_bottom()*
+
+    Move to the bottom of the picker
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_to_middle({prompt_bufnr})              *actions.move_to_middle()*
+
+    Move to the middle of the picker
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.move_to_top({prompt_bufnr})                    *actions.move_to_top()*
+
+    Move to the top of the picker
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.remove_selection({prompt_bufnr})          *actions.remove_selection()*
+
+    Remove current entry from multi select
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+actions.toggle_selection({prompt_bufnr})          *actions.toggle_selection()*
+
+    Toggle current entry status for multi select
+
+
+    Parameters: ~
+        {prompt_bufnr} (number)  The prompt bufnr
+
+
+
+================================================================================
                                                               *telescope.layout*
 
 
diff --git a/lua/telescope/actions/init.lua b/lua/telescope/actions/init.lua
index 2fa7ce2..79bcb23 100644
--- a/lua/telescope/actions/init.lua
+++ b/lua/telescope/actions/init.lua
@@ -1,4 +1,11 @@
--- Actions functions that are useful for people creating their own mappings.
+---@tag telescope.actions
+
+-- TODO: Add @module to make it so we can have the prefix.
+--@module telescope.actions
+
+---@brief [[
+--- Actions functions that are useful for people creating their own mappings.
+---@brief ]]
 
 local a = vim.api
 
@@ -40,7 +47,6 @@ local action_is_deprecated = function(name, err)
   )
 end
 
---- Get the current entry
 function actions.get_selected_entry()
   -- TODO(1.0): Remove
   action_is_deprecated("get_selected_entry")
@@ -53,7 +59,6 @@ function actions.get_current_line()
   return action_state.get_current_line()
 end
 
---- Get the current picker object for the prompt
 function actions.get_current_picker(prompt_bufnr)
   -- TODO(1.0): Remove
   action_is_deprecated("get_current_picker")
@@ -61,15 +66,33 @@ function actions.get_current_picker(prompt_bufnr)
 end
 
 --- Move the selection to the next entry
+---@param prompt_bufnr number: The prompt bufnr
 function actions.move_selection_next(prompt_bufnr)
   action_set.shift_selection(prompt_bufnr, 1)
 end
 
 --- Move the selection to the previous entry
+---@param prompt_bufnr number: The prompt bufnr
 function actions.move_selection_previous(prompt_bufnr)
   action_set.shift_selection(prompt_bufnr, -1)
 end
 
+--- Move the selection to the entry that has a worse score
+---@param prompt_bufnr number: The prompt bufnr
+function actions.move_selection_worse(prompt_bufnr)
+  local picker = action_state.get_current_picker(prompt_bufnr)
+  action_set.shift_selection(prompt_bufnr, p_scroller.worse(picker.sorting_strategy))
+end
+
+--- Move the selection to the entry that has a better score
+---@param prompt_bufnr number: The prompt bufnr
+function actions.move_selection_better(prompt_bufnr)
+  local picker = action_state.get_current_picker(prompt_bufnr)
+  action_set.shift_selection(prompt_bufnr, p_scroller.better(picker.sorting_strategy))
+end
+
+--- Move to the top of the picker
+---@param prompt_bufnr number: The prompt bufnr
 function actions.move_to_top(prompt_bufnr)
   local current_picker = actions.get_current_picker(prompt_bufnr)
   current_picker:set_selection(p_scroller.top(current_picker.sorting_strategy,
@@ -78,6 +101,8 @@ function actions.move_to_top(prompt_bufnr)
   ))
 end
 
+--- Move to the middle of the picker
+---@param prompt_bufnr number: The prompt bufnr
 function actions.move_to_middle(prompt_bufnr)
   local current_picker = actions.get_current_picker(prompt_bufnr)
   current_picker:set_selection(p_scroller.middle(
@@ -87,6 +112,8 @@ function actions.move_to_middle(prompt_bufnr)
   ))
 end
 
+--- Move to the bottom of the picker
+---@param prompt_bufnr number: The prompt bufnr
 function actions.move_to_bottom(prompt_bufnr)
   local current_picker = actions.get_current_picker(prompt_bufnr)
   current_picker:set_selection(p_scroller.bottom(current_picker.sorting_strategy,
@@ -95,16 +122,22 @@ function actions.move_to_bottom(prompt_bufnr)
   ))
 end
 
+--- Add current entry to multi select
+---@param prompt_bufnr number: The prompt bufnr
 function actions.add_selection(prompt_bufnr)
   local current_picker = action_state.get_current_picker(prompt_bufnr)
   current_picker:add_selection(current_picker:get_selection_row())
 end
 
+--- Remove current entry from multi select
+---@param prompt_bufnr number: The prompt bufnr
 function actions.remove_selection(prompt_bufnr)
   local current_picker = action_state.get_current_picker(prompt_bufnr)
   current_picker:remove_selection(current_picker:get_selection_row())
 end
 
+--- Toggle current entry status for multi select
+---@param prompt_bufnr number: The prompt bufnr
 function actions.toggle_selection(prompt_bufnr)
   local current_picker = action_state.get_current_picker(prompt_bufnr)
   current_picker:toggle_selection(current_picker:get_selection_row())
@@ -357,7 +390,7 @@ actions.send_selected_to_qflist = function(prompt_bufnr)
   local picker = action_state.get_current_picker(prompt_bufnr)
 
   local qf_entries = {}
-  for entry in pairs(picker.multi_select) do
+  for _, entry in ipairs(picker:get_multi_selection()) do
     table.insert(qf_entries, entry_to_qf(entry))
   end
 
diff --git a/lua/telescope/mappings.lua b/lua/telescope/mappings.lua
index 9490309..2e586d4 100644
--- a/lua/telescope/mappings.lua
+++ b/lua/telescope/mappings.lua
@@ -24,8 +24,9 @@ mappings.default_mappings = config.values.default_mappings or {
       ["<C-u>"] = actions.preview_scrolling_up,
       ["<C-d>"] = actions.preview_scrolling_down,
 
-      -- TODO: When we implement multi-select, you can turn this back on :)
-      -- ["<Tab>"] = actions.add_selection,
+      ["<Tab>"] = actions.toggle_selection + actions.move_selection_worse,
+      ["<C-q>"] = actions.send_to_qflist + actions.open_qflist,
+      ["<M-q>"] = actions.send_selected_to_qflist + actions.open_qflist,
     },
 
     n = {
@@ -35,6 +36,10 @@ mappings.default_mappings = config.values.default_mappings or {
       ["<C-v>"] = actions.select_vertical,
       ["<C-t>"] = actions.select_tab,
 
+      ["<Tab>"] = actions.toggle_selection + actions.move_selection_worse,
+      ["<C-q>"] = actions.send_to_qflist + actions.open_qflist,
+      ["<M-q>"] = actions.send_selected_to_qflist + actions.open_qflist,
+
       -- TODO: This would be weird if we switch the ordering.
       ["j"] = actions.move_selection_next,
       ["k"] = actions.move_selection_previous,
diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index c04ce17..3992361 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -19,6 +19,7 @@ local p_highlights = require('telescope.pickers.highlights')
 local p_scroller = require('telescope.pickers.scroller')
 
 local EntryManager = require('telescope.entry_manager')
+local MultiSelect = require('telescope.pickers.multi')
 
 local get_default = utils.get_default
 
@@ -73,6 +74,7 @@ function Picker:new(opts)
     cwd = opts.cwd,
 
     _completion_callbacks = {},
+    _multi = MultiSelect:new(),
 
     track = get_default(opts.track, false),
     stats = {},
@@ -245,7 +247,7 @@ function Picker:clear_extra_rows(results_bufnr)
     log.debug(msg)
   end
 
-  log.debug("Clearing:", worst_line)
+  log.trace("Clearing:", worst_line)
 end
 
 function Picker:highlight_displayed_rows(results_bufnr, prompt)
@@ -295,7 +297,7 @@ function Picker:highlight_one_row(results_bufnr, prompt, display, row)
   end
 
   local entry = self.manager:get_entry(self:get_index(row))
-  self.highlighter:hi_multiselect(row, entry)
+  self.highlighter:hi_multiselect(row, self:is_multi_selected(entry))
 end
 
 function Picker:can_select_row(row)
@@ -579,56 +581,36 @@ end
 
 function Picker:add_selection(row)
   local entry = self.manager:get_entry(self:get_index(row))
-  self.multi_select[entry] = true
+  self._multi:add(entry)
 
-  self.highlighter:hi_multiselect(row, entry)
+  self.highlighter:hi_multiselect(row, true)
 end
 
-function Picker:add_selection(row)
+function Picker:remove_selection(row)
   local entry = self.manager:get_entry(self:get_index(row))
-  self.multi_select[entry] = true
+  self._multi:drop(entry)
 
-  self.highlighter:hi_multiselect(row, entry)
+  self.highlighter:hi_multiselect(row, false)
 end
 
-function Picker:remove_selection(row)
-  local entry = self.manager:get_entry(self:get_index(row))
-  self.multi_select[entry] = nil
+function Picker:is_multi_selected(entry)
+  return self._multi:is_selected(entry)
+end
 
-  self.highlighter:hi_multiselect(row, entry)
+function Picker:get_multi_selection()
+  return self._multi:get()
 end
 
 function Picker:toggle_selection(row)
   local entry = self.manager:get_entry(self:get_index(row))
+  self._multi:toggle(entry)
 
-  if self.multi_select[entry] then
-    self:remove_selection(row)
-  else
-    self:add_selection(row)
-  end
-end
-
-function Picker:display_multi_select(results_bufnr)
-  for entry, _ in pairs(self.multi_select) do
-    local index = self.manager:find_entry(entry)
-    if index then
-      vim.api.nvim_buf_add_highlight(
-        results_bufnr,
-        a.nvim_create_namespace('telescope_selection'),
-        "TelescopeMultiSelection",
-        self:get_row(index),
-        0,
-        -1
-      )
-    end
-  end
+  self.highlighter:hi_multiselect(row, self._multi:is_selected(entry))
 end
 
 function Picker:reset_selection()
   self._selection_entry = nil
   self._selection_row = nil
-
-  self.multi_select = {}
 end
 
 function Picker:_reset_prefix_color(hl_group)
@@ -695,45 +677,50 @@ function Picker:set_selection(row)
     end
   end
 
-  local entry = self.manager:get_entry(self:get_index(row))
   local results_bufnr = self.results_bufnr
+  if not a.nvim_buf_is_valid(results_bufnr) then
+    return
+  end
 
   if row > a.nvim_buf_line_count(results_bufnr) then
-    error(string.format(
+    log.debug(string.format(
       "Should not be possible to get row this large %s %s",
         row,
         a.nvim_buf_line_count(results_bufnr)
     ))
-  end
 
-  state.set_global_key("selected_entry", entry)
-
-  if not vim.api.nvim_buf_is_valid(results_bufnr) then
     return
   end
 
+  local entry = self.manager:get_entry(self:get_index(row))
+  state.set_global_key("selected_entry", entry)
+
   -- TODO: Probably should figure out what the rows are that made this happen...
   --        Probably something with setting a row that's too high for this?
   --        Not sure.
   local set_ok, set_errmsg = pcall(function()
-    if not a.nvim_buf_is_valid(results_bufnr) then
-      return
-    end
-
     local prompt = self:_get_prompt()
 
     -- Handle adding '> ' to beginning of selections
     if self._selection_row then
-      local display, display_highlights = entry_display.resolve(self, self._selection_entry)
+      -- Only change the first couple characters, nvim_buf_set_text leaves the existing highlights
+      a.nvim_buf_set_text(
+        results_bufnr,
+        self._selection_row, 0,
+        self._selection_row, #self.entry_prefix,
+        { self.entry_prefix }
+      )
+      self.highlighter:hi_multiselect(
+        self._selection_row,
+        self:is_multi_selected(self._selection_entry)
+      )
 
-      if display then
-        display = self.entry_prefix .. display
-        a.nvim_buf_set_lines(results_bufnr, self._selection_row, self._selection_row + 1, false, {display})
+      -- local display = a.nvim_buf_get_lines(results_bufnr, old_row, old_row + 1, false)[1]
+      -- display = '  ' .. display
+      -- a.nvim_buf_set_lines(results_bufnr, old_row, old_row + 1, false, {display})
 
-        self.highlighter:hi_display(self._selection_row, self.entry_prefix, display_highlights)
-        self.highlighter:hi_sorter(self._selection_row, prompt, display)
-        self.highlighter:hi_multiselect(self._selection_row, self._selection_entry)
-      end
+      -- self.highlighter:hi_display(old_row, '  ', display_highlights)
+      -- self.highlighter:hi_sorter(old_row, prompt, display)
     end
 
     local caret = self.selection_caret
@@ -755,7 +742,8 @@ function Picker:set_selection(row)
     self.highlighter:hi_selection(row, caret:sub(1, -2))
     self.highlighter:hi_display(row, caret, display_highlights)
     self.highlighter:hi_sorter(row, prompt, display)
-    self.highlighter:hi_multiselect(row, entry)
+
+    self.highlighter:hi_multiselect(row, self:is_multi_selected(entry))
   end)
 
   if not set_ok then
@@ -819,7 +807,7 @@ function Picker:entry_adder(index, entry, _, insert)
     end
 
     if self.request_number ~= scheduled_request then
-      log.debug("Cancelling request number:", self.request_number, " // ", scheduled_request)
+      log.trace("Cancelling request number:", self.request_number, " // ", scheduled_request)
       return
     end
 
diff --git a/lua/telescope/pickers/highlights.lua b/lua/telescope/pickers/highlights.lua
index 62ffb4b..bbb06e4 100644
--- a/lua/telescope/pickers/highlights.lua
+++ b/lua/telescope/pickers/highlights.lua
@@ -1,4 +1,5 @@
 local a = vim.api
+local log = require('telescope.log')
 
 local highlights = {}
 
@@ -76,25 +77,30 @@ function Highlighter:hi_selection(row, caret)
   )
 end
 
-function Highlighter:hi_multiselect(row, entry)
+function Highlighter:hi_multiselect(row, is_selected)
   local results_bufnr = assert(self.picker.results_bufnr, "Must have a results bufnr")
 
-  if self.picker.multi_select[entry] then
+  if is_selected then
     vim.api.nvim_buf_add_highlight(
-      results_bufnr,
-      ns_telescope_multiselection,
-      "TelescopeMultiSelection",
-      row,
-      0,
-      -1
+      results_bufnr, ns_telescope_multiselection, "TelescopeMultiSelection", row, 0, -1
     )
   else
-    vim.api.nvim_buf_clear_namespace(
-      results_bufnr,
-      ns_telescope_multiselection,
-      row,
-      row + 1
+    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 a bout the gravity of the extmark?
+    if #existing_marks > 0 then
+      log.trace("Clearning row: ", row)
+
+      vim.api.nvim_buf_clear_namespace(
+        results_bufnr,
+        ns_telescope_multiselection,
+        row,
+        row + 1
+      )
+    end
   end
 end
 
diff --git a/lua/telescope/pickers/multi.lua b/lua/telescope/pickers/multi.lua
new file mode 100644
index 0000000..2896c08
--- /dev/null
+++ b/lua/telescope/pickers/multi.lua
@@ -0,0 +1,51 @@
+
+local MultiSelect = {}
+MultiSelect.__index = MultiSelect
+
+function MultiSelect:new()
+  return setmetatable({
+    _entries = {}
+  }, MultiSelect)
+end
+
+function MultiSelect:get()
+  local marked_entries = {}
+  for entry, count in pairs(self._entries) do
+    table.insert(marked_entries, {count, entry})
+  end
+
+  table.sort(marked_entries, function(left, right)
+    return left[1] < right[1]
+  end)
+
+  local selections = {}
+  for _, entry in ipairs(marked_entries) do
+    table.insert(selections, entry[2])
+  end
+
+  return selections
+end
+
+function MultiSelect:is_selected(entry)
+  return self._entries[entry]
+end
+
+local multi_select_count = 0
+function MultiSelect:add(entry)
+  multi_select_count = multi_select_count + 1
+  self._entries[entry] = multi_select_count
+end
+
+function MultiSelect:drop(entry)
+  self._entries[entry] = nil
+end
+
+function MultiSelect:toggle(entry)
+  if self:is_selected(entry) then
+    self:drop(entry)
+  else
+    self:add(entry)
+  end
+end
+
+return MultiSelect
diff --git a/lua/telescope/pickers/scroller.lua b/lua/telescope/pickers/scroller.lua
index b77fc96..d2aed2c 100644
--- a/lua/telescope/pickers/scroller.lua
+++ b/lua/telescope/pickers/scroller.lua
@@ -99,4 +99,16 @@ scroller.bottom = function(sorting_strategy, max_results, num_results)
   return max_results - 1
 end
 
+scroller.better = function(sorting_strategy)
+  if sorting_strategy == 'ascending' then
+    return -1
+  else
+    return 1
+  end
+end
+
+scroller.worse = function(sorting_strategy)
+  return -(scroller.better(sorting_strategy))
+end
+
 return scroller
diff --git a/lua/tests/automated/pickers/scrolling_spec.lua b/lua/tests/automated/pickers/scrolling_spec.lua
index ab82378..1a86822 100644
--- a/lua/tests/automated/pickers/scrolling_spec.lua
+++ b/lua/tests/automated/pickers/scrolling_spec.lua
@@ -2,6 +2,9 @@ require('plenary.reload').reload_module('telescope')
 
 local tester = require('telescope.pickers._test')
 
+local log = require('telescope.log')
+log.use_console = false
+
 describe('scrolling strategies', function()
   it('should handle cycling for full list', function()
     tester.run_file [[find_files__scrolling_descending_cycle]]
diff --git a/lua/tests/automated/scroller_spec.lua b/lua/tests/automated/scroller_spec.lua
index c650aac..7c2a557 100644
--- a/lua/tests/automated/scroller_spec.lua
+++ b/lua/tests/automated/scroller_spec.lua
@@ -1,5 +1,8 @@
 local p_scroller = require('telescope.pickers.scroller')
 
+local log = require('telescope.log')
+log.use_console = false
+
 local eq = assert.are.same
 
 describe('scroller', function()
diff --git a/scripts/gendocs.lua b/scripts/gendocs.lua
index 0cf8a80..fdf3c5a 100644
--- a/scripts/gendocs.lua
+++ b/scripts/gendocs.lua
@@ -11,6 +11,7 @@ docs.test = function()
     "./lua/telescope/init.lua",
     "./lua/telescope/builtin/init.lua",
     "./lua/telescope/pickers/layout_strategies.lua",
+    "./lua/telescope/actions/init.lua",
   }
 
   table.sort(input_files, function(a, b)