feat: Add more sorter hooks (#752)

* feat: Add more sorter hooks

* fix breaking conni brain

2 files changed, 35 insertions, 12 deletions

diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index 7e658c3..00e1049 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -401,9 +401,7 @@ function Picker:find()
         self.finder = finder
       end
 
-      if self.sorter then
-        self.sorter:_start(prompt)
-      end
+      if self.sorter then self.sorter:_start(prompt) end
 
       -- TODO: Entry manager should have a "bulk" setter. This can prevent a lot of redraws from display
       self.manager = EntryManager:new(self.max_results, self.entry_adder, self.stats)
@@ -422,9 +420,6 @@ function Picker:find()
     end
   end)
 
-  -- on_lines(nil, nil, nil, 0, 1)
-  status_updater()
-
   -- Register attach
   vim.api.nvim_buf_attach(prompt_bufnr, false, {
     on_lines = tx.send,
@@ -442,7 +437,9 @@ function Picker:find()
     end,
   })
 
+  if self.sorter then self.sorter:_init() end
   async_lib.run(main_loop())
+  status_updater()
 
   -- TODO: Use WinLeave as well?
   local on_buf_leave = string.format(
@@ -993,13 +990,13 @@ function Picker:get_result_completor(results_bufnr, find_id, prompt, status_upda
     local current_line = vim.api.nvim_get_current_line():sub(self.prompt_prefix:len() + 1)
     state.set_global_key('current_line', current_line)
 
+    status_updater()
+
     self:clear_extra_rows(results_bufnr)
     self:highlight_displayed_rows(results_bufnr, prompt)
-
+    if self.sorter then self.sorter:_finish(prompt) end
 
     self:_on_complete()
-
-    status_updater()
   end
 end
 
@@ -1037,13 +1034,14 @@ function pickers.on_close_prompt(prompt_bufnr)
   local status = state.get_status(prompt_bufnr)
   local picker = status.picker
 
+  if picker.sorter then
+    picker.sorter:_destroy()
+  end
+
   if picker.previewer then
     picker.previewer:teardown()
   end
 
-  -- TODO: This is an attempt to clear all the memory stuff we may have left.
-  -- vim.api.nvim_buf_detach(prompt_bufnr)
-
   picker.close_windows(status)
 end
 
diff --git a/lua/telescope/sorters.lua b/lua/telescope/sorters.lua
index 4147f59..bf4a735 100644
--- a/lua/telescope/sorters.lua
+++ b/lua/telescope/sorters.lua
@@ -38,6 +38,10 @@ Sorter.__index = Sorter
 ---@field highlighter function: Highlights results to display them pretty
 ---@field discard boolean: Whether this is a discardable style sorter or not.
 ---@field score function: Override the score function if desired.
+---@field init function: Function to run when creating sorter
+---@field start function: Function to run on every new prompt
+---@field finish function: Function to run after every new prompt
+---@field destroy function: Functo to run when destroying sorter
 function Sorter:new(opts)
   opts = opts or {}
 
@@ -45,6 +49,13 @@ function Sorter:new(opts)
     score = opts.score,
     state = {},
     tags = opts.tags,
+
+    -- State management
+    init    = opts.init,
+    start   = opts.start,
+    finish  = opts.finish,
+    destroy = opts.destroy,
+
     filter_function = opts.filter_function,
     scoring_function = opts.scoring_function,
     highlighter = opts.highlighter,
@@ -56,12 +67,22 @@ function Sorter:new(opts)
   }, Sorter)
 end
 
+function Sorter:_init()
+  if self.init then self:init() end
+end
+
+function Sorter:_destroy()
+  if self.destroy then self:destroy() end
+end
+
 -- TODO: We could make this a bit smarter and cache results "as we go" and where they got filtered.
 --          Then when we hit backspace, we don't have to re-caculate everything.
 --          Prime did a lot of the hard work already, but I don't want to copy as much memory around
 --              as he did in his example.
 --              Example can be found in ./scratch/prime_prompt_cache.lua
 function Sorter:_start(prompt)
+  if self.start then self:start(prompt) end
+
   if not self.discard then
     return
   end
@@ -80,6 +101,10 @@ function Sorter:_start(prompt)
   self._discard_state.prompt = prompt
 end
 
+function Sorter:_finish(prompt)
+  if self.finish then self:finish(prompt) end
+end
+
 -- TODO: Consider doing something that makes it so we can skip the filter checks
 --          if we're not discarding. Also, that means we don't have to check otherwise as well :)
 function Sorter:score(prompt, entry, cb_add, cb_filter)