lua: vim.wait allows control of fast events (#13053)

* lua: vim.wait allows control of fast events

* fixup: remove requirement of function for easier waiting

* fixup: lint

* fixup: bfredl comments

runtime/doc/lua.txt

@@ -737,13 +737,20 @@ vim.defer_fn({fn}, {timeout})                                    *vim.defer_fn*
     Returns: ~
         |vim.loop|.new_timer() object
 
-vim.wait({time}, {callback} [, {interval}])                         *vim.wait()*
+vim.wait({time} [, {callback}, {interval}, {fast_only}])          *vim.wait()*
         Wait for {time} in milliseconds until {callback} returns `true`.
 
         Executes {callback} immediately and at approximately {interval}
         milliseconds (default 200). Nvim still processes other events during
         this time.
 
+    Parameters: ~
+        {time}      Number of milliseconds to wait
+        {callback}  Optional callback. Waits until {callback} returns true
+        {interval}  (Approximate) number of milliseconds to wait between polls
+        {fast_only} If true, only |api-fast| events will be processed.
+                        If called from while in an |api-fast| event, will
+                        automatically be set to `true`.
 
     Returns: ~
         If {callback} returns `true` during the {time}:

src/nvim/lua/executor.c

@@ -299,8 +299,12 @@ static int nlua_wait(lua_State *lstate)
     return luaL_error(lstate, "timeout must be > 0");
   }
 
+  int lua_top = lua_gettop(lstate);
+
   // Check if condition can be called.
-  bool is_function = (lua_type(lstate, 2) == LUA_TFUNCTION);
+  bool is_function = false;
+  if (lua_top >= 2 && !lua_isnil(lstate, 2)) {
+    is_function = (lua_type(lstate, 2) == LUA_TFUNCTION);
 
     // Check if condition is callable table
     if (!is_function && luaL_getmetafield(lstate, 2, "__call") != 0) {
@@ -309,35 +313,52 @@ static int nlua_wait(lua_State *lstate)
     }
 
     if (!is_function) {
-    lua_pushliteral(lstate, "vim.wait: condition must be a function");
+      lua_pushliteral(
+          lstate,
+          "vim.wait: if passed, condition must be a function");
       return lua_error(lstate);
     }
+  }
 
   intptr_t interval = 200;
-  if (lua_gettop(lstate) >= 3) {
+  if (lua_top >= 3 && !lua_isnil(lstate, 3)) {
     interval = luaL_checkinteger(lstate, 3);
     if (interval < 0) {
       return luaL_error(lstate, "interval must be > 0");
     }
   }
 
+  bool fast_only = false;
+  if (lua_top >= 4) {
+    fast_only =  lua_toboolean(lstate, 4);
+  }
+
+  MultiQueue *loop_events = fast_only || in_fast_callback > 0
+    ? main_loop.fast_events : main_loop.events;
+
   TimeWatcher *tw = xmalloc(sizeof(TimeWatcher));
 
   // Start dummy timer.
   time_watcher_init(&main_loop, tw, NULL);
-  tw->events = main_loop.events;
+  tw->events = loop_events;
   tw->blockable = true;
-  time_watcher_start(tw, dummy_timer_due_cb,
+  time_watcher_start(
-                     (uint64_t)interval, (uint64_t)interval);
+      tw,
+      dummy_timer_due_cb,
+      (uint64_t)interval,
+      (uint64_t)interval);
 
   int pcall_status = 0;
   bool callback_result = false;
 
   LOOP_PROCESS_EVENTS_UNTIL(
       &main_loop,
-      main_loop.events,
+      loop_events,
       (int)timeout,
-      nlua_wait_condition(lstate, &pcall_status, &callback_result) || got_int);
+      is_function ? nlua_wait_condition(
+          lstate,
+          &pcall_status,
+          &callback_result) : false || got_int);
 
   // Stop dummy timer
   time_watcher_stop(tw);

test/functional/lua/vim_spec.lua

@@ -1214,6 +1214,23 @@ describe('lua stdlib', function()
       ]])
     end)
 
+    it('should not process non-fast events when commanded', function()
+      eq({wait_result = false}, exec_lua[[
+        start_time = get_time()
+
+        vim.g.timer_result = false
+        timer = vim.loop.new_timer()
+        timer:start(100, 0, vim.schedule_wrap(function()
+          vim.g.timer_result = true
+        end))
+
+        wait_result = vim.wait(300, function() return vim.g.timer_result end, nil, true)
+
+        return {
+          wait_result = wait_result,
+        }
+      ]])
+    end)
     it('should work with vim.defer_fn', function()
       eq({time = true, wait_result = true}, exec_lua[[
         start_time = get_time()
@@ -1228,15 +1245,6 @@ describe('lua stdlib', function()
       ]])
     end)
 
-    it('should require functions to be passed', function()
-      local pcall_result = exec_lua [[
-        return {pcall(function() vim.wait(1000, 13) end)}
-      ]]
-
-      eq(pcall_result[1], false)
-      matches('condition must be a function', pcall_result[2])
-    end)
-
     it('should not crash when callback errors', function()
       local pcall_result = exec_lua [[
         return {pcall(function() vim.wait(1000, function() error("As Expected") end) end)}
@@ -1246,6 +1254,31 @@ describe('lua stdlib', function()
       matches('As Expected', pcall_result[2])
     end)
 
+    it('if callback is passed, it must be a function', function()
+      local pcall_result = exec_lua [[
+        return {pcall(function() vim.wait(1000, 13) end)}
+      ]]
+
+      eq(pcall_result[1], false)
+      matches('if passed, condition must be a function', pcall_result[2])
+    end)
+
+    it('should allow waiting with no callback, explicit', function()
+      eq(true, exec_lua [[
+        local start_time = vim.loop.hrtime()
+        vim.wait(50, nil)
+        return vim.loop.hrtime() - start_time > 25000
+      ]])
+    end)
+
+    it('should allow waiting with no callback, implicit', function()
+      eq(true, exec_lua [[
+        local start_time = vim.loop.hrtime()
+        vim.wait(50)
+        return vim.loop.hrtime() - start_time > 25000
+      ]])
+    end)
+
     it('should call callbacks exactly once if they return true immediately', function()
       eq(true, exec_lua [[
         vim.g.wait_count = 0