feat(lua): vim.wait() returns callback results #35588

Problem:
The callback passed to `vim.wait` cannot return results directly, it
must set upvalues or globals.

    local rv1, rv2, rv3
    local ok = vim.wait(200, function()
      rv1, rv2, rv3 = 'a', 42, { ok = { 'yes' } }
      return true
    end)

Solution:
Let the callback return values after the first "status" result.

    local ok, rv1, rv2, rv3 = vim.wait(200, function()
      return true, 'a', 42, { ok = { 'yes' } }
    end)

runtime/doc/lua.txt

@@ -843,32 +843,29 @@ vim.ui_detach({ns})                                          *vim.ui_detach()*
       • {ns}  (`integer`) Namespace ID
 
 vim.wait({time}, {callback}, {interval}, {fast_only})             *vim.wait()*
-    Wait for {time} in milliseconds until {callback} returns `true`.
+    Waits up to `time` milliseconds, until `callback` returns `true`
+    (success). Executes `callback` immediately, then at intervals of
+    approximately `interval` milliseconds (default 200). Returns all
+    `callback` results on success.
 
-    Executes {callback} immediately and at approximately {interval}
+    Nvim processes other events while waiting. Cannot be called during an
-    milliseconds (default 200). Nvim still processes other events during this
+    |api-fast| event.
-    time.
-
-    Cannot be called while in an |api-fast| event.
 
     Examples: >lua
-        ---
+        -- Wait for 100 ms, allowing other events to process.
-        -- Wait for 100 ms, allowing other events to process
+        vim.wait(100)
-        vim.wait(100, function() end)
 
-        ---
+        -- Wait up to 1000 ms or until `vim.g.foo` is true, at intervals of ~500 ms.
-        -- Wait for 100 ms or until global variable set.
+        vim.wait(1000, function() return vim.g.foo end, 500)
-        vim.wait(100, function() return vim.g.waiting_for_var end)
 
-        ---
+        -- Wait up to 100 ms or until `vim.g.foo` is true, and get the callback results.
-        -- Wait for 1 second or until global variable set, checking every ~500 ms
+        local ok, rv1, rv2, rv3 = vim.wait(100, function()
-        vim.wait(1000, function() return vim.g.waiting_for_var end, 500)
+          return vim.g.foo, 'a', 42, { ok = { 'yes' } }
+        end)
 
-        ---
+        -- Schedule a function to set a value in 100ms. This would wait 10s if blocked, but actually
-        -- Schedule a function to set a value in 100ms
+        -- only waits 100ms because `vim.wait` processes other events while waiting.
         vim.defer_fn(function() vim.g.timer_result = true end, 100)
-
-        -- Would wait ten seconds if results blocked. Actually only waits  100 ms
         if vim.wait(10000, function() return vim.g.timer_result end) then
           print('Only waiting a little bit of time!')
         end
@@ -886,10 +883,10 @@ vim.wait({time}, {callback}, {interval}, {fast_only})             *vim.wait()*
     Return (multiple): ~
         (`boolean`)
         (`-1|-2?`)
-        • If {callback} returns `true` during the {time}: `true, nil`
+        • If callback returns `true` before timeout: `true, nil, ...`
-        • If {callback} never returns `true` during the {time}: `false, -1`
+        • On timeout: `false, -1`
-        • If {callback} is interrupted during the {time}: `false, -2`
+        • On interrupt: `false, -2`
-        • If {callback} errors, the error is raised.
+        • On error: the error is raised.
 
 
 ==============================================================================

runtime/doc/news.txt

@@ -247,6 +247,7 @@ LSP
 
 LUA
 
+• |vim.wait()| returns the callback results.
 • Lua type annotations for `vim.uv`.
 • |vim.hl.range()| now allows multiple timed highlights.
 • |vim.tbl_extend()| and |vim.tbl_deep_extend()| now accept a function behavior argument.

runtime/lua/vim/_meta/builtin.lua

@@ -179,34 +179,30 @@ function vim.iconv(str, from, to, opts) end
 --- @param fn fun()
 function vim.schedule(fn) end
 
---- Wait for {time} in milliseconds until {callback} returns `true`.
+--- Waits up to `time` milliseconds, until `callback` returns `true` (success). Executes
+--- `callback` immediately, then at intervals of approximately `interval` milliseconds (default
+--- 200). Returns all `callback` results on success.
 ---
