From 079773ebee5f18f55c72856049c1a332dabe28e2 Mon Sep 17 00:00:00 2001
From: TJ DeVries <devries.timothyj@gmail.com>
Date: Sat, 26 Sep 2020 12:36:58 -0400
Subject: [PATCH] fix: No longer lose highlights when moving through results

---
 lua/telescope/pickers.lua | 91 +++++++++++++++++++++------------------
 1 file changed, 50 insertions(+), 41 deletions(-)

diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index ebcb8ea..b6fba4d 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -227,7 +227,6 @@ function Picker:clear_extra_rows(results_bufnr)
   if self.sorting_strategy == 'ascending' then
     local num_results = self.manager:num_results()
     local worst_line = self.max_results - num_results
-    log.info(self.max_results, num_results, worst_line)
 
     if worst_line <= 0 then
       return
@@ -253,34 +252,38 @@ function Picker:highlight_displayed_rows(results_bufnr, prompt)
   vim.api.nvim_buf_clear_namespace(results_bufnr, ns_telescope_matching, 0, -1)
 
   local displayed_rows = vim.api.nvim_buf_get_lines(results_bufnr, 0, -1, false)
-  for row = 1, #displayed_rows do
-    local display = displayed_rows[row]
+  for row_index = 1, #displayed_rows do
+    local display = displayed_rows[row_index]
 
-    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:highlight_one_row(results_bufnr, prompt, display, row_index - 1)
+  end
+end
 
-        vim.api.nvim_buf_add_highlight(
-          results_bufnr,
-          ns_telescope_matching,
-          highlight,
-          row - 1,
-          start - 1,
-          finish
-        )
+function Picker:highlight_one_row(results_bufnr, prompt, display, row)
+  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
+
+      vim.api.nvim_buf_add_highlight(
+        results_bufnr,
+        ns_telescope_matching,
+        highlight,
+        row,
+        start - 1,
+        finish
+      )
     end
   end
 end
@@ -597,7 +600,7 @@ function Picker:close_windows(status)
 end
 
 function Picker:get_selection()
-  return self._selection
+  return self._selection_entry
 end
 
 function Picker:get_selection_row()
@@ -609,7 +612,7 @@ function Picker:move_selection(change)
 end
 
 function Picker:reset_selection()
-  self._selection = nil
+  self._selection_entry = nil
   self._selection_row = nil
 end
 
@@ -640,16 +643,27 @@ function Picker:set_selection(row)
   --        Probably something with setting a row that's too high for this?
   --        Not sure.
   local set_ok, set_errmsg = pcall(function()
+    local prompt = vim.api.nvim_buf_get_lines(self.prompt_bufnr, 0, 1, false)[1]
+
     -- Handle adding '> ' to beginning of selections
     if self._selection_row then
       local old_selection = a.nvim_buf_get_lines(results_bufnr, self._selection_row, self._selection_row + 1, false)[1]
 
       if old_selection then
-        a.nvim_buf_set_lines(results_bufnr, self._selection_row, self._selection_row + 1, false, {'  ' .. old_selection:sub(3)})
+        local old_display = '  ' .. old_selection:sub(3)
+        a.nvim_buf_set_lines(results_bufnr, self._selection_row, self._selection_row + 1, false, {old_display})
+
+        if prompt and self.sorter.highlighter then
+          self:highlight_one_row(results_bufnr, prompt, old_display, self._selection_row)
+        end
       end
     end
 
-    a.nvim_buf_set_lines(results_bufnr, row, row + 1, false, {'> ' .. (a.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1] or ''):sub(3)})
+    local display = '> ' .. (a.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1] or ''):sub(3)
+
+    -- TODO: You should go back and redraw the highlights for this line from the sorter.
+    -- That's the only smart thing to do.
+    a.nvim_buf_set_lines(results_bufnr, row, row + 1, false, {display})
 
     a.nvim_buf_clear_namespace(results_bufnr, ns_telescope_selection, 0, -1)
     a.nvim_buf_add_highlight(
@@ -660,6 +674,10 @@ function Picker:set_selection(row)
       0,
       -1
     )
+
+    if prompt and self.sorter.highlighter then
+      self:highlight_one_row(results_bufnr, prompt, display, row)
+    end
   end)
 
   if not set_ok then
@@ -667,21 +685,12 @@ function Picker:set_selection(row)
     return
   end
 
-  -- if self._match_id then
-  --   -- vim.fn.matchdelete(self._match_id)
-  --   vim.fn.clearmatches(results_win)
-  -- end
-
-  -- self._match_id = vim.fn.matchaddpos("Conceal", { {row + 1, 1, 2} }, 0, -1, { window = results_win, conceal = ">" })
-  if self._selection == entry and self._selection_row == row then
+  if self._selection_entry == entry and self._selection_row == row then
     return
   end
 
-  -- TODO: Don't let you go over / under the buffer limits
-  -- TODO: Make sure you start exactly at the bottom selected
-
   -- TODO: Get row & text in the same obj
-  self._selection = entry
+  self._selection_entry = entry
   self._selection_row = row
 
   if status.preview_win and self.previewer then