vim.task proposal

task.lua

---@enum vim.task.TaskStatus
local TaskStatus = {
  pending = 'pending',
  success = 'success',
  failure = 'failure',
  aborted = 'aborted',
}

---Internal registry for tracking asynchronous tasks.
---@type table<thread, vim.task.Task>
local threads = {}

---The interrupt key
local interrupt_key = vim.keycode('<C-c>')

---The namespace of the interrupt.
local interrupt_ns = vim.api.nvim_create_namespace('vim.task.interrupt')

---Convert callback-style values to pcall-style values.
---@param err? unknown
---@vararg ...
---@return boolean, ...
local function callback2pcall(err, ...)
  if err then
    return false, err
  end
  return true, ...
end

---Split ok and data part on tuple.
---@param ok boolean
---@vararg ...
---@return boolean, unknown[]
local function ok_data(ok, ...)
  return ok, { ... }
end

---Get current time in milliseconds.
---@return integer
local function now_ms()
  return vim.uv.hrtime() / 1e6
end

---Check if a value is callable.
---@param v unknown
---@return boolean
local function is_callable(v)
  if type(v) == 'function' then
    return true
  end
  local mt = getmetatable(v)
  if mt and type(mt.__call) == 'function' then
    return true
  end
  return false
end

---A standard callback signature, typically used for asynchronous operations.
---The first argument `err` is for an error value (if any), followed by any results.
---@alias vim.task.Callback fun(err?: unknown, ...: unknown)

---Represents the context for a yielding.
---This object provides methods for interacting with the asynchronous task it belongs to.
---@class vim.task.YieldingContext
---@field is_finished fun(): boolean
---@field on_finished fun(callback: fun(...: unknown))
---@field resume vim.task.Callback
---@field __call fun(self: vim.task.YieldingContext, err?: unknown, ...: unknown)

---A function representing a yielding.
---This is an alternative form of yielding that is a pure function.
---@alias vim.task.YieldingFunction fun(ctx: vim.task.YieldingContext)

---A callable object representing a yielding.
---This allows an asynchronous operation to be passed as an object that can be awaited.
---@class vim.task.YieldingCallable
---@field __call fun(self: unknown, ctx: vim.task.YieldingContext)

---The primary asynchronous primitive.
---When called within an `async.spawn` context, it pauses execution until the provided yielding resolves.
---@alias vim.task.Yielding vim.task.YieldingCallable | vim.task.YieldingFunction

---Internal state of an asynchronous task.
---@class vim.task.Task.State
---@field status vim.task.TaskStatus
---@field data? unknown[]
---@field callbacks? table<vim.task.TaskStatus, fun(...: unknown)[]>
---@field detached boolean

---The task object returned by `async.spawn`.
---This object represents an asynchronous operation, allowing users to synchronize and wait for its completion or an error.
---As it extends `vim.task.Yielding`, users can directly `await` a task object, e.g., `async.await(task)`.
---@class vim.task.Task: vim.task.YieldingCallable
---@field co thread
---@field co_parent thread
---@field is_finished fun(): boolean
---@field on_finished fun(callback: fun(...: unknown))
---@field is_detached fun(): boolean
---@field abort fun(reason: string)
---@field detach fun(): vim.task.Task
---@field sync fun(timeout?: integer): unknown

local M = {}

---Internal marker used to identify that a yielded value is an asynchronous yielding.
local YieldingMarker = { 'YieldingMarker' }

---Dispatches status changes for an asynchronous task's context.
---This function updates the task's state and triggers any registered callbacks for the new status.
---@param state vim.task.Task.State
---@param status vim.task.TaskStatus
---@vararg ...
local function dispatch_status(state, status, ...)
  if state.status ~= TaskStatus.pending then
    return
  end
  state.status = status
  state.data = { ... }
  if state.callbacks and state.callbacks[status] then
    for _, callback in ipairs(state.callbacks[status]) do
      if state.status == TaskStatus.success then
        callback(nil, ...)
      else
        callback(...)
      end
    end
    state.callbacks = nil
  end
