From 4e35b30224171d475d25f97e9384d61ec60983bd Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Sun, 27 Nov 2022 08:46:50 +0000
Subject: [PATCH 01/23] Update README.md

---
 README.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index b77adc6..54621d6 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,8 @@
 # async.nvim
 Small async library for Neovim plugins
 
+[API documentation](async.md)
+
 ## Example
 
 Take the current function that uses a callback style function to run a system process.
@@ -19,8 +21,8 @@ end
 
 If we want to emulate something like:
 
-```lua
-echo foo && echo bar && echo baz
+```bash
+echo 'foo' && echo 'bar' && echo 'baz'
 ```
 
 Would need to be implemented as:

From e261f3bfaf3011a3a500b69fbe1c33b96d225cd5 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Sun, 27 Nov 2022 08:48:32 +0000
Subject: [PATCH 02/23] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 54621d6..6c9c1d3 100644
--- a/README.md
+++ b/README.md
@@ -62,7 +62,7 @@ local run_job_a = a.wrap(run_job, 3)
 
 Now we need to create a top level function to initialize the async context. To do this we can use `void` or `sync`.
 
-Note: the main difference between `void` and `sync` is that `sync` functions can be called with a callback (like the `run_job` in a non-async context, however the user must provide the number of agurments.
+Note: the main difference between `void` and `sync` is that `sync` functions can be called with a callback, like the original `run_job` in a non-async context, however the user must provide the number of agurments.
 
 For this example we will use `void`:
 

From 9bd4cbb69af3397b432c8274c36d30598b494c16 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Sun, 27 Nov 2022 08:48:50 +0000
Subject: [PATCH 03/23] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 6c9c1d3..6b83a31 100644
--- a/README.md
+++ b/README.md
@@ -62,7 +62,7 @@ local run_job_a = a.wrap(run_job, 3)
 
 Now we need to create a top level function to initialize the async context. To do this we can use `void` or `sync`.
 
-Note: the main difference between `void` and `sync` is that `sync` functions can be called with a callback, like the original `run_job` in a non-async context, however the user must provide the number of agurments.
+Note: the main difference between `void` and `sync` is that `sync` functions can be called with a callback, like the original `run_job` in a non-async context, however the user must provide the number of arguments.
 
 For this example we will use `void`:
 

From 1e03fe1d45b009a0df3226426478762bab694d72 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Fri, 6 Jan 2023 09:58:52 +0000
Subject: [PATCH 04/23] rename sync to create

---
 async.md      | 2 +-
 lua/async.lua | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/async.md b/async.md
index d1864cf..b390fb4 100644
--- a/async.md
+++ b/async.md
@@ -5,7 +5,7 @@ Small async library for Neovim plugins
 
 ## Functions
 
-### `sync(func, argc)`
+### `create(func, argc)`
 
 Use this to create a function which executes in an async context but
  called from a non-async context.  Inherently this cannot return anything
diff --git a/lua/async.lua b/lua/async.lua
index 01c089f..d8aa62f 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -58,7 +58,7 @@ end
 --- since it is non-blocking
 --- @tparam function func
 --- @tparam number argc The number of arguments of func. Defaults to 0
-function M.sync(func, argc)
+function M.create(func, argc)
   argc = argc or 0
   return function(...)
     if coroutine.running() ~= main_co_or_nil then

From 7dc08cda69e8e6ea80cc1678c6bd6be90f59c74e Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Fri, 6 Jan 2023 10:15:47 +0000
Subject: [PATCH 05/23] feat: add running()

---
 async.md      |  6 ++++++
 lua/async.lua | 27 ++++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/async.md b/async.md
index b390fb4..1fa3d63 100644
--- a/async.md
+++ b/async.md
@@ -5,6 +5,12 @@ Small async library for Neovim plugins
 
 ## Functions
 
+### `running()`
+
+Returns whether the current execution context is async.
+
+
+---
 ### `create(func, argc)`
 
 Use this to create a function which executes in an async context but
diff --git a/lua/async.lua b/lua/async.lua
index d8aa62f..a611716 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -1,9 +1,13 @@
 --- Small async library for Neovim plugins
 --- @module async
 
+-- Store all the async threads in a weak table so we don't prevent them from
+-- being garbage collected
+local threads = setmetatable({}, { __mode = 'kv' })
+
 local M = {}
 
--- Coroutine.running() was changed between Lua 5.1 and 5.2:
+-- Note: coroutine.running() was changed between Lua 5.1 and 5.2:
 -- - 5.1: Returns the running coroutine, or nil when called by the main thread.
 -- - 5.2: Returns the running coroutine plus a boolean, true when the running
 --   coroutine is the main one.
@@ -11,10 +15,20 @@ local M = {}
 -- For LuaJIT, 5.2 behaviour is enabled with LUAJIT_ENABLE_LUA52COMPAT
 --
 -- We need to handle both.
-local main_co_or_nil = coroutine.running()
+
+--- Returns whether the current execution context is async.
+---
+--- @return boolean|nil
+function M.running()
+  local current = coroutine.running()
+  if current and threads[current] then
+    return true
+  end
+end
 
 local function execute(func, callback, ...)
   local co = coroutine.create(func)
+  threads[co] = true
 
   local function step(...)
     local ret = {coroutine.resume(co, ...)}
@@ -61,7 +75,7 @@ end
 function M.create(func, argc)
   argc = argc or 0
   return function(...)
-    if coroutine.running() ~= main_co_or_nil then
+    if M.running() then
       return func(...)
     end
     local callback = select(argc+1, ...)
@@ -74,7 +88,7 @@ end
 --- @tparam function func
 function M.void(func)
   return function(...)
-    if coroutine.running() ~= main_co_or_nil then
+    if M.running() then
       return func(...)
     end
     execute(func, nil, ...)
@@ -89,7 +103,7 @@ end
 function M.wrap(func, argc, protected)
   assert(argc)
   return function(...)
-    if coroutine.running() == main_co_or_nil then
+    if not M.running() then
       return func(...)
     end
     return coroutine.yield(argc, protected, func, ...)
@@ -130,6 +144,9 @@ function M.join(n, interrupt_check, thunks)
     end
   end
 
+  if not M.running() then
+    return run
+  end
   return coroutine.yield(1, false, run)
 end
 

From 5a60d514a5201c360ba9938f2673120c2ec257ea Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Fri, 6 Jan 2023 11:41:13 +0000
Subject: [PATCH 06/23] feat: add async handles

with methods cancel() and is_cancelled()
---
 lua/async.lua | 61 ++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 51 insertions(+), 10 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index a611716..c912abf 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -3,7 +3,7 @@
 
 -- Store all the async threads in a weak table so we don't prevent them from
 -- being garbage collected
-local threads = setmetatable({}, { __mode = 'kv' })
+local handles = setmetatable({}, { __mode = 'k' })
 
 local M = {}
 
@@ -21,25 +21,57 @@ local M = {}
 --- @return boolean|nil
 function M.running()
   local current = coroutine.running()
-  if current and threads[current] then
+  if current and handles[current] then
     return true
   end
 end
 
-local function execute(func, callback, ...)
+local Handle = {}
+
+function Handle.new(func)
   local co = coroutine.create(func)
-  threads[co] = true
+  local handle setmetatable({ co = co }, { __index = {
+    cancel = function(self, cb)
+      self.cancelled_cb = cb or function() end
+    end
+  }})
+  handles[co] = handle
+  return handle
+end
+
+-- Analogous to uv.close
+function Handle:cancel(callback)
+  vim.validate{ callback = { callback , 'function', true } }
+  self.cancelled_cb = callback or function() end
+end
+
+-- Analogous to uv.is_closing
+function Handle:is_cancelled()
+  return self.cancelled_cb == nil
+end
+
+local function execute(func, callback, ...)
+  vim.validate{
+    func = { func, 'function' },
+    callback = { callback , 'function', true }
+  }
+  local handle = Handle.new(func)
 
   local function step(...)
-    local ret = {coroutine.resume(co, ...)}
+    if handle.cancelled_cb then
+      handle.cancelled_cb()
+      return
+    end
+
+    local ret = {coroutine.resume(handle.co, ...)}
     local stat, nargs, protected, err_or_fn = unpack(ret)
 
     if not stat then
       error(string.format("The coroutine failed with this message: %s\n%s",
-        err_or_fn, debug.traceback(co)))
+        err_or_fn, debug.traceback(handle.co)))
     end
 
-    if coroutine.status(co) == 'dead' then
+    if coroutine.status(handle.co) == 'dead' then
       if callback then
         callback(unpack(ret, 4))
       end
@@ -65,6 +97,7 @@ local function execute(func, callback, ...)
   end
 
   step(...)
+  return handle
 end
 
 --- Use this to create a function which executes in an async context but
@@ -73,13 +106,17 @@ end
 --- @tparam function func
 --- @tparam number argc The number of arguments of func. Defaults to 0
 function M.create(func, argc)
+  vim.validate{
+    func = { func , 'function' },
+    argc = { argc, 'number', true }
+  }
   argc = argc or 0
   return function(...)
     if M.running() then
       return func(...)
     end
     local callback = select(argc+1, ...)
-    execute(func, callback, unpack({...}, 1, argc))
+    return execute(func, callback, unpack({...}, 1, argc))
   end
 end
 
@@ -87,11 +124,12 @@ end
 --- called from a non-async context.
 --- @tparam function func
 function M.void(func)
+  vim.validate{ func = { func , 'function' } }
   return function(...)
     if M.running() then
       return func(...)
     end
-    execute(func, nil, ...)
+    return execute(func, nil, ...)
   end
 end
 
@@ -101,7 +139,10 @@ end
 --- @tparam boolean protected call the function in protected mode (like pcall)
 --- @return function Returns an async function
 function M.wrap(func, argc, protected)
-  assert(argc)
+  vim.validate{
+    argc = { argc, 'number' },
+    protected = { protected, 'boolean', true }
+  }
   return function(...)
     if not M.running() then
       return func(...)

From 62935d4f01d1617f198e459930de7ec908089bd1 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Fri, 6 Jan 2023 15:49:47 +0000
Subject: [PATCH 07/23] feat: improve canceling

---
 lua/async.lua | 78 +++++++++++++++++++++++++++++----------------------
 1 file changed, 45 insertions(+), 33 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index c912abf..191a378 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -26,28 +26,10 @@ function M.running()
   end
 end
 
-local Handle = {}
-
-function Handle.new(func)
-  local co = coroutine.create(func)
-  local handle setmetatable({ co = co }, { __index = {
-    cancel = function(self, cb)
-      self.cancelled_cb = cb or function() end
-    end
-  }})
-  handles[co] = handle
-  return handle
-end
-
--- Analogous to uv.close
-function Handle:cancel(callback)
-  vim.validate{ callback = { callback , 'function', true } }
-  self.cancelled_cb = callback or function() end
-end
-
--- Analogous to uv.is_closing
-function Handle:is_cancelled()
-  return self.cancelled_cb == nil
+local function is_async_handle(handle)
+  if handle and handle.cancel and handle.is_cancelled then
+    return true
+  end
 end
 
 local function execute(func, callback, ...)
@@ -55,23 +37,51 @@ local function execute(func, callback, ...)
     func = { func, 'function' },
     callback = { callback , 'function', true }
   }
-  local handle = Handle.new(func)
 
-  local function step(...)
-    if handle.cancelled_cb then
-      handle.cancelled_cb()
-      return
+  local co = coroutine.create(func)
+
+  -- Handle for an object currently running on the event loop.
+  -- The coroutine is paused while this is active.
+  -- Must provide methods cancel() and is_cancelled()
+  local cur_exec_handle
+
+  -- Handle for the user. Since cur_exec_handle will change every
+  -- step() we need to provide access to it through a proxy
+  local handle = {}
+
+  -- Analogous to uv.close
+  function handle:cancel(cb)
+    vim.validate{ callback = { cb , 'function', true } }
+    -- Cancel anything running on the event loop
+    if cur_exec_handle and not cur_exec_handle:is_cancelled() then
+      cur_exec_handle:cancel(cb)
+    end
+  end
+
+  -- Analogous to uv.is_closing
+  function handle:is_cancelled()
+    return cur_exec_handle and cur_exec_handle:is_cancelled()
+  end
+
+  local function set_executing_handle(h)
+    if is_async_handle(h) then
+      cur_exec_handle = h
     end
+  end
 
-    local ret = {coroutine.resume(handle.co, ...)}
+  setmetatable(handle, { __index = handle })
+  handles[co] = handle
+
+  local function step(...)
+    local ret = {coroutine.resume(co, ...)}
     local stat, nargs, protected, err_or_fn = unpack(ret)
 
     if not stat then
       error(string.format("The coroutine failed with this message: %s\n%s",
-        err_or_fn, debug.traceback(handle.co)))
+        err_or_fn, debug.traceback(co)))
     end
 
-    if coroutine.status(handle.co) == 'dead' then
+    if coroutine.status(co) == 'dead' then
       if callback then
         callback(unpack(ret, 4))
       end
@@ -86,13 +96,15 @@ local function execute(func, callback, ...)
       args[nargs] = function(...)
         step(true, ...)
       end
-      local ok, err = pcall(err_or_fn, unpack(args, 1, nargs))
+      local ok, err_or_handle = pcall(err_or_fn, unpack(args, 1, nargs))
       if not ok then
-        step(false, err)
+        step(false, err_or_handle)
+      else
+        set_executing_handle(err_or_handle)
       end
     else
       args[nargs] = step
-      err_or_fn(unpack(args, 1, nargs))
+      set_executing_handle(err_or_fn(unpack(args, 1, nargs)))
     end
   end
 

From 9de7bda630cadb9feca42ea33e918a568231d060 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Fri, 6 Jan 2023 16:05:13 +0000
Subject: [PATCH 08/23] doc: add example for cancelling

---
 README.md | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/README.md b/README.md
index 6b83a31..55ed574 100644
--- a/README.md
+++ b/README.md
@@ -86,3 +86,45 @@ main()
 
 We can now call `run_job_a` in linear imperative fashion without needing to define callbacks.
 The arguments provided to the callback in the original function are simply returned by the async version.
+
+## The `async_t` handle
+
+This library supports cancelling async functions that are currently running. This is done via the `async_t` handle interface.
+The handle must provide the methods `cancel()` and `is_cancelled()`, and the purpose of these is to allow the cancelled async function to run any cleanup and free any resources it has created.
+
+### Example use with `vim.loop.spawn`:
+
+Typically applications to `vim.loop.spawn` make use of `stdio` pipes for communicating. This involves creating `uv_pipe_t` objects.
+If a job is cancelled then these objects must be closed.
+
+```lua
+local function run_job = async.wrap(function(cmd, args, callback)
+  local stdout = vim.loop.new_pipe(false)
+
+  local raw_handle
+  raw_handle = vim.loop.spawn(cmd, { args  = args, stdio = { nil, stdout }},
+    function(code)
+      stdout:close()
+      raw_handle:close()
+      callback(code)
+    end
+  )
+
+  local handle = {}
+
+  handle.is_cancelled = function(_)
+    return raw_handle.is_closing()
+  end
+
+  handle.cancel = function(_, cb)
+    raw_handle:close(function()
+        stdout:close(cb)
+    end)
+  end
+
+  return handle
+end)
+```
+
+So even if `run_job` is called in a deep function stack, calling `cancel()` on any parent async function will allow the job to be cancelled safely.
+

From 6f2bea24e5b800b0268d35a636878eb5ca7c68e0 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Fri, 6 Jan 2023 16:31:36 +0000
Subject: [PATCH 09/23] reorder arguments to join()

---
 async.md      | 4 ++--
 lua/async.lua | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/async.md b/async.md
index 1fa3d63..98e0525 100644
--- a/async.md
+++ b/async.md
@@ -44,16 +44,16 @@ Creates an async function with a callback style function.
 * `protected` (`boolean`):  call the function in protected mode (like pcall)
 
 ---
-### `join(n, interrupt_check, thunks)`
+### `join(thunks, n, interrupt_check)`
 
 Run a collection of async functions (`thunks`) concurrently and return when
  all have finished.
 
 #### Parameters:
 
+* `thunks` (`function[]`):
 * `n` (`integer`):  Max number of thunks to run concurrently
 * `interrupt_check` (`function`):  Function to abort thunks between calls
-* `thunks` (`function[]`):
 
 ---
 ### `curry(fn, ...)`
diff --git a/lua/async.lua b/lua/async.lua
index 191a378..e3c3a82 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -165,10 +165,10 @@ end
 
 --- Run a collection of async functions (`thunks`) concurrently and return when
 --- all have finished.
+--- @tparam function[] thunks
 --- @tparam integer n Max number of thunks to run concurrently
 --- @tparam function interrupt_check Function to abort thunks between calls
---- @tparam function[] thunks
-function M.join(n, interrupt_check, thunks)
+function M.join(thunks, n, interrupt_check )
   local function run(finish)
     if #thunks == 0 then
       return finish()

From f2104a5da56755cc034f2954b20fda4f31cf2073 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 12:21:07 +0000
Subject: [PATCH 10/23] feat: Add async.run and async.wait

- Add strict option to `void`, `create` and `wrap`
---
 async.md      | 35 +++++++++++++++++++++--
 lua/async.lua | 79 +++++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 102 insertions(+), 12 deletions(-)

diff --git a/async.md b/async.md
index 98e0525..33d0789 100644
--- a/async.md
+++ b/async.md
@@ -11,7 +11,30 @@ Returns whether the current execution context is async.
 
 
 ---
-### `create(func, argc)`
+### `run(func, callback, ...)`
+
+Run a function in an async context.
+
+#### Parameters:
+
+* `func` (`function`):
+* `callback` (`function`):
+* `...` (`any`):  Arguments for func
+
+---
+### `wait(argc, protected, func, ...)`
+
+Wait on a callback style function
+
+#### Parameters:
+
+* `argc` (`integer?`):  The number of arguments of func. Must be included.
+* `protected` (`boolean?`):  call the function in protected mode (like pcall)
+* `func` (`function`):  callback style function to execute
+* `...` (`any`):  Arguments for func
+
+---
+### `create(func, argc, strict)`
 
 Use this to create a function which executes in an async context but
  called from a non-async context.  Inherently this cannot return anything
@@ -21,9 +44,10 @@ Use this to create a function which executes in an async context but
 
 * `func` (`function`):
 * `argc` (`number`):  The number of arguments of func. Defaults to 0
+* `strict` (`boolean`):  Error when called in non-async context
 
 ---
-### `void(func)`
+### `void(func, strict)`
 
 Create a function which executes in an async context but
  called from a non-async context.
@@ -31,17 +55,22 @@ Create a function which executes in an async context but
 #### Parameters:
 
 * `func` (`function`):
+* `strict` (`boolean`):  Error when called in non-async context
 
 ---
-### `wrap(func, argc, protected)`
+### `wrap(func, argc, protected, strict)`
 
 Creates an async function with a callback style function.
 
+ TODO(lewis6991): Remove protected
+
+
 #### Parameters:
 
 * `func` (`function`):  A callback style function to be converted. The last argument must be the callback.
 * `argc` (`integer`):  The number of arguments of func. Must be included.
 * `protected` (`boolean`):  call the function in protected mode (like pcall)
+* `strict` (`boolean`):  Error when called in non-async context
 
 ---
 ### `join(thunks, n, interrupt_check)`
diff --git a/lua/async.lua b/lua/async.lua
index e3c3a82..11911b8 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -32,7 +32,11 @@ local function is_async_handle(handle)
   end
 end
 
-local function execute(func, callback, ...)
+--- Run a function in an async context.
+--- @tparam function func
+--- @tparam function callback
+--- @tparam any ... Arguments for func
+function M.run(func, callback, ...)
   vim.validate{
     func = { func, 'function' },
     callback = { callback , 'function', true }
@@ -112,54 +116,111 @@ local function execute(func, callback, ...)
   return handle
 end
 
+local function wait(argc, func, protected, ...)
+  vim.validate{
+    argc = { argc, 'number'},
+    func = { func, 'function' },
+    protected = { protected, 'boolean' },
+  }
+  return coroutine.yield(argc, func, protected, ...)
+end
+
+--- Wait on a callback style function
+---
+--- @tparam integer? argc The number of arguments of func. Must be included.
+--- @tparam boolean? protected call the function in protected mode (like pcall)
+--- @tparam function func callback style function to execute
+--- @tparam any ... Arguments for func
+function M.wait(...)
+  local nargs, args = select('#', ...), {...}
+
+  local argstart
+  local protected = false
+  local argc
+  for i = 1, math.min(nargs, 3) do
+    if type(args[i]) == 'function' then
+      argstart = i + 1
+      break
+    end
+    if type(args[i]) == 'boolean' then
+      protected = args[i]
+    end
+    if type(args[i]) == 'number' then
+      argc = args[i]
+    end
+  end
+
+  assert(argstart)
+
+  local func = args[argstart - 1]
+  argc = argc or nargs - argstart + 1
+
+  return wait(argc, func, protected, unpack(args, argstart, nargs))
+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
 --- @tparam function func
 --- @tparam number argc The number of arguments of func. Defaults to 0
-function M.create(func, argc)
+--- @tparam boolean strict Error when called in non-async context
+function M.create(func, argc, strict)
   vim.validate{
-    func = { func , 'function' },
+    func = { func, 'function' },
     argc = { argc, 'number', true }
   }
   argc = argc or 0
   return function(...)
     if M.running() then
+      if strict then
+        error('This function must run in a non-async context')
+      end
       return func(...)
     end
     local callback = select(argc+1, ...)
-    return execute(func, callback, unpack({...}, 1, argc))
+    return M.run(func, callback, unpack({...}, 1, argc))
   end
 end
 
 --- Create a function which executes in an async context but
 --- called from a non-async context.
 --- @tparam function func
-function M.void(func)
+--- @tparam boolean strict Error when called in non-async context
+function M.void(func, strict)
   vim.validate{ func = { func , 'function' } }
   return function(...)
     if M.running() then
+      if strict then
+        error('This function must run in a non-async context')
+      end
       return func(...)
     end
-    return execute(func, nil, ...)
+    return M.run(func, nil, ...)
   end
 end
 
 --- Creates an async function with a callback style function.
+---
+--- TODO(lewis6991): Remove protected
+---
 --- @tparam function func A callback style function to be converted. The last argument must be the callback.
 --- @tparam integer argc The number of arguments of func. Must be included.
 --- @tparam boolean protected call the function in protected mode (like pcall)
+--- @tparam boolean strict Error when called in non-async context
 --- @return function Returns an async function
-function M.wrap(func, argc, protected)
+function M.wrap(func, argc, protected, strict)
   vim.validate{
     argc = { argc, 'number' },
     protected = { protected, 'boolean', true }
   }
   return function(...)
     if not M.running() then
+      if strict then
+        error('This function must run in an async context')
+      end
       return func(...)
     end
-    return coroutine.yield(argc, protected, func, ...)
+    return M.wait(argc, protected, func, ...)
   end
 end
 
@@ -200,7 +261,7 @@ function M.join(thunks, n, interrupt_check )
   if not M.running() then
     return run
   end
-  return coroutine.yield(1, false, run)
+  return M.wait(1, false, run)
 end
 
 --- Partially applying arguments to an async function

From 56be481d78add1033fabe0f6fc08917f15ba29aa Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 12:49:56 +0000
Subject: [PATCH 11/23] doc: add return types

---
 async.md      | 16 ++++++++++++++++
 ldoc.ltp      | 19 +++++++++++++++++++
 lua/async.lua |  6 ++++--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/async.md b/async.md
index 33d0789..5ef5c15 100644
--- a/async.md
+++ b/async.md
@@ -10,6 +10,10 @@ Small async library for Neovim plugins
 Returns whether the current execution context is async.
 
 
+#### Returns
+
+  `boolean?`
+
 ---
 ### `run(func, callback, ...)`
 
@@ -21,6 +25,10 @@ Run a function in an async context.
 * `callback` (`function`):
 * `...` (`any`):  Arguments for func
 
+#### Returns
+
+  `async_t`: Handle
+
 ---
 ### `wait(argc, protected, func, ...)`
 
@@ -46,6 +54,10 @@ Use this to create a function which executes in an async context but
 * `argc` (`number`):  The number of arguments of func. Defaults to 0
 * `strict` (`boolean`):  Error when called in non-async context
 
+#### Returns
+
+  `function(...):async_t`
+
 ---
 ### `void(func, strict)`
 
@@ -72,6 +84,10 @@ Creates an async function with a callback style function.
 * `protected` (`boolean`):  call the function in protected mode (like pcall)
 * `strict` (`boolean`):  Error when called in non-async context
 
+#### Returns
+
+  `function`: Returns an async function
+
 ---
 ### `join(thunks, n, interrupt_check)`
 
diff --git a/ldoc.ltp b/ldoc.ltp
index 49d5f73..68f1e15 100644
--- a/ldoc.ltp
+++ b/ldoc.ltp
@@ -43,6 +43,25 @@ $(lev3) $(subnames):
 >         end
 >       end -- for
 >     end -- if params
+> if item.retgroups then
+
+$(lev3) Returns
+
+>   for _, group in ldoc.ipairs(item.retgroups) do
+>     for r in group:iter() do
+>       local type, ctypes = item:return_type(r)
+>       if type ~= '' then
+>         if r.text ~= '' then
+  `$(type)`: $(r.text)
+>         else
+  `$(type)`
+>         end
+>       else
+  $(r.text)
+>       end
+>     end
+>   end
+> end
 
 ---
 >   end
diff --git a/lua/async.lua b/lua/async.lua
index 11911b8..4230bd5 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -18,7 +18,7 @@ local M = {}
 
 --- Returns whether the current execution context is async.
 ---
---- @return boolean|nil
+--- @treturn boolean?
 function M.running()
   local current = coroutine.running()
   if current and handles[current] then
@@ -36,6 +36,7 @@ end
 --- @tparam function func
 --- @tparam function callback
 --- @tparam any ... Arguments for func
+--- @treturn async_t Handle
 function M.run(func, callback, ...)
   vim.validate{
     func = { func, 'function' },
@@ -164,6 +165,7 @@ end
 --- @tparam function func
 --- @tparam number argc The number of arguments of func. Defaults to 0
 --- @tparam boolean strict Error when called in non-async context
+--- @treturn function(...):async_t
 function M.create(func, argc, strict)
   vim.validate{
     func = { func, 'function' },
@@ -207,7 +209,7 @@ end
 --- @tparam integer argc The number of arguments of func. Must be included.
 --- @tparam boolean protected call the function in protected mode (like pcall)
 --- @tparam boolean strict Error when called in non-async context
---- @return function Returns an async function
+--- @treturn function Returns an async function
 function M.wrap(func, argc, protected, strict)
   vim.validate{
     argc = { argc, 'number' },

From 46af0495f663ebde2e33a4fa67f2ed2e87056e57 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:11:35 +0000
Subject: [PATCH 12/23] feat: allow to work with builtin pcall

---
 lua/async.lua | 60 +++++++++++++++++++++++++--------------------------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index 4230bd5..5d3eb9f 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -79,7 +79,7 @@ function M.run(func, callback, ...)
 
   local function step(...)
     local ret = {coroutine.resume(co, ...)}
-    local stat, nargs, protected, err_or_fn = unpack(ret)
+    local stat, nargs, err_or_fn = unpack(ret)
 
     if not stat then
       error(string.format("The coroutine failed with this message: %s\n%s",
@@ -96,56 +96,59 @@ function M.run(func, callback, ...)
     assert(type(err_or_fn) == 'function', "type error :: expected func")
 
     local args = {select(5, unpack(ret))}
-
-    if protected then
-      args[nargs] = function(...)
-        step(true, ...)
-      end
-      local ok, err_or_handle = pcall(err_or_fn, unpack(args, 1, nargs))
-      if not ok then
-        step(false, err_or_handle)
-      else
-        set_executing_handle(err_or_handle)
-      end
-    else
-      args[nargs] = step
-      set_executing_handle(err_or_fn(unpack(args, 1, nargs)))
-    end
+    args[nargs] = step
+    set_executing_handle(err_or_fn(unpack(args, 1, nargs)))
   end
 
   step(...)
   return handle
 end
 
-local function wait(argc, func, protected, ...)
+local function wait(argc, func, ...)
   vim.validate{
     argc = { argc, 'number'},
     func = { func, 'function' },
-    protected = { protected, 'boolean' },
   }
-  return coroutine.yield(argc, func, protected, ...)
+
+  -- Always run the wrapped functions in xpcall and re-raise the error in the
+  -- coroutine
+  local function pfunc(...)
+    local args = { ... }
+    local cb = args[argc]
+    args[argc] = function(...)
+      cb(true, ...)
+    end
+    xpcall(func, function(err)
+      cb(false, err, debug.traceback())
+    end, unpack(args, 1, argc))
+  end
+
+  local ret = {coroutine.yield(argc, pfunc, ...)}
+
+  local ok = ret[1]
+  if not ok then
+    local _, err, traceback = unpack(ret)
+    error(string.format("Wrapped function failed: %s\n%s", err, traceback))
+  end
+
+  return unpack(ret, 2)
 end
 
 --- Wait on a callback style function
 ---
 --- @tparam integer? argc The number of arguments of func. Must be included.
---- @tparam boolean? protected call the function in protected mode (like pcall)
 --- @tparam function func callback style function to execute
 --- @tparam any ... Arguments for func
 function M.wait(...)
   local nargs, args = select('#', ...), {...}
 
   local argstart
-  local protected = false
   local argc
   for i = 1, math.min(nargs, 3) do
     if type(args[i]) == 'function' then
       argstart = i + 1
       break
     end
-    if type(args[i]) == 'boolean' then
-      protected = args[i]
-    end
     if type(args[i]) == 'number' then
       argc = args[i]
     end
@@ -156,7 +159,7 @@ function M.wait(...)
   local func = args[argstart - 1]
   argc = argc or nargs - argstart + 1
 
-  return wait(argc, func, protected, unpack(args, argstart, nargs))
+  return wait(argc, func, unpack(args, argstart, nargs))
 end
 
 --- Use this to create a function which executes in an async context but
@@ -203,17 +206,14 @@ end
 
 --- Creates an async function with a callback style function.
 ---
---- TODO(lewis6991): Remove protected
----
 --- @tparam function func A callback style function to be converted. The last argument must be the callback.
 --- @tparam integer argc The number of arguments of func. Must be included.
 --- @tparam boolean protected call the function in protected mode (like pcall)
 --- @tparam boolean strict Error when called in non-async context
 --- @treturn function Returns an async function
-function M.wrap(func, argc, protected, strict)
+function M.wrap(func, argc, strict)
   vim.validate{
     argc = { argc, 'number' },
-    protected = { protected, 'boolean', true }
   }
   return function(...)
     if not M.running() then
@@ -222,7 +222,7 @@ function M.wrap(func, argc, protected, strict)
       end
       return func(...)
     end
-    return M.wait(argc, protected, func, ...)
+    return M.wait(argc, func, ...)
   end
 end
 

From a6e1874d2e4e52559c08c5c9db63a9e26c8150c9 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:13:46 +0000
Subject: [PATCH 13/23] refactor: simplify wait

---
 lua/async.lua | 21 +++------------------
 1 file changed, 3 insertions(+), 18 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index 5d3eb9f..e0e0664 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -140,26 +140,11 @@ end
 --- @tparam function func callback style function to execute
 --- @tparam any ... Arguments for func
 function M.wait(...)
-  local nargs, args = select('#', ...), {...}
-
-  local argstart
-  local argc
-  for i = 1, math.min(nargs, 3) do
-    if type(args[i]) == 'function' then
-      argstart = i + 1
-      break
-    end
-    if type(args[i]) == 'number' then
-      argc = args[i]
-    end
+  if type(select(1, ...)) == 'number' then
+    return wait(...)
   end
 
-  assert(argstart)
-
-  local func = args[argstart - 1]
-  argc = argc or nargs - argstart + 1
-
-  return wait(argc, func, unpack(args, argstart, nargs))
+  return wait(select('#', ...) - 1, ...)
 end
 
 --- Use this to create a function which executes in an async context but

From 462347ef65b3af356461f862e65f846028a806eb Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:15:11 +0000
Subject: [PATCH 14/23] feat: improve handle detection

---
 lua/async.lua | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lua/async.lua b/lua/async.lua
index e0e0664..0b974d5 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -27,7 +27,7 @@ function M.running()
 end
 
 local function is_async_handle(handle)
-  if handle and handle.cancel and handle.is_cancelled then
+  if handle and vim.is_callable(handle.cancel) and vim.is_callable(handle.is_cancelled) then
     return true
   end
 end

From f05fba96f9579516c0ec39bf535b9f9e759c2c50 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:18:50 +0000
Subject: [PATCH 15/23] refactor

---
 lua/async.lua | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index 0b974d5..1fa9219 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -68,12 +68,6 @@ function M.run(func, callback, ...)
     return cur_exec_handle and cur_exec_handle:is_cancelled()
   end
 
-  local function set_executing_handle(h)
-    if is_async_handle(h) then
-      cur_exec_handle = h
-    end
-  end
-
   setmetatable(handle, { __index = handle })
   handles[co] = handle
 
@@ -97,7 +91,11 @@ function M.run(func, callback, ...)
 
     local args = {select(5, unpack(ret))}
     args[nargs] = step
-    set_executing_handle(err_or_fn(unpack(args, 1, nargs)))
+
+    local r = err_or_fn(unpack(args, 1, nargs))
+    if is_async_handle(r) then
+      cur_exec_handle = r
+    end
   end
 
   step(...)

From baac8bd0b09e729a963567f6c728038a7deb1936 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:19:12 +0000
Subject: [PATCH 16/23] docs: update

---
 async.md | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/async.md b/async.md
index 5ef5c15..f54e41a 100644
--- a/async.md
+++ b/async.md
@@ -30,14 +30,13 @@ Run a function in an async context.
   `async_t`: Handle
 
 ---
-### `wait(argc, protected, func, ...)`
+### `wait(argc, func, ...)`
 
 Wait on a callback style function
 
 #### Parameters:
 
 * `argc` (`integer?`):  The number of arguments of func. Must be included.
-* `protected` (`boolean?`):  call the function in protected mode (like pcall)
 * `func` (`function`):  callback style function to execute
 * `...` (`any`):  Arguments for func
 
@@ -74,9 +73,6 @@ Create a function which executes in an async context but
 
 Creates an async function with a callback style function.
 
- TODO(lewis6991): Remove protected
-
-
 #### Parameters:
 
 * `func` (`function`):  A callback style function to be converted. The last argument must be the callback.

From 82513a52cb2adb2041dd21bde4ab393fb3776315 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:22:36 +0000
Subject: [PATCH 17/23] refactor: comments

---
 async.md      | 2 +-
 lua/async.lua | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/async.md b/async.md
index f54e41a..e568637 100644
--- a/async.md
+++ b/async.md
@@ -36,7 +36,7 @@ Wait on a callback style function
 
 #### Parameters:
 
-* `argc` (`integer?`):  The number of arguments of func. Must be included.
+* `argc` (`integer?`):  The number of arguments of func.
 * `func` (`function`):  callback style function to execute
 * `...` (`any`):  Arguments for func
 
diff --git a/lua/async.lua b/lua/async.lua
index 1fa9219..2c41fcf 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -109,7 +109,7 @@ local function wait(argc, func, ...)
   }
 
   -- Always run the wrapped functions in xpcall and re-raise the error in the
-  -- coroutine
+  -- coroutine. This makes pcall work as normal.
   local function pfunc(...)
     local args = { ... }
     local cb = args[argc]
@@ -134,7 +134,7 @@ end
 
 --- Wait on a callback style function
 ---
---- @tparam integer? argc The number of arguments of func. Must be included.
+--- @tparam integer? argc The number of arguments of func.
 --- @tparam function func callback style function to execute
 --- @tparam any ... Arguments for func
 function M.wait(...)
@@ -142,6 +142,7 @@ function M.wait(...)
     return wait(...)
   end
 
+  -- Asume argc is equal to the number of passed arguments.
   return wait(select('#', ...) - 1, ...)
 end
 

From 2c5954b5d49933ac415765d4147b070cdc273b1b Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 13:45:02 +0000
Subject: [PATCH 18/23] fix: use table.maxn

---
 lua/async.lua | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index 2c41fcf..a27c6a9 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -82,7 +82,7 @@ function M.run(func, callback, ...)
 
     if coroutine.status(co) == 'dead' then
       if callback then
-        callback(unpack(ret, 4))
+        callback(unpack(ret, 4, table.maxn(ret)))
       end
       return
     end
@@ -129,7 +129,7 @@ local function wait(argc, func, ...)
     error(string.format("Wrapped function failed: %s\n%s", err, traceback))
   end
 
-  return unpack(ret, 2)
+  return unpack(ret, 2, table.maxn(ret))
 end
 
 --- Wait on a callback style function

From f903979044d7a393ab90e8c9546b89b4957e962f Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Mon, 9 Jan 2023 17:05:30 +0000
Subject: [PATCH 19/23] refactor

---
 lua/async.lua | 63 +++++++++++++++++++++++++--------------------------
 1 file changed, 31 insertions(+), 32 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index a27c6a9..b36ef63 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -26,12 +26,36 @@ function M.running()
   end
 end
 
-local function is_async_handle(handle)
-  if handle and vim.is_callable(handle.cancel) and vim.is_callable(handle.is_cancelled) then
+local function is_Async_T(handle)
+  if handle
+    and type(handle) == 'table'
+    and vim.is_callable(handle.cancel)
+    and vim.is_callable(handle.is_cancelled) then
     return true
   end
 end
 
+local Async_T = {}
+
+-- Analogous to uv.close
+function Async_T:cancel(cb)
+  -- Cancel anything running on the event loop
+  if self._current and not self._current:is_cancelled() then
+    self._current:cancel(cb)
+  end
+end
+
+function Async_T.new(co)
+  local handle = setmetatable({}, { __index = Async_T })
+  handles[co] = handle
+  return handle
+end
+
+-- Analogous to uv.is_closing
+function Async_T:is_cancelled()
+  return self._current and self._current:is_cancelled()
+end
+
 --- Run a function in an async context.
 --- @tparam function func
 --- @tparam function callback
@@ -44,32 +68,7 @@ function M.run(func, callback, ...)
   }
 
   local co = coroutine.create(func)
-
-  -- Handle for an object currently running on the event loop.
-  -- The coroutine is paused while this is active.
-  -- Must provide methods cancel() and is_cancelled()
-  local cur_exec_handle
-
-  -- Handle for the user. Since cur_exec_handle will change every
-  -- step() we need to provide access to it through a proxy
-  local handle = {}
-
-  -- Analogous to uv.close
-  function handle:cancel(cb)
-    vim.validate{ callback = { cb , 'function', true } }
-    -- Cancel anything running on the event loop
-    if cur_exec_handle and not cur_exec_handle:is_cancelled() then
-      cur_exec_handle:cancel(cb)
-    end
-  end
-
-  -- Analogous to uv.is_closing
-  function handle:is_cancelled()
-    return cur_exec_handle and cur_exec_handle:is_cancelled()
-  end
-
-  setmetatable(handle, { __index = handle })
-  handles[co] = handle
+  local handle = Async_T.new(co)
 
   local function step(...)
     local ret = {coroutine.resume(co, ...)}
@@ -89,12 +88,12 @@ function M.run(func, callback, ...)
 
     assert(type(err_or_fn) == 'function', "type error :: expected func")
 
-    local args = {select(5, unpack(ret))}
+    local args = {select(4, unpack(ret))}
     args[nargs] = step
 
     local r = err_or_fn(unpack(args, 1, nargs))
-    if is_async_handle(r) then
-      cur_exec_handle = r
+    if is_Async_T(r) then
+      handle._current = r
     end
   end
 
@@ -215,7 +214,7 @@ end
 --- @tparam function[] thunks
 --- @tparam integer n Max number of thunks to run concurrently
 --- @tparam function interrupt_check Function to abort thunks between calls
-function M.join(thunks, n, interrupt_check )
+function M.join(thunks, n, interrupt_check)
   local function run(finish)
     if #thunks == 0 then
       return finish()

From 3c84e0e073d1a1070561a07084acdad36c963654 Mon Sep 17 00:00:00 2001
From: Yuriy Artemyev <anuvyklack@gmail.com>
Date: Sat, 11 Feb 2023 00:58:38 +0300
Subject: [PATCH 20/23] fix: coroutine returns error in second argument, not in
 third

---
 lua/async.lua | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index b36ef63..6af18a6 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -72,11 +72,12 @@ function M.run(func, callback, ...)
 
   local function step(...)
     local ret = {coroutine.resume(co, ...)}
-    local stat, nargs, err_or_fn = unpack(ret)
+    local ok = ret[1]
 
-    if not stat then
-      error(string.format("The coroutine failed with this message: %s\n%s",
-        err_or_fn, debug.traceback(co)))
+    if not ok then
+      local err = ret[2]
+      error(string.format("The coroutine failed with this message:\n%s\n%s",
+                          err, debug.traceback(co)))
     end
 
     if coroutine.status(co) == 'dead' then
@@ -86,12 +87,14 @@ function M.run(func, callback, ...)
       return
     end
 
-    assert(type(err_or_fn) == 'function', "type error :: expected func")
-
+    local nargs, fn = ret[2], ret[3]
     local args = {select(4, unpack(ret))}
+
+    assert(type(fn) == 'function', "type error :: expected func")
+
     args[nargs] = step
 
-    local r = err_or_fn(unpack(args, 1, nargs))
+    local r = fn(unpack(args, 1, nargs))
     if is_Async_T(r) then
       handle._current = r
     end

From 55cd3d3f51cdd6be9a31e762c92a7c39335baedd Mon Sep 17 00:00:00 2001
From: Yuriy Artemyev <anuvyklack@gmail.com>
Date: Sat, 11 Feb 2023 01:01:38 +0300
Subject: [PATCH 21/23] chore: remove outdated function parameter

---
 lua/async.lua | 1 -
 1 file changed, 1 deletion(-)

diff --git a/lua/async.lua b/lua/async.lua
index 6af18a6..ff82eb6 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -194,7 +194,6 @@ end
 ---
 --- @tparam function func A callback style function to be converted. The last argument must be the callback.
 --- @tparam integer argc The number of arguments of func. Must be included.
---- @tparam boolean protected call the function in protected mode (like pcall)
 --- @tparam boolean strict Error when called in non-async context
 --- @treturn function Returns an async function
 function M.wrap(func, argc, strict)

From bad4edbb2917324cd11662dc0209ce53f6c8bc23 Mon Sep 17 00:00:00 2001
From: Yuriy Artemyev <anuvyklack@gmail.com>
Date: Sat, 11 Feb 2023 01:04:09 +0300
Subject: [PATCH 22/23] style

---
 lua/async.lua | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index ff82eb6..c213060 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -62,9 +62,9 @@ end
 --- @tparam any ... Arguments for func
 --- @treturn async_t Handle
 function M.run(func, callback, ...)
-  vim.validate{
+  vim.validate {
     func = { func, 'function' },
-    callback = { callback , 'function', true }
+    callback = { callback, 'function', true }
   }
 
   local co = coroutine.create(func)
@@ -105,8 +105,8 @@ function M.run(func, callback, ...)
 end
 
 local function wait(argc, func, ...)
-  vim.validate{
-    argc = { argc, 'number'},
+  vim.validate {
+    argc = { argc, 'number' },
     func = { func, 'function' },
   }
 
@@ -156,7 +156,7 @@ end
 --- @tparam boolean strict Error when called in non-async context
 --- @treturn function(...):async_t
 function M.create(func, argc, strict)
-  vim.validate{
+  vim.validate {
     func = { func, 'function' },
     argc = { argc, 'number', true }
   }
@@ -168,7 +168,7 @@ function M.create(func, argc, strict)
       end
       return func(...)
     end
-    local callback = select(argc+1, ...)
+    local callback = select(argc + 1, ...)
     return M.run(func, callback, unpack({...}, 1, argc))
   end
 end
@@ -178,7 +178,7 @@ end
 --- @tparam function func
 --- @tparam boolean strict Error when called in non-async context
 function M.void(func, strict)
-  vim.validate{ func = { func , 'function' } }
+  vim.validate { func = { func, 'function' } }
   return function(...)
     if M.running() then
       if strict then
@@ -197,7 +197,7 @@ end
 --- @tparam boolean strict Error when called in non-async context
 --- @treturn function Returns an async function
 function M.wrap(func, argc, strict)
-  vim.validate{
+  vim.validate {
     argc = { argc, 'number' },
   }
   return function(...)
@@ -228,7 +228,7 @@ function M.join(thunks, n, interrupt_check)
     local ret = {}
 
     local function cb(...)
-      ret[#ret+1] = {...}
+      ret[#ret + 1] = {...}
       to_go = to_go - 1
       if to_go == 0 then
         finish(ret)
@@ -260,7 +260,7 @@ function M.curry(fn, ...)
   return function(...)
     local other = {...}
     for i = 1, select('#', ...) do
-      args[nargs+i] = other[i]
+      args[nargs + i] = other[i]
     end
     fn(unpack(args))
   end

From 0fd6f438cc77b34e063ddc4c0c38360a794ab0a3 Mon Sep 17 00:00:00 2001
From: Yuriy Artemyev <anuvyklack@gmail.com>
Date: Sat, 11 Feb 2023 15:36:54 +0300
Subject: [PATCH 23/23] refactor: change assumption in the `wait` function

---
 lua/async.lua | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/lua/async.lua b/lua/async.lua
index c213060..c786e0d 100644
--- a/lua/async.lua
+++ b/lua/async.lua
@@ -140,12 +140,13 @@ end
 --- @tparam function func callback style function to execute
 --- @tparam any ... Arguments for func
 function M.wait(...)
-  if type(select(1, ...)) == 'number' then
+  if type(...) == 'number' then
     return wait(...)
+  else
+    -- Assume argc is equal to the number of passed arguments (- 1 for function
+    -- that is first argument, + 1 for callback that hasn't been passed).
+    return wait(select('#', ...), ...)
   end
-
-  -- Asume argc is equal to the number of passed arguments.
-  return wait(select('#', ...) - 1, ...)
 end
 
 --- Use this to create a function which executes in an async context but