---- Executes {callback} immediately and at approximately {interval}
+--- Nvim processes other events while waiting.
---- milliseconds (default 200). Nvim still processes other events during
+--- Cannot be called during an |api-fast| event.
---- this time.
----
---- Cannot be called while in an |api-fast| event.
 ---
 --- Examples:
 ---
 --- ```lua
---- ---
+--- -- Wait for 100 ms, allowing other events to process.
---- -- Wait for 100 ms, allowing other events to process
+--- vim.wait(100)
---- vim.wait(100, function() end)
 ---
---- ---
+--- -- Wait up to 1000 ms or until `vim.g.foo` is true, at intervals of ~500 ms.
---- -- Wait for 100 ms or until global variable set.
+--- vim.wait(1000, function() return vim.g.foo end, 500)
---- vim.wait(100, function() return vim.g.waiting_for_var end)
 ---
---- ---
+--- -- Wait up to 100 ms or until `vim.g.foo` is true, and get the callback results.
---- -- Wait for 1 second or until global variable set, checking every ~500 ms
+--- local ok, rv1, rv2, rv3 = vim.wait(100, function()
---- vim.wait(1000, function() return vim.g.waiting_for_var end, 500)
+---   return vim.g.foo, 'a', 42, { ok = { 'yes' } }
+--- end)
 ---
---- ---
+--- -- Schedule a function to set a value in 100ms. This would wait 10s if blocked, but actually
---- -- Schedule a function to set a value in 100ms
+--- -- only waits 100ms because `vim.wait` processes other events while waiting.
 --- vim.defer_fn(function() vim.g.timer_result = true end, 100)
----
---- -- Would wait ten seconds if results blocked. Actually only waits  100 ms
 --- if vim.wait(10000, function() return vim.g.timer_result end) then
 ---   print('Only waiting a little bit of time!')
 --- end
@@ -216,11 +212,11 @@ function vim.schedule(fn) end
 --- @param callback? fun(): boolean Optional callback. Waits until {callback} returns true
 --- @param interval? integer (Approximate) number of milliseconds to wait between polls
 --- @param fast_only? boolean If true, only |api-fast| events will be processed.
---- @return boolean, nil|-1|-2
+--- @return boolean, nil|-1|-2, ...
----     - If {callback} returns `true` during the {time}: `true, nil`
+---     - If callback returns `true` before timeout: `true, nil, ...`
----     - If {callback} never returns `true` during the {time}: `false, -1`
+---     - On timeout: `false, -1`
----     - If {callback} is interrupted during the {time}: `false, -2`
+---     - On interrupt: `false, -2`
----     - If {callback} errors, the error is raised.
+---     - On error: the error is raised.
 function vim.wait(time, callback, interval, fast_only) end
 
 --- Subscribe to |ui-events|, similar to |nvim_ui_attach()| but receive events in a Lua callback.

src/nvim/lua/executor.c

@@ -175,11 +175,18 @@ int nlua_pcall(lua_State *lstate, int nargs, int nresults)
   lua_getfield(lstate, -1, "traceback");
   lua_remove(lstate, -2);
   lua_insert(lstate, -2 - nargs);
+  int pre_top = lua_gettop(lstate);
   int status = lua_pcall(lstate, nargs, nresults, -2 - nargs);
   if (status) {
     lua_remove(lstate, -2);
   } else {
-    lua_remove(lstate, -1 - nresults);
+    if (nresults == LUA_MULTRET) {
+      int new_top = lua_gettop(lstate);
+      int actual_nres = new_top - pre_top + nargs + 1;
+      lua_remove(lstate, -1 - actual_nres);
+    } else {
+      lua_remove(lstate, -1 - nresults);
+    }
   }
   return status;
 }
@@ -415,16 +422,28 @@ static void dummy_timer_close_cb(TimeWatcher *tw, void *data)
   xfree(tw);
 }
 
-static bool nlua_wait_condition(lua_State *lstate, int *status, bool *callback_result)
+static bool nlua_wait_condition(lua_State *lstate, int *status, bool *callback_result,
+                                int *nresults)
 {
+  int top = lua_gettop(lstate);
   lua_pushvalue(lstate, 2);
-  *status = nlua_pcall(lstate, 0, 1);
+  *status = nlua_pcall(lstate, 0, LUA_MULTRET);
   if (*status) {
     return true;  // break on error, but keep error on stack
   }
-  *callback_result = lua_toboolean(lstate, -1);
+  *nresults = lua_gettop(lstate) - top;
-  lua_pop(lstate, 1);
+  if (*nresults == 0) {
-  return *callback_result;  // break if true
+    *callback_result = false;
+    return false;
+  }
+  *callback_result = lua_toboolean(lstate, top + 1);
+  if (!*callback_result) {
+    lua_settop(lstate, top);
+    return false;
+  }
+  lua_remove(lstate, top + 1);
+  (*nresults)--;
+  return true;  // break if true
 }
 
 /// "vim.wait(timeout, condition[, interval])" function
@@ -454,8 +473,7 @@ static int nlua_wait(lua_State *lstate)
     }
 
     if (!is_function) {
-      lua_pushliteral(lstate,
+      lua_pushliteral(lstate, "vim.wait: callback must be callable");
-                      "vim.wait: if passed, condition must be a function");
       return lua_error(lstate);
     }
   }
@@ -488,6 +506,7 @@ static int nlua_wait(lua_State *lstate)
 
   int pcall_status = 0;
   bool callback_result = false;
+  int nresults = 0;
 
   // Flush screen updates before blocking.
   ui_flush();
@@ -497,7 +516,8 @@ static int nlua_wait(lua_State *lstate)
                             (int)timeout,
                             got_int || (is_function ? nlua_wait_condition(lstate,
                                                                           &pcall_status,
-                                                                          &callback_result)
+                                                                          &callback_result,
+                                                                          &nresults)
                                                     : false));
 
   // Stop dummy timer
@@ -508,18 +528,26 @@ static int nlua_wait(lua_State *lstate)
     return lua_error(lstate);
   } else if (callback_result) {
     lua_pushboolean(lstate, 1);
-    lua_pushnil(lstate);
+    if (nresults == 0) {
+      lua_pushnil(lstate);
+      nresults = 1;
+    } else {
+      lua_insert(lstate, -1 - nresults);
+    }
+    return nresults + 1;
   } else if (got_int) {
     got_int = false;
     vgetc();
     lua_pushboolean(lstate, 0);
     lua_pushinteger(lstate, -2);
+    return 2;
   } else {
     lua_pushboolean(lstate, 0);
     lua_pushinteger(lstate, -1);
+    return 2;
   }
 
-  return 2;
+  abort();
 }
 
 static nlua_ref_state_t *nlua_new_ref_state(lua_State *lstate, bool is_thread)

test/functional/lua/vim_spec.lua

@@ -2096,6 +2096,31 @@ stack traceback:
       exec_lua [[vim.wait(100, function() return true end)]]
     end)
 
+    it('returns all (multiple) callback results', function()
+      eq({ true, false }, exec_lua [[return { vim.wait(200, function() return true, false end) }]])
+      eq(
+        { true, 'a', 42, { ok = { 'yes' } } },
+        exec_lua [[
+          local ok, rv1, rv2, rv3 = vim.wait(200, function()
+            return true, 'a', 42, { ok = { 'yes' } }
+          end)
+
+          return { ok, rv1, rv2, rv3 }
+        ]]
+      )
+    end)
+
+    it('does not return callback results on timeout', function()
+      eq(
+        { false, -1 },
+        exec_lua [[
+          return { vim.wait(1, function()
+            return false, 'a', 42, { ok = { 'yes' } }
+          end) }
+        ]]
+      )
+    end)
+
     it('waits the expected time if false', function()
       eq(
         { time = true, wait_result = { false, -1 } },
@@ -2184,38 +2209,36 @@ stack traceback:
       eq({ false, '[string "<nvim>"]:1: As Expected' }, { result[1], remove_trace(result[2]) })
     end)
 
-    it('if callback is passed, it must be a function', function()
+    it('callback must be a function', function()
       eq(
-        { false, 'vim.wait: if passed, condition must be a function' },
+        { false, 'vim.wait: callback must be callable' },
-        exec_lua [[
+        exec_lua [[return {pcall(function() vim.wait(1000, 13) end)}]]
-        return {pcall(function() vim.wait(1000, 13) end)}
-      ]]
       )
     end)
 
-    it('allows waiting with no callback, explicit', function()
+    it('waits if callback arg is nil', function()
       eq(
         true,
         exec_lua [[
         local start_time = vim.uv.hrtime()
-        vim.wait(50, nil)
+        vim.wait(50, nil) -- select('#', ...) == 1
         return vim.uv.hrtime() - start_time > 25000
       ]]
       )
     end)
 
-    it('allows waiting with no callback, implicit', function()
+    it('waits if callback arg is omitted', function()
       eq(
         true,
         exec_lua [[
         local start_time = vim.uv.hrtime()
-        vim.wait(50)
+        vim.wait(50) -- select('#', ...) == 0
         return vim.uv.hrtime() - start_time > 25000
       ]]
       )
     end)
 
-    it('calls callbacks exactly once if they return true immediately', function()
+    it('invokes callback exactly once if it returns true immediately', function()
       eq(
         true,
         exec_lua [[