async pickers take 3

lua/telescope/finders.lua

@@ -2,6 +2,13 @@ local Job = require('plenary.job')
 
 local make_entry = require('telescope.make_entry')
 local log = require('telescope.log')
+local co = coroutine
+local a = require('plenary.async_lib')
+local async = a.async
+local await = a.await
+local channel = a.util.channel
+local utils = require('telescope.utils')
+local Executor = utils.Executor
 
 local finders = {}
 
@@ -124,7 +131,7 @@ function OneshotJobFinder:new(opts)
     _started = false,
   }, self)
 
-  obj._find = coroutine.wrap(function(finder, _, process_result, process_complete)
+  obj._find = coroutine.wrap(function(finder, _, process_result, process_complete, cancel_rx)
     local num_execution = 1
     local num_results = 0
 
@@ -176,18 +183,45 @@ function OneshotJobFinder:new(opts)
 
     job:start()
 
+    local counter = 0
     while true do
-      finder, _, process_result, process_complete = coroutine.yield()
+      cancelled = false
+      finder, _, process_result, process_complete, cancel_rx = coroutine.yield()
       num_execution = num_execution + 1
-
       local current_count = num_results
-      for index = 1, current_count do
-        process_result(results[index])
-      end
 
-      if completed then
-        process_complete()
-      end
+      local update_cancelled = async(function()
+        await(cancel_rx())
+        -- print('cancelled!', counter)
+        counter = counter + 1
+        cancelled = true
+      end)
+
+      a.run(update_cancelled())
+
+      local exe = Executor.new {
+        -- how many tasks to run for each loop
+        run_task_amount = 1,
+      }
+
+      exe:add(co.create(function()
+        for index = 1, current_count do
+          if cancelled then
+            -- print('canelling')
+            return
+          end
+
+          process_result(results[index])
+
+          co.yield()
+        end
+
+        if completed then
+          process_complete()
+        end
+      end))
+
+      exe:run()
     end
   end)
 

lua/telescope/pickers.lua

@@ -26,6 +26,10 @@ local get_default = utils.get_default
 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')
+local as = require('plenary.async_lib')
+local async = as.async
+local await = as.await
+local channel = as.util.channel
 
 local pickers = {}
 
@@ -395,74 +399,93 @@ function Picker:find()
   local debounced_status = debounce.throttle_leading(status_updater, 50)
 
   self.request_number = 0
-  local on_lines = function(_, _, _, first_line, last_line)
-    self.request_number = self.request_number + 1
-    self:_reset_track()
+  self.__on_lines = on_lines
 
-    if not vim.api.nvim_buf_is_valid(prompt_bufnr) then
-      log.debug("ON_LINES: Invalid prompt_bufnr", prompt_bufnr)
-      return
-    end
+  -- on_lines(nil, nil, nil, 0, 1)
+  status_updater()
 
-    if not first_line then first_line = 0 end
-    if not last_line then last_line = 1 end
+  local tx, rx = channel.mpsc()
+  local should_stop = false
 
-    if first_line > 0 or last_line > 1 then
-      log.debug("ON_LINES: Bad range", first_line, last_line)
-      return
-    end
+  -- Register attach
+  vim.api.nvim_buf_attach(prompt_bufnr, false, {
+    on_lines = tx.send,
+    on_detach = function()
+      should_stop = true
+      -- -- TODO: Can we add a "cleanup" / "teardown" function that completely removes these.
+      self.finder = nil
+      self.previewer = nil
+      self.sorter = nil
+      self.manager = nil
 
-    local original_prompt = self:_get_prompt()
-    local on_input_result = self._on_input_filter_cb(original_prompt) or {}
+      -- -- TODO: Should we actually do this?
+      collectgarbage(); collectgarbage()
+    end,
+  })
 
-    local prompt = on_input_result.prompt or original_prompt
-    local finder = on_input_result.updated_finder
+  local main_loop = async(function()
+    local cancel_tx, cancel_rx = channel.oneshot()
+
+    while not should_stop do
+      -- this creates a new timer, should be able to reuse it
+      -- adjust polling rate based on number of entries
+      await(as.util.sleep(50)) -- polling rate
+      local _, _, _, first_line, last_line = await(rx.last())
+      -- these next two lines are extremely important
+      -- they will cancel the previous find
+      cancel_tx()
+      cancel_tx, cancel_rx = channel.oneshot()
+      await(as.scheduler())
+
+      self.request_number = self.request_number + 1
+      self:_reset_track()
+
+      if not vim.api.nvim_buf_is_valid(prompt_bufnr) then
+        log.debug("ON_LINES: Invalid prompt_bufnr", prompt_bufnr)
+        return
+      end
 
-    if finder then
-      self.finder:close()
-      self.finder = finder
-    end
+      if not first_line then first_line = 0 end
+      if not last_line then last_line = 1 end
 
-    if self.sorter then
-      self.sorter:_start(prompt)
-    end
+      if first_line > 0 or last_line > 1 then
+        log.debug("ON_LINES: Bad range", first_line, last_line)
+        return
+      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, self.request_number)
+      local original_prompt = self:_get_prompt()
+      local on_input_result = self._on_input_filter_cb(original_prompt) or {}
 
-    local process_result = self:get_result_processor(prompt, debounced_status)
-    local process_complete = self:get_result_completor(self.results_bufnr, prompt, status_updater)
+      local prompt = on_input_result.prompt or original_prompt
+      local finder = on_input_result.updated_finder
 
-    local ok, msg = pcall(function()
-      self.finder(prompt, process_result, vim.schedule_wrap(process_complete))
-    end)
+      if finder then
+        self.finder:close()
+        self.finder = finder
+      end
 
-    if not ok then
-      log.warn("Failed with msg: ", msg)
-    end
-  end
+      if self.sorter then
+        self.sorter:_start(prompt)
+      end
 
-  self.__on_lines = on_lines
+      -- 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, self.request_number)
+      self.manager = EntryManager:new(self.max_results, nil, self.stats, self.request_number)
 
