feat: cmp async (#1583)

lua/cmp/utils/async.lua

@@ -1,4 +1,5 @@
 local feedkeys = require('cmp.utils.feedkeys')
+local config = require('cmp.config')
 
 local async = {}
 
@@ -9,6 +10,7 @@ local async = {}
 ---@field public stop function
 ---@field public __call function
 
+---@type uv_timer_t[]
 local timers = {}
 
 vim.api.nvim_create_autocmd('VimLeavePre', {
@@ -27,7 +29,8 @@ vim.api.nvim_create_autocmd('VimLeavePre', {
 ---@return cmp.AsyncThrottle
 async.throttle = function(fn, timeout)
   local time = nil
-  local timer = vim.loop.new_timer()
+  local timer = assert(vim.loop.new_timer())
+  local _async = nil ---@type Async?
   timers[#timers + 1] = timer
   return setmetatable({
     running = false,
@@ -37,9 +40,15 @@ async.throttle = function(fn, timeout)
         return not self.running
       end)
     end,
-    stop = function()
-      time = nil
+    stop = function(reset_time)
+      if reset_time ~= false then
+        time = nil
+      end
       timer:stop()
+      if _async then
+        _async:cancel()
+        _async = nil
+      end
     end,
   }, {
     __call = function(self, ...)
@@ -50,12 +59,23 @@ async.throttle = function(fn, timeout)
       end
 
       self.running = true
-      timer:stop()
+      self.stop(false)
       timer:start(math.max(1, self.timeout - (vim.loop.now() - time)), 0, function()
         vim.schedule(function()
           time = nil
-          fn(unpack(args))
-          self.running = false
+          local ret = fn(unpack(args))
+          if async.is_async(ret) then
+            ---@cast ret Async
+            _async = ret
+            _async:await(function(_, error)
+              self.running = false
+              if error and error ~= 'abort' then
+                vim.notify(error, vim.log.levels.ERROR)
+              end
+            end)
+          else
+            self.running = false
+          end
         end)
       end)
     end,
@@ -147,4 +167,124 @@ async.debounce_next_tick_by_keymap = function(callback)
   end
 end
 
+local Scheduler = {}
+Scheduler._queue = {}
+Scheduler._executor = assert(vim.loop.new_check())
+
+function Scheduler.step()
+  local budget = config.get().performance.async_budget * 1e6
+  local start = vim.loop.hrtime()
+  while #Scheduler._queue > 0 and vim.loop.hrtime() - start < budget do
+    local a = table.remove(Scheduler._queue, 1)
+    a:_step()
+    if a.running then
+      table.insert(Scheduler._queue, a)
+    end
+  end
+  if #Scheduler._queue == 0 then
+    return Scheduler._executor:stop()
+  end
+end
+
+---@param a Async
+function Scheduler.add(a)
+  table.insert(Scheduler._queue, a)
+  if not Scheduler._executor:is_active() then
+    Scheduler._executor:start(vim.schedule_wrap(Scheduler.step))
+  end
+end
+
+--- @alias AsyncCallback fun(result?:any, error?:string)
+
+--- @class Async
+--- @field running boolean
+--- @field result? any
+--- @field error? string
+--- @field callbacks AsyncCallback[]
+--- @field thread thread
+local Async = {}
+Async.__index = Async
+
+function Async.new(fn)
+  local self = setmetatable({}, Async)
+  self.callbacks = {}
+  self.running = true
+  self.thread = coroutine.create(fn)
+  Scheduler.add(self)
+  return self
+end
+
+---@param result? any
+---@param error? string
+function Async:_done(result, error)
+  self.running = false
+  self.result = result
+  self.error = error
+  for _, callback in ipairs(self.callbacks) do
+    callback(result, error)
+  end
+end
+
+function Async:_step()
+  local ok, res = coroutine.resume(self.thread)
+  if not ok then
+    return self:_done(nil, res)
+  elseif res == 'abort' then
+    return self:_done(nil, 'abort')
+  elseif coroutine.status(self.thread) == 'dead' then
+    return self:_done(res)
+  end
+end
+
+function Async:cancel()
+  self.running = false
+end
+
+---@param cb AsyncCallback
+function Async:await(cb)
+  if not cb then
+    error('callback is required')
+  end
+  if self.running then
+    table.insert(self.callbacks, cb)
+  else
+    cb(self.result, self.error)
+  end
+end
+
+function Async:sync()
+  while self.running do
+    vim.wait(10)
+  end
+  return self.error and error(self.error) or self.result
+end
+
+--- @return boolean
+function async.is_async(obj)
+  return obj and type(obj) == 'table' and getmetatable(obj) == Async
+end
+
+--- @return fun(...): Async
+function async.wrap(fn)
+  return function(...)
+    local args = { ... }
+    return Async.new(function()
+      return fn(unpack(args))
+    end)
+  end
+end
+
+-- This will yield when called from a coroutine
+function async.yield(...)
+  if not coroutine.isyieldable() then
+    error('Trying to yield from a non-yieldable context')
+    return ...
+  end
+  return coroutine.yield(...)
+end
+
+function async.abort()
+  return async.yield('abort')
+end
+
 return async