end

---Registers a callback to listen for specific status changes in a task's context.
---If the status has already been reached, the callback is invoked immediately.
---@param state vim.task.Task.State
---@param status vim.task.TaskStatus
---@param callback fun(...: unknown)
local function listen_status(state, status, callback)
  if state.status == TaskStatus.pending then
    state.callbacks = state.callbacks or {}
    state.callbacks[status] = state.callbacks[status] or {}
    table.insert(state.callbacks[status], callback)
  elseif state.status == status then
    if state.status == TaskStatus.success then
      callback(nil, unpack(state.data))
    else
      callback(unpack(state.data))
    end
  end
end

---Recursive check for floating tasks.
---@param co thread
local function check_floating_tasks(co)
  for _, task in pairs(threads) do
    if task.co_parent == co and not task.is_detached() then
      if not task.is_finished() or check_floating_tasks(task.co) then
        return true
      end
    end
  end
  return false
end

---Executes the next step in the asynchronous coroutine.
---This function handles resuming the coroutine after an `await` or an error, propagating results or errors accordingly.
---@param task vim.task.Task
---@param co thread
---@param ok boolean
---@vararg unknown
local function step(settle, task, co, ok, ...)
  local data = { ... }
  while true do
    local maybe_marker = data[1] --[[@as unknown]]
    local yielding = data[2] --[[@as vim.task.Yielding]]
    if maybe_marker ~= YieldingMarker or not is_callable(yielding) then
      if coroutine.status(co) == 'dead' then
        if ok then
          if check_floating_tasks(co) then
            return settle('Task has finished, but there are still floating tasks.')
          end
          return settle(nil, unpack(data))
        end
        return settle(unpack(data))
      end
      return settle(debug.traceback(co, 'Unexpected coroutine.yield'))
    end

    -- yield.
    local settled = false
    local yield_ok
    local yield_data
    local yok, yerr = pcall(function()
      local sync = true
      ---@type vim.task.YieldingContext
      local ctx = setmetatable({
        is_finished = task.is_finished,
        on_finished = task.on_finished,
        resume = function(err, ...)
          if settled or task.is_finished() then
            return
          end
          settled = true
          if sync then
            yield_ok, yield_data = ok_data(coroutine.resume(co, callback2pcall(err, ...)))
          else
            step(settle, task, co, coroutine.resume(co, callback2pcall(err, ...)))
          end
        end
      }, {
        __call = function(self, err, ...)
          self.resume(err, ...)
        end
      })
      yielding(ctx)
      sync = false
    end)
    if not yok then
      return settle(yerr)
    end
    if yield_ok == nil then
      return
    end
    ok = yield_ok
    data = yield_data
  end
end


---Creates an task and executes a function within it.
---This is the entry point for spawning new tasks.
---@generic T: ...
---@param task_fn fun(...: T): unknown?
---@param ... T
---@return vim.task.Task
function M.spawn(task_fn, ...)
  ---@type vim.task.Task.State
  local state = {
    status = TaskStatus.pending,
    data = nil,
    callbacks = nil,
    detached = false,
  }

  local co = coroutine.create(task_fn)

  ---@type vim.task.Task
  local task
  task = setmetatable({
    co = co,
    co_parent = (coroutine.running()),
    is_finished = function()
      return state.status ~= TaskStatus.pending
    end,
    on_finished = function(callback)
      listen_status(state, TaskStatus.success, callback)
      listen_status(state, TaskStatus.failure, callback)
      listen_status(state, TaskStatus.aborted, callback)
    end,
    is_detached = function()
      return state.detached
    end,
    abort = function(reason)
      threads[co] = nil

      local traceback
      if M.in_context() then
        traceback = debug.traceback((coroutine.running()), reason, 3)
      else
        traceback = debug.traceback(reason, 2)
      end
      dispatch_status(state, TaskStatus.aborted, traceback)
    end,
    sync = function(timeout)
      timeout = timeout or math.huge

      vim.on_key(function(_, typed)
        if typed == interrupt_key then
          task.abort('Keyboard interrupt')
          return ''
        end
      end, interrupt_ns)

      local start_ms = now_ms()
      repeat
        vim.wait(16, function()
          return state.status ~= TaskStatus.pending
        end)
      until (state.status ~= TaskStatus.pending) or (now_ms() - start_ms >= timeout)

      vim.on_key(nil, interrupt_ns)

      if state.status == TaskStatus.pending then
        error('Task.sync has timed out.', 2)
      end
      if state.status == TaskStatus.failure then
        error(state.data[1], 2)
      end
      if state.status == TaskStatus.aborted then
        error(state.data[1], 2)
      end
      return unpack(state.data)
    end,
    detach = function()
      state.detached = true
      return task
    end,
  }, {
    __call = function(self, callback)
      self.on_finished(callback)
    end,
  })

  -- propagate abort.
  local parent_task = threads[(coroutine.running())]
  if parent_task then
    parent_task.on_finished(function(err)
      if not task.is_finished() and not task.is_detached() then
        task.abort(err or 'Parent task has already finished.')
      end
    end)
  end

  threads[co] = task

  step(function(err, ...)
    threads[co] = nil

    if err then
      dispatch_status(state, TaskStatus.failure, err)
    else
      dispatch_status(state, TaskStatus.success, ...)
    end
  end, task, co, coroutine.resume(co, ...))

  return task
end

---Checks if the current coroutine is running within task context.
---@return boolean
function M.in_context()
  return threads[coroutine.running()] ~= nil
end

---Executes an yielding and waits for its resolution.
---This function can only be called from within an task context.
---@async
---@param yielding vim.task.Yielding
---@return unknown
function M.yield(yielding)
  if not M.in_context() then
    error('vim.task.yield can only be called within an task context.', 2)
  end
  local ok, data = ok_data(coroutine.yield(YieldingMarker, yielding))
  if not ok then
    error(data[1], 2)
  end
  return unpack(data)
end

---Awaits the success of all provided tasks.
---The resulting value is an array containing the results of each task in order.
---If any task fails, the entire entire task operation will fail.
---@async
---@param tasks vim.task.Task[]
---@param map_fn? fun(...: any): any
---@return unknown[]
function M.map(tasks, map_fn)
  map_fn = map_fn or function(...) return ... end

  local co = coroutine.running()
  return M.yield(function(ctx)
    if #tasks == 0 then
      return ctx.resume(nil, {})
    end
    local remain = #tasks
    local values = {}
    local settled = false
    for i, task in ipairs(tasks) do
      task(function(err, ...)
        if err and not settled then
          settled = true
          for _, t in ipairs(tasks) do
            if t ~= task and not t.is_detached() and co == t.co_parent then
              t.abort('vim.task.map: remaining tasks aborted after some task failure')
            end
          end
          ctx.resume(err)
          return
        end
        values[i] = map_fn(...)
        remain = remain - 1
        if remain == 0 then
          ctx.resume(nil, values)
        end
      end)
    end
  end)
end

---Awaits the first task to resolve among the provided tasks.
---Returns the result of the first task that successfully completes or errors.
---Subsequent completions/errors from other tasks are ignored.
---@async
---@param tasks vim.task.Task[]
---@return unknown
function M.any(tasks)
  local co = (coroutine.running())
  return M.yield(function(ctx)
    if #tasks == 0 then
      return ctx.resume(nil)
    end
    local settled = false
    for _, task in ipairs(tasks) do
      task(function(err, ...)
        if settled then
          return
        end
        settled = true

        for _, t in ipairs(tasks) do
          if t ~= task and not t.is_detached() and co == t.co_parent then
            t.abort('vim.task.any: remaining tasks aborted after winner finished')
          end
        end
        ctx.resume(err, ...)
      end)
    end
  end)
end

---Awaits for a specified duration, creating a new timeout context.
---This function pauses the current task for the given number of milliseconds.
---@async
---@param timeout integer # The duration to wait in milliseconds.
function M.timeout(timeout)
  M.yield(function(ctx)
    local timer = assert(vim.uv.new_timer())
    ctx.on_finished(function()
      timer:stop()
      timer:close()
    end)
    timer:start(timeout, 0, function()
      if ctx.is_finished() then
        return
      end
      vim.schedule(ctx.resume)
    end)
  end)
end

---Schedules the yielding to run in the next event loop iteration.
---This effectively yields control to the Neovim event loop and resumes the task in the immediate future.
---@async
function M.schedule()
  M.yield(function(ctx)
    vim.schedule(ctx.resume)
  end)
end

vim.task = M

------------------------------------------------------------
---↓↓↓↓ playground ↓↓↓↓
------------------------------------------------------------

---Gets file system statistics for a given path asynchronously.
---@param path string # The file path.
---@return uv.fs_stat.result
local function fs_stat(path)
  return vim.task.yield(function(ctx)
    vim.uv.fs_stat(path, ctx.resume)
  end)
end

local print_with_time = (function()
  local s = now_ms()
  return function(...)
    local e = now_ms()
    vim.print(string.format('[%sms] ', e - s) .. vim.inspect(...))
    s = e
  end
end)()

---@param name string
---@param fn fun(): vim.task.Task
---@param expects { ok: boolean, match?: string }
local function playground(name, fn, expects)
  vim.print('\n')
  print_with_time('--- ' .. name .. ' ---')
  local output = { pcall(fn().sync, 10 * 1000) }
  print_with_time(vim.inspect(output))
  assert(output[1] == expects.ok, vim.inspect(output[2]))
  if expects.match then
    assert(string.match(vim.inspect(output[2]), expects.match),
      ('got: %s, wants: %s'):format(vim.inspect(output[2]), expects.match))
  end
  vim.print('\n')
end

playground('usage: timeout', function()
  return vim.task.spawn(function()
    vim.task.timeout(100)
    return 'Hello!'
  end)
end, {
  ok = true,
  match = 'Hello!',
})

playground('usage: fs_stat', function()
  return vim.task.spawn(function()
    return fs_stat(vim.fs.normalize('~/Develop/Repo/dotfiles/dot_config/nvim/init.lua'))
  end)
end, {
  ok = true,
  match = 'type = "file"',
})

playground('with await sync error', function()
  return vim.task.spawn(function()
    vim.task.yield(function()
      error('An error occurred.')
    end)
  end)
end, {
  ok = false,
  match = 'An error occurred.',
})

playground('with task async error', function()
  return vim.task.spawn(function()
    vim.task.timeout(100)
    error('An error occurred.')
  end)
end, {
  ok = false,
  match = 'An error occurred.',
})

playground('with task sync error', function()
  return vim.task.spawn(function()
    error('An error occurred.')
  end)
end, {
  ok = false,
  match = 'An error occurred.',
})

playground('with abort', function()
  local task = vim.task.spawn(function()
    vim.task.timeout(100)
    error('do not reach here')
  end)
  task.abort('abort by parent')
  return task
end, {
  ok = false,
  match = 'abort by parent',
})

