async.lua

local co = coroutine
local vararg = require "plenary.vararg"
local errors = require "plenary.errors"
local traceback_error = errors.traceback_error
local f = require "plenary.functional"

---@class PlenaryAsyncAsync
local M = {}

---@param fn any
---@return boolean
local function is_callable(fn)
  return type(fn) == "function" or (type(fn) == "table" and type(getmetatable(fn)["__call"]) == "function")
end

---because we can't store varargs
---@param step function
---@param thread thread
---@param callback? function
---@param ... any
local function callback_or_next(step, thread, callback, ...)
  local stat = f.first(...)

  if not stat then
    error(string.format("The coroutine failed with this message: %s", f.second(...)))
  end

  if co.status(thread) == "dead" then
    if callback == nil then
      return
    end
    callback(select(2, ...))
  else
    local returned_function = f.second(...)
    local nargs = f.third(...)

    assert(is_callable(returned_function), "type error :: expected func")
    returned_function(vararg.rotate(nargs, step, select(4, ...)))
  end
end

---Executes a future with a callback when it is done
---@param func function the future to execute
---@param callback? function the callback to call when done
local execute = function(func, callback, ...)
  assert(is_callable(func), "type error :: expected func")

  local thread = co.create(func)

  local step
  step = function(...)
    callback_or_next(step, thread, callback, co.resume(thread, ...))
  end

  step(...)
end

---A function including async logic
---@alias PlenaryAsyncFunction async fun(...): ...

local add_leaf_function
do
  ---A table to store all leaf async functions
  ---@type table<PlenaryAsyncFunction, integer>
  _PlenaryLeafTable = setmetatable({}, {
    __mode = "k",
  })

  ---@param async_func PlenaryAsyncFunction
  ---@param argc integer
  add_leaf_function = function(async_func, argc)
    assert(_PlenaryLeafTable[async_func] == nil, "Async function should not already be in the table")
    _PlenaryLeafTable[async_func] = argc
  end

  ---@param async_func PlenaryAsyncFunction
  ---@return boolean
  function M.is_leaf_function(async_func)
    return _PlenaryLeafTable[async_func] ~= nil
  end

  ---@param async_func PlenaryAsyncFunction
  ---@return integer
  function M.get_leaf_function_argc(async_func)
    return _PlenaryLeafTable[async_func]
  end
end

---Creates an async function with a callback style function.
---@param func function A callback style function to be converted. The last argument must be the callback.
---@param argc integer The number of arguments of func. Must be included.
---@return PlenaryAsyncFunction leaf Returns an leaf
M.wrap = function(func, argc)
  if not is_callable(func) then
    traceback_error("type error :: expected func, got " .. type(func))
  end

  if type(argc) ~= "number" then
    traceback_error("type error :: expected number, got " .. type(argc))
  end

  ---@type PlenaryAsyncFunction
  local function leaf(...)
    local nargs = select("#", ...)

    if nargs == argc then
      return func(...)
    else
      return co.yield(func, argc, ...)
    end
  end

  add_leaf_function(leaf, argc)

  return leaf
end

---Use this to either run a future concurrently and then do something else
---or use it to run a future with a callback in a non async context
---@param func PlenaryAsyncFunction|function
---@param callback function
M.run = function(func, callback)
  if M.is_leaf_function(func) then
    func(callback)
  else
    execute(func, callback)
  end
end

---Use this to create a function which executes in an async context but
---called from a non-async context. Inherently this cannot return anything
---since it is non-blocking
---@param func function
M.void = function(func)
  return function(...)
    execute(func, nil, ...)
  end
end

return M