-  on_lines(nil, nil, nil, 0, 1)
-  status_updater()
+      local process_result = self:get_result_processor(prompt, debounced_status)
+      local process_complete = self:get_result_completor(self.results_bufnr, prompt, status_updater)
 
-  -- Register attach
-  vim.api.nvim_buf_attach(prompt_bufnr, false, {
-    on_lines = on_lines,
-    on_detach = vim.schedule_wrap(function()
-      on_lines = nil
+      local ok, msg = pcall(function()
+        self.finder(prompt, process_result, vim.schedule_wrap(process_complete), cancel_rx)
+      end)
 
-      -- TODO: Can we add a "cleanup" / "teardown" function that completely removes these.
-      self.finder = nil
-      self.previewer = nil
-      self.sorter = nil
-      self.manager = nil
+      if not ok then
+        log.warn("Failed with msg: ", msg)
+      end
+    end
+  end)
 
-      -- TODO: Should we actually do this?
-      collectgarbage(); collectgarbage()
-    end),
-  })
+  as.run(main_loop())
 
   -- TODO: Use WinLeave as well?
   local on_buf_leave = string.format(
@@ -986,6 +1009,12 @@ end
 
 function Picker:get_result_completor(results_bufnr, prompt, status_updater)
   return function()
+    local counter = 1
+    for entry in self.manager:iter() do
+      self:entry_adder(counter, entry, nil, false)
+      counter = counter + 1
+    end
+
     if self:is_done() then return end
 
     local selection_strategy = self.selection_strategy or 'reset'

lua/telescope/utils.lua

@@ -2,6 +2,8 @@ local has_devicons, devicons = pcall(require, 'nvim-web-devicons')
 
 local pathlib = require('telescope.path')
 local Job     = require('plenary.job')
+local co = coroutine
+local uv = vim.loop
 
 local utils = {}
 
@@ -392,4 +394,89 @@ utils.get_devicons = (function()
   end
 end)()
 
+local Executor = {}
+Executor.__index = Executor
+
+function Executor.new(opts)
+  opts = opts or {}
+
+  local self = setmetatable({}, Executor)
+
+  self.tasks = opts.tasks or {}
+  self.mode = opts.mode or "next"
+  self.index = opts.start_idx or 1
+  self.run_task_amount = opts.run_task_amount or 1
+  self.idle = uv.new_idle()
+
+  return self
+end
+
+function Executor:run()
+  self.idle:start(vim.schedule_wrap(function()
+    if #self.tasks == 0 then
+      self.idle:stop()
+      return
+    end
+
+    for _ = 1, self.run_task_amount do
+      if self.mode == "finish" then
+        self:step_finish()
+      else
+        self:step()
+      end
+    end
+  end))
+end
+
+function Executor:close()
+  self.idle:stop()
+  self.tasks = {}
+end
+
+function Executor:step_finish()
+  if #self.tasks == 0 then return end
+  local curr_task = self.tasks[self.index]
+  if curr_task == nil then
+    self.index = 1
+    curr_task = self.tasks[self.index]
+  end
+
+  local _, _ = co.resume(curr_task)
+  if co.status(curr_task) == "dead" then
+    table.remove(self.tasks, self.index)
+
+    self.index = self.index + 1
+  end
+end
+
+function Executor:step()
+  if #self.tasks == 0 then return end
+  local curr_task = self.tasks[self.index]
+  if curr_task == nil then
+    self.index = 1
+    curr_task = self.tasks[self.index]
+  end
+
+  local _, _ = co.resume(curr_task[1], unpack(curr_task[2]))
+  if co.status(curr_task[1]) == "dead" then
+    table.remove(self.tasks, self.index)
+  end
+
+  self.index = self.index + 1
+end
+
+function Executor:get_current_task()
+  return self.tasks[self.index]
+end
+
+function Executor:remove_task(idx)
+  table.remove(self.tasks, idx)
+end
+
+function Executor:add(task, ...)
+  table.insert(self.tasks, {task, {...}})
+end
+
+utils.Executor = Executor
+
 return utils