playground('vim.task.map', function()
  return vim.task.spawn(function()
    print_with_time('Running tasks...')
    vim.task.map({
      vim.task.spawn(vim.task.timeout, 100),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
    print_with_time('Completed tasks.')
  end)
end, {
  ok = true,
})

playground('vim.task.map: early error', function()
  return vim.task.spawn(function()
    print_with_time('Running tasks...')
    vim.task.map({
      vim.task.yield(function(ctx)
        ctx.resume('Early error')
      end),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
    print_with_time('Completed tasks.')
  end)
end, {
  ok = false,
  match = 'Early error',
})

playground('vim.task.any', function()
  return vim.task.spawn(function()
    print_with_time('Running tasks...')
    vim.task.any({
      vim.task.spawn(vim.task.timeout, 100),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
    print_with_time('Completed tasks.')
  end)
end, {
  ok = true,
})

playground('vim.task.any: early error', function()
  return vim.task.spawn(function()
    print_with_time('Running tasks...')
    vim.task.any({
      vim.task.yield(function(ctx)
        ctx.resume('Early error')
      end),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
    print_with_time('Completed tasks.')
  end)
end, {
  ok = false,
  match = 'Early error',
})

playground('dependencies: dependent children task should be finished when parent is finished', function()
  local timeout
  local task = vim.task.spawn(function()
    timeout = vim.task.spawn(vim.task.timeout, 200)
    vim.task.yield(timeout)
  end)
  task.abort('abort')
  assert(timeout.is_finished() == true)
  return task
end, {
  ok = false,
  match = 'abort',
})

playground('dependencies: detached children task should not be finished when parent is finished', function()
  local timeout
  local task = vim.task.spawn(function()
    timeout = vim.task.spawn(vim.task.timeout, 200).detach()
    vim.task.yield(timeout)
  end)
  task.abort('abort')
  assert(timeout.is_finished() == false)
  return task
end, {
  ok = false,
  match = 'abort',
})

playground('dependencies: parent task should be failure if it has floating sub-tasks', function()
  local task = vim.task.spawn(function()
    vim.task.spawn(vim.task.timeout, 200) -- not joined
  end)
  return task
end, {
  ok = false,
  match = 'Task has finished, but there are still floating tasks.',
})

playground('dependencies: parent task should not be failure if it has floating sub-tasks but detached', function()
  local task = vim.task.spawn(function()
    vim.task.spawn(vim.task.timeout, 200).detach() -- not joined
  end)
  return task
end, {
  ok = true,
})

playground('interrupt by <C-c>', function()
  return vim.task.spawn(function()
    vim.task.timeout(math.huge)
  end)
end, {
  ok = false,
  match = 'Keyboard interrupt',
})

playground('coroutine.yield is forbidden', function()
  return vim.task.spawn(function()
    coroutine.yield('This will cause an error.')
  end)
end, {
  ok = false,
  match = 'Unexpected coroutine.yield',
})

playground('cleanup threads table on success status', function()
  local task = vim.task.spawn(function()
    vim.task.map({
      vim.task.spawn(vim.task.timeout, 100),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
  end)
  task.sync()
  assert(#vim.tbl_keys(threads) == 0, 'Threads table should be empty after abort')
  return task
end, {
  ok = true,
})

playground('cleanup threads table on failure status', function()
  local task = vim.task.spawn(function()
    vim.task.map({
      vim.task.spawn(vim.task.timeout, 100),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
    error('An error occurred.')
  end)
  pcall(task.sync)
  assert(#vim.tbl_keys(threads) == 0, 'Threads table should be empty after abort')
  return task
end, {
  ok = false,
  match = 'An error occurred.',
})

playground('cleanup threads table on aborted status', function()
  local task = vim.task.spawn(function()
    vim.task.map({
      vim.task.spawn(vim.task.timeout, 100),
      vim.task.spawn(vim.task.timeout, 200),
      vim.task.spawn(vim.task.timeout, 300),
    })
  end)
  task.abort('abort')
  assert(#vim.tbl_keys(threads) == 0, 'Threads table should be empty after abort')
  return task
end, {
  ok = false,
  match = 'abort',
})

playground('sync yield does not need new stack frame', function()
  local function deep(n)
    if n == 0 then
      return 'done'
    end
    vim.task.yield(function(ctx)
      ctx.resume(nil)
    end)
    return deep(n - 1)
  end
  return vim.task.spawn(function()
    return deep(10000)
  end)
end, {
  ok = true,
  match = 'done',
})