feat: Add TelescopeMatching a new highlight for matches from sorters

2 files changed, 77 insertions, 8 deletions

diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index 27d6021..1bdc269 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -24,6 +24,9 @@ local extend = function(opts, defaults)
   return result
 end
 
+local ns_telescope_selection = a.nvim_create_namespace('telescope_selection')
+local ns_telescope_matching = a.nvim_create_namespace('telescope_matching')
+
 local pickers = {}
 
 -- TODO: Add motions to keybindings
@@ -239,6 +242,47 @@ function Picker:clear_extra_rows(results_bufnr)
   end
 end
 
+function Picker:highlight_displayed_rows(results_bufnr, prompt)
+  if not self.sorter.highlighter then
+    return
+  end
+
+
+  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]
+
+    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 - 1,
+          start - 1,
+          finish
+        )
+      end
+    end
+  end
+end
+
 function Picker:can_select_row(row)
   if self.sorting_strategy == 'ascending' then
     return row <= self.manager:num_results()
@@ -414,6 +458,7 @@ function Picker:find()
       end
 
       self:clear_extra_rows(results_bufnr)
+      self:highlight_displayed_rows(results_bufnr, prompt)
 
       PERF("Filtered Amount    ", filtered_amount)
       PERF("Displayed Amount   ", displayed_amount)
@@ -549,8 +594,6 @@ function Picker:close_windows(status)
   state.clear_status(status.prompt_bufnr)
 end
 
-local ns_telescope_selection = a.nvim_create_namespace('telescope_selection')
-
 function Picker:get_selection()
   return self._selection
 end
diff --git a/lua/telescope/sorters.lua b/lua/telescope/sorters.lua
index 6371bbe..b41ab91 100644
--- a/lua/telescope/sorters.lua
+++ b/lua/telescope/sorters.lua
@@ -3,6 +3,23 @@ local util = require('telescope.utils')
 
 local sorters = {}
 
+local ngram_highlighter = function(ngram_len, prompt, display)
+  local highlights = {}
+  display = display:lower()
+
+  for disp_index = 1, #display do
+    local char = display:sub(disp_index, disp_index + ngram_len - 1)
+    if prompt:find(char, 1, true) then
+      table.insert(highlights, {
+        start = disp_index,
+        finish = disp_index + ngram_len - 1
+      })
+    end
+  end
+
+  return highlights
+end
+
 
 local Sorter = {}
 Sorter.__index = Sorter
@@ -21,6 +38,7 @@ function Sorter:new(opts)
   return setmetatable({
     state = {},
     scoring_function = opts.scoring_function,
+    highlighter = opts.highlighter,
   }, Sorter)
 end
 
@@ -204,12 +222,16 @@ sorters.get_fuzzy_file = function(opts)
       end
 
       return 1 / denominator
-    end
+    end,
+
+    highlighter = function(_, prompt, display)
+      return ngram_highlighter(ngram_len, prompt, display)
+    end,
   }
 end
 
 sorters.get_generic_fuzzy_sorter = function()
-  local ngramlen = 2
+  local ngram_len = 2
 
   local cached_ngrams = {}
 
@@ -238,14 +260,14 @@ sorters.get_generic_fuzzy_sorter = function()
     -- line (entry.ordinal)
     -- entry (the whole entry)
     scoring_function = function(_, prompt, line, _)
-      if prompt == 0 or #prompt < ngramlen then
+      if prompt == 0 or #prompt < ngram_len then
         return 0
       end
 
       local prompt_lower = prompt:lower()
       local line_lower = line:lower()
 
-      local prompt_ngrams = overlapping_ngrams(prompt_lower, ngramlen)
+      local prompt_ngrams = overlapping_ngrams(prompt_lower, ngram_len)
 
       local N = #prompt
 
@@ -273,7 +295,7 @@ sorters.get_generic_fuzzy_sorter = function()
         -- biases for shorter strings
         -- TODO(ashkan): this can bias towards repeated finds of the same
         -- subpattern with overlapping_ngrams
-        + 3 * match_count * ngramlen / #line
+        + 3 * match_count * ngram_len / #line
         + consecutive_matches
         + N / (contains_string or (2 * #line))
         -- + 30/(c1 or 2*N)
@@ -288,7 +310,11 @@ sorters.get_generic_fuzzy_sorter = function()
       end
 
       return 1 / denominator
-    end
+    end,
+
+    highlighter = function(_, prompt, display)
+      return ngram_highlighter(ngram_len, prompt, display)
+    end,
   }
 end