Merge pull request #35536 from bfredl/skipahead

perf(highlight): allow decoration providers to skip ranges without data

fixes #35644

runtime/doc/api.txt

@@ -3362,6 +3362,12 @@ nvim_set_decoration_provider({ns_id}, {opts})
                    included. >
                     ["range", winid, bufnr, begin_row, begin_col, end_row, end_col]
 <
+                   In addition to returning a boolean, it is also allowed to
+                   return a `skip_row, skip_col` pair of integers. This
+                   implies that this function does not need to be called until
+                   a range which continues beyond the skipped position. A
+                   single integer return value `skip_row` is short for
+                   `skip_row, 0`
                  • on_end: called at the end of a redraw cycle >
                     ["end", tick]
 <

runtime/lua/vim/_meta/api.lua

@@ -2175,6 +2175,13 @@ function vim.api.nvim_set_current_win(window) end
 ---   ```
 ---     ["range", winid, bufnr, begin_row, begin_col, end_row, end_col]
 ---   ```
+---
+---   In addition to returning a boolean, it is also allowed to
+---   return a `skip_row, skip_col` pair of integers. This implies
+---   that this function does not need to be called until a range
+---   which continues beyond the skipped position. A single integer
+---   return value `skip_row` is short for `skip_row, 0`
+---
 --- - on_end: called at the end of a redraw cycle
 ---   ```
 ---     ["end", tick]

runtime/lua/vim/_meta/api_keysets.lua

@@ -423,6 +423,7 @@ error('Cannot require a meta file')
 --- @field undo_restore? boolean
 --- @field url? string
 --- @field scoped? boolean
+--- @field _subpriority? integer
 
 --- @class vim.api.keyset.user_command
 --- @field addr? any

runtime/lua/vim/treesitter/highlighter.lua

@@ -53,15 +53,12 @@ function TSHighlighterQuery:query()
   return self._query
 end
 
----@alias MarkInfo { start_line: integer, start_col: integer, opts: vim.api.keyset.set_extmark }
-
 ---@class (private) vim.treesitter.highlighter.State
 ---@field tstree TSTree
 ---@field next_row integer
 ---@field next_col integer
 ---@field iter vim.treesitter.highlighter.Iter?
 ---@field highlighter_query vim.treesitter.highlighter.Query
----@field prev_marks MarkInfo[]
 
 ---@nodoc
 ---@class vim.treesitter.highlighter
@@ -238,7 +235,6 @@ function TSHighlighter:prepare_highlight_states(win, srow, erow)
       next_col = 0,
       iter = nil,
       highlighter_query = hl_query,
-      prev_marks = {},
     })
   end)
 end
@@ -330,44 +326,6 @@ local function get_spell(capture_name)
   return nil, 0
 end
 
----Adds the mark to the buffer, clipped by the line.
----Queues the remainder if the mark continues after the line.
----@param m MarkInfo
----@param buf integer
----@param range_start_row integer
----@param range_start_col integer
----@param range_end_row integer
----@param range_end_col integer
----@param next_marks MarkInfo[]
-local function add_mark(
-  m,
-  buf,
-  range_start_row,
-  range_start_col,
-  range_end_row,
-  range_end_col,
-  next_marks
-)
-  local cur_start_l = m.start_line
-  local cur_start_c = m.start_col
-  if cmp_lt(cur_start_l, cur_start_c, range_start_row, range_start_col) then
-    cur_start_l = range_start_row
-    cur_start_c = range_start_col
-  end
-
-  local cur_opts = m.opts
-  if cmp_lt(range_end_row, range_end_col, cur_opts.end_line, cur_opts.end_col) then
-    cur_opts = vim.deepcopy(cur_opts, true)
-    cur_opts.end_line = range_end_row
-    cur_opts.end_col = range_end_col
-    table.insert(next_marks, m)
-  end
-
-  if cmp_lt(cur_start_l, cur_start_c, cur_opts.end_line, cur_opts.end_col) then
-    api.nvim_buf_set_extmark(buf, ns, cur_start_l, cur_start_c, cur_opts)
-  end
-end
-
 ---@param self vim.treesitter.highlighter
 ---@param win integer
 ---@param buf integer
@@ -398,7 +356,14 @@ local function on_range_impl(
   for i = range_start_row, range_end_row - 1 do
     self._conceal_checked[i] = self._conceal_line or nil
   end
+
+  local MAX_ROW = 2147483647 -- sentinel for skipping to the end of file
+  local skip_until_row = MAX_ROW
+  local skip_until_col = 0
+
+  local subtree_counter = 0
   self:for_each_highlight_state(win, function(state)
+    subtree_counter = subtree_counter + 1
     local root_node = state.tstree:root()
     ---@type { [1]: integer, [2]: integer, [3]: integer, [4]: integer }
     local root_range = { root_node:range() }
@@ -409,25 +374,15 @@ local function on_range_impl(
         { range_start_row, range_start_col, range_end_row, range_end_col }
       )
     then
+      if cmp_lt(root_range[1], root_range[2], skip_until_row, skip_until_col) then
+        skip_until_row = root_range[1]
+        skip_until_col = root_range[2]
+      end
       return
     end
 
     local tree_region = state.tstree:included_ranges(true)
 
-    local next_marks = {}
-
-    for _, mark in ipairs(state.prev_marks) do
-      add_mark(
-        mark,
-        buf,
-        range_start_row,
-        range_start_col,
-        range_end_row,
-        range_end_col,
-        next_marks
-      )
-    end
-
     local next_row = state.next_row
     local next_col = state.next_col
 
@@ -488,7 +443,7 @@ local function on_range_impl(
           local url = get_url(match, buf, capture, metadata)
 
           if hl and not on_conceal and (not on_spell or spell ~= nil) then
-            local opts = {
+            api.nvim_buf_set_extmark(buf, ns, start_row, start_col, {
               end_line = end_row,
               end_col = end_col,
               hl_group = hl,
@@ -497,17 +452,8 @@ local function on_range_impl(
               conceal = conceal,
               spell = spell,
               url = url,
-            }
+              _subpriority = subtree_counter,
-            local mark = { start_line = start_row, start_col = start_col, opts = opts }
+            })
-            add_mark(
-              mark,
-              buf,
-              range_start_row,
-              range_start_col,
-              range_end_row,
-              range_end_col,
-              next_marks
-            )
           end
 
           if
@@ -525,8 +471,12 @@ local function on_range_impl(
 
     state.next_row = next_row
     state.next_col = next_col
-    state.prev_marks = next_marks
+    if cmp_lt(next_row, next_col, skip_until_row, skip_until_col) then
+      skip_until_row = next_row
+      skip_until_col = next_col
+    end
   end)
+  return skip_until_row, skip_until_col
 end
 
 ---@private
@@ -542,7 +492,7 @@ function TSHighlighter._on_range(_, win, buf, br, bc, er, ec, _)
     return
   end
 
-  on_range_impl(self, win, buf, br, bc, er, ec, false, false)
+  return on_range_impl(self, win, buf, br, bc, er, ec, false, false)
 end
 
 ---@private

src/nvim/api/extmark.c

@@ -836,6 +836,15 @@ Integer nvim_buf_set_extmark(Buffer buffer, Integer ns_id, Integer line, Integer
       col2 = c;
     }
 
+    DecorPriority subpriority = 0;
+    if (HAS_KEY(opts, set_extmark, _subpriority)) {
+      VALIDATE_RANGE((opts->_subpriority >= 0 && opts->_subpriority <= UINT16_MAX),
+                     "_subpriority", {
+        goto error;
+      });
+      subpriority = (DecorPriority)opts->_subpriority;
+    }
+
     if (kv_size(virt_text.data.virt_text)) {
       decor_range_add_virt(&decor_state, r, c, line2, col2, decor_put_vt(virt_text, NULL), true);
     }
@@ -845,7 +854,8 @@ Integer nvim_buf_set_extmark(Buffer buffer, Integer ns_id, Integer line, Integer
     if (has_hl) {
       DecorSignHighlight sh = decor_sh_from_inline(hl);
       sh.url = url;
-      decor_range_add_sh(&decor_state, r, c, line2, col2, &sh, true, (uint32_t)ns_id, id);
+      decor_range_add_sh(&decor_state, r, c, line2, col2, &sh, true, (uint32_t)ns_id, id,
+                         subpriority);
     }
   } else {
     if (opts->ephemeral) {
@@ -1051,6 +1061,13 @@ void nvim_buf_clear_namespace(Buffer buffer, Integer ns_id, Integer line_start,
 ///               ```
 ///                 ["range", winid, bufnr, begin_row, begin_col, end_row, end_col]
 ///               ```
+///
+///               In addition to returning a boolean, it is also allowed to
+///               return a `skip_row, skip_col` pair of integers. This implies
+///               that this function does not need to be called until a range
+///               which continues beyond the skipped position. A single integer
+///               return value `skip_row` is short for `skip_row, 0`
+///
 ///             - on_end: called at the end of a redraw cycle
 ///               ```
 ///                 ["end", tick]

src/nvim/api/keysets_defs.h

@@ -62,6 +62,8 @@ typedef struct {
   Boolean undo_restore;
   String url;
   Boolean scoped;
+
+  Integer _subpriority;
 } Dict(set_extmark);
 
 typedef struct {

src/nvim/decoration.c

@@ -553,12 +553,12 @@ static void decor_range_add_from_inline(DecorState *state, int start_row, int st
     uint32_t idx = decor.data.ext.sh_idx;
     while (idx != DECOR_ID_INVALID) {
       DecorSignHighlight *sh = &kv_A(decor_items, idx);
-      decor_range_add_sh(state, start_row, start_col, end_row, end_col, sh, owned, ns, mark_id);
+      decor_range_add_sh(state, start_row, start_col, end_row, end_col, sh, owned, ns, mark_id, 0);
       idx = sh->next;
     }
   } else {
     DecorSignHighlight sh = decor_sh_from_inline(decor.data.hl);
-    decor_range_add_sh(state, start_row, start_col, end_row, end_col, &sh, owned, ns, mark_id);
+    decor_range_add_sh(state, start_row, start_col, end_row, end_col, &sh, owned, ns, mark_id, 0);
   }
 }
 
@@ -619,14 +619,15 @@ void decor_range_add_virt(DecorState *state, int start_row, int start_col, int e
     .data.vt = vt,
     .attr_id = 0,
     .owned = owned,
-    .priority = vt->priority,
+    .priority_internal = ((DecorPriorityInternal)vt->priority << 16),
     .draw_col = -10,
   };
   decor_range_insert(state, &range);
 }
 
 void decor_range_add_sh(DecorState *state, int start_row, int start_col, int end_row, int end_col,
-                        DecorSignHighlight *sh, bool owned, uint32_t ns, uint32_t mark_id)
+                        DecorSignHighlight *sh, bool owned, uint32_t ns, uint32_t mark_id,
+                        DecorPriority subpriority)
 {
   if (sh->flags & kSHIsSign) {
     return;
@@ -638,7 +639,7 @@ void decor_range_add_sh(DecorState *state, int start_row, int start_col, int end
     .data.sh = *sh,
     .attr_id = 0,
     .owned = owned,
-    .priority = sh->priority,
+    .priority_internal = ((DecorPriorityInternal)sh->priority << 16) + subpriority,
     .draw_col = -10,
   };
 
@@ -731,7 +732,7 @@ next_mark:
       break;
     }
     int const ordering = r->ordering;
-    DecorPriority const priority = r->priority;
+    DecorPriorityInternal const priority = r->priority_internal;
 
     int begin = 0;
     int end = cur_end;
@@ -739,7 +740,8 @@ next_mark:
       int mid = begin + ((end - begin) >> 1);
       int mi = indices[mid];
       DecorRange *mr = &slots[mi].range;
-      if (mr->priority < priority || (mr->priority == priority && mr->ordering < ordering)) {
+      if (mr->priority_internal < priority
+          || (mr->priority_internal == priority && mr->ordering < ordering)) {
         begin = mid + 1;
       } else {
         end = mid;

src/nvim/decoration.h

@@ -37,7 +37,7 @@ typedef struct {
   int end_row;
   int end_col;
   int ordering;  ///< range insertion order
-  DecorPriority priority;
+  DecorPriorityInternal priority_internal;
   bool owned;  ///< ephemeral decoration, free memory immediately
   DecorRangeKind kind;
   // next pointers MUST NOT be used, these are separate ranges

src/nvim/decoration_defs.h

@@ -36,6 +36,7 @@ enum {
 typedef kvec_t(struct virt_line { VirtText line; int flags; }) VirtLines;
 
 typedef uint16_t DecorPriority;
+typedef uint32_t DecorPriorityInternal;
 #define DECOR_PRIORITY_BASE 0x1000
 
 /// Keep in sync with hl_mode_str[] in decoration.h
@@ -145,6 +146,9 @@ typedef struct {
     kDecorProviderDisabled = 4,
   } state;
 
+  int win_skip_row;
+  int win_skip_col;
+
   LuaRef redraw_start;
   LuaRef redraw_buf;
   LuaRef redraw_win;
@@ -159,3 +163,8 @@ typedef struct {
 
   uint8_t error_count;
 } DecorProvider;
+
+#define DECORATION_PROVIDER_INIT(ns_id) (DecorProvider) \
+  { ns_id, kDecorProviderDisabled, 0, 0, LUA_NOREF, LUA_NOREF, \
+    LUA_NOREF, LUA_NOREF, LUA_NOREF, LUA_NOREF, \
+    LUA_NOREF, LUA_NOREF, -1, false, false, 0 }

src/nvim/decoration_provider.c

@@ -23,11 +23,6 @@
 
 static kvec_t(DecorProvider) decor_providers = KV_INITIAL_VALUE;
 
-#define DECORATION_PROVIDER_INIT(ns_id) (DecorProvider) \
-  { ns_id, kDecorProviderDisabled, LUA_NOREF, LUA_NOREF, \
-    LUA_NOREF, LUA_NOREF, LUA_NOREF, LUA_NOREF, \
-    LUA_NOREF, LUA_NOREF, -1, false, false, 0 }
-
 static void decor_provider_error(DecorProvider *provider, const char *name, const char *msg)
 {
   const char *ns = describe_ns(provider->ns_id, "(UNKNOWN PLUGIN)");
@@ -38,21 +33,28 @@ static void decor_provider_error(DecorProvider *provider, const char *name, cons
 // Note we pass in a provider index as this function may cause decor_providers providers to be
 // reallocated so we need to be careful with DecorProvider pointers
 static bool decor_provider_invoke(int provider_idx, const char *name, LuaRef ref, Array args,
-                                  bool default_true)
+                                  bool default_true, Array *res)
 {
   Error err = ERROR_INIT;
 
   textlock++;
-  Object ret = nlua_call_ref(ref, name, args, kRetNilBool, NULL, &err);
+  Object ret = nlua_call_ref(ref, name, args, res ? kRetMulti : kRetNilBool, NULL, &err);
   textlock--;
 
   // We get the provider here via an index in case the above call to nlua_call_ref causes
   // decor_providers to be reallocated.
   DecorProvider *provider = &kv_A(decor_providers, provider_idx);
-  if (!ERROR_SET(&err)
+  if (!ERROR_SET(&err)) {
-      && api_object_to_bool(ret, "provider %s retval", default_true, &err)) {
     provider->error_count = 0;
+    if (res) {
+      assert(ret.type == kObjectTypeArray);
+      *res = ret.data.array;
       return true;
+    } else {
+      if (api_object_to_bool(ret, "provider %s retval", default_true, &err)) {
+        return true;
+      }
+    }
   }
 
   if (ERROR_SET(&err) && provider->error_count < CB_MAX_ERROR) {
@@ -65,7 +67,7 @@ static bool decor_provider_invoke(int provider_idx, const char *name, LuaRef ref
   }
 
   api_clear_error(&err);
-  api_free_object(ret);
+  api_free_object(ret);  // TODO(bfredl): wants to be on an arena
   return false;
 }
 
@@ -81,7 +83,7 @@ void decor_providers_invoke_spell(win_T *wp, int start_row, int start_col, int e
       ADD_C(args, INTEGER_OBJ(start_col));
       ADD_C(args, INTEGER_OBJ(end_row));
       ADD_C(args, INTEGER_OBJ(end_col));
-      decor_provider_invoke((int)i, "spell", p->spell_nav, args, true);
+      decor_provider_invoke((int)i, "spell", p->spell_nav, args, true, NULL);
     }
   }
 }
@@ -97,7 +99,7 @@ bool decor_providers_invoke_conceal_line(win_T *wp, int row)
       ADD_C(args, INTEGER_OBJ(wp->handle));
       ADD_C(args, INTEGER_OBJ(wp->w_buffer->handle));
       ADD_C(args, INTEGER_OBJ(row));
-      decor_provider_invoke((int)i, "conceal_line", p->conceal_line, args, true);
+      decor_provider_invoke((int)i, "conceal_line", p->conceal_line, args, true, NULL);
     }
   }
   return wp->w_buffer->b_marktree->n_keys > keys;
@@ -114,7 +116,7 @@ void decor_providers_start(void)
     if (p->state != kDecorProviderDisabled && p->redraw_start != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 2);
       ADD_C(args, INTEGER_OBJ((int)display_tick));
-      bool active = decor_provider_invoke((int)i, "start", p->redraw_start, args, true);
+      bool active = decor_provider_invoke((int)i, "start", p->redraw_start, args, true, NULL);
       kv_A(decor_providers, i).state = active ? kDecorProviderActive : kDecorProviderRedrawDisabled;
     } else if (p->state != kDecorProviderDisabled) {
       kv_A(decor_providers, i).state = kDecorProviderActive;
@@ -147,6 +149,9 @@ void decor_providers_invoke_win(win_T *wp)
       p->state = kDecorProviderActive;
     }
 
+    p->win_skip_row = 0;
+    p->win_skip_col = 0;
+
     if (p->state == kDecorProviderActive && p->redraw_win != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 4);
       ADD_C(args, WINDOW_OBJ(wp->handle));
@@ -154,7 +159,8 @@ void decor_providers_invoke_win(win_T *wp)
       // TODO(bfredl): we are not using this, but should be first drawn line?
       ADD_C(args, INTEGER_OBJ(wp->w_topline - 1));
       ADD_C(args, INTEGER_OBJ(botline - 1));
-      if (!decor_provider_invoke((int)i, "win", p->redraw_win, args, true)) {
+      // TODO(bfredl): could skip a call if retval was interpreted like range?
+      if (!decor_provider_invoke((int)i, "win", p->redraw_win, args, true, NULL)) {
         kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
       }
     }
@@ -178,7 +184,7 @@ void decor_providers_invoke_line(win_T *wp, int row)
       ADD_C(args, WINDOW_OBJ(wp->handle));
       ADD_C(args, BUFFER_OBJ(wp->w_buffer->handle));
       ADD_C(args, INTEGER_OBJ(row));
-      if (!decor_provider_invoke((int)i, "line", p->redraw_line, args, true)) {
+      if (!decor_provider_invoke((int)i, "line", p->redraw_line, args, true, NULL)) {
         // return 'false' or error: skip rest of this window
         kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
       }
@@ -195,6 +201,10 @@ void decor_providers_invoke_range(win_T *wp, int start_row, int start_col, int e
   for (size_t i = 0; i < kv_size(decor_providers); i++) {
     DecorProvider *p = &kv_A(decor_providers, i);
     if (p->state == kDecorProviderActive && p->redraw_range != LUA_NOREF) {
+      if (p->win_skip_row > end_row || (p->win_skip_row == end_row && p->win_skip_col >= end_col)) {
+        continue;
+      }
+
       MAXSIZE_TEMP_ARRAY(args, 6);
       ADD_C(args, WINDOW_OBJ(wp->handle));
       ADD_C(args, BUFFER_OBJ(wp->w_buffer->handle));
@@ -202,10 +212,34 @@ void decor_providers_invoke_range(win_T *wp, int start_row, int start_col, int e
       ADD_C(args, INTEGER_OBJ(start_col));
       ADD_C(args, INTEGER_OBJ(end_row));
       ADD_C(args, INTEGER_OBJ(end_col));
-      if (!decor_provider_invoke((int)i, "range", p->redraw_range, args, true)) {
+      Array res = ARRAY_DICT_INIT;
-        // return 'false' or error: skip rest of this window
+      bool status = decor_provider_invoke((int)i, "range", p->redraw_range, args, true, &res);
-        kv_A(decor_providers, i).state = kDecorProviderWinDisabled;
+      p = &kv_A(decor_providers, i);  // lua call might have reallocated decor_providers
+
+      if (!status) {
+        // error: skip rest of this window
+        p->state = kDecorProviderWinDisabled;
+      } else if (res.size >= 1) {
+        Object first = res.items[0];
+        if (first.type == kObjectTypeBoolean) {
+          if (first.data.boolean == false) {
+            p->state = kDecorProviderWinDisabled;
           }
+        } else if (first.type == kObjectTypeInteger) {
+          Integer row = first.data.integer;
+          Integer col = 0;
+          if (res.size >= 2) {
+            Object second = res.items[1];
+            if (second.type == kObjectTypeInteger) {
+              col = second.data.integer;
+            }
+          }
+          p->win_skip_row = (int)row;
+          p->win_skip_col = (int)col;
+        }
+      }
+
+      api_free_array(res);
 
       hl_check_ns();
     }
@@ -226,7 +260,7 @@ void decor_providers_invoke_buf(buf_T *buf)
       MAXSIZE_TEMP_ARRAY(args, 2);
       ADD_C(args, BUFFER_OBJ(buf->handle));
       ADD_C(args, INTEGER_OBJ((int64_t)display_tick));
-      decor_provider_invoke((int)i, "buf", p->redraw_buf, args, true);
+      decor_provider_invoke((int)i, "buf", p->redraw_buf, args, true, NULL);
     }
   }
 }
@@ -243,7 +277,7 @@ void decor_providers_invoke_end(void)
     if (p->state != kDecorProviderDisabled && p->redraw_end != LUA_NOREF) {
       MAXSIZE_TEMP_ARRAY(args, 1);
       ADD_C(args, INTEGER_OBJ((int)display_tick));
-      decor_provider_invoke((int)i, "end", p->redraw_end, args, true);
+      decor_provider_invoke((int)i, "end", p->redraw_end, args, true, NULL);
     }
   }
   decor_check_to_be_deleted();

src/nvim/lua/executor.c

@@ -181,12 +181,9 @@ int nlua_pcall(lua_State *lstate, int nargs, int nresults)
     lua_remove(lstate, -2);
   } else {
     if (nresults == LUA_MULTRET) {
-      int new_top = lua_gettop(lstate);
+      nresults = lua_gettop(lstate) - (pre_top - nargs - 1);
-      int actual_nres = new_top - pre_top + nargs + 1;
-      lua_remove(lstate, -1 - actual_nres);
-    } else {
-      lua_remove(lstate, -1 - nresults);
     }
+    lua_remove(lstate, -1 - nresults);
   }
   return status;
 }
@@ -1551,6 +1548,7 @@ Object nlua_exec(const String str, const char *chunkname, const Array args, LuaR
 {
   lua_State *const lstate = global_lstate;
 
+  int top = lua_gettop(lstate);
   const char *name = (chunkname && chunkname[0]) ? chunkname : "<nvim>";
   if (luaL_loadbuffer(lstate, str.data, str.size, name)) {
     size_t len;
@@ -1570,7 +1568,7 @@ Object nlua_exec(const String str, const char *chunkname, const Array args, LuaR
     return NIL;
   }
 
-  return nlua_call_pop_retval(lstate, mode, arena, err);
+  return nlua_call_pop_retval(lstate, mode, arena, top, err);
 }
 
 bool nlua_ref_is_function(LuaRef ref)
@@ -1601,10 +1599,16 @@ Object nlua_call_ref(LuaRef ref, const char *name, Array args, LuaRetMode mode,
   return nlua_call_ref_ctx(false, ref, name, args, mode, arena, err);
 }
 
+static int mode_ret(LuaRetMode mode)
+{
+  return mode == kRetMulti ? LUA_MULTRET : 1;
+}
+
 Object nlua_call_ref_ctx(bool fast, LuaRef ref, const char *name, Array args, LuaRetMode mode,
                          Arena *arena, Error *err)
 {
   lua_State *const lstate = global_lstate;
+  int top = lua_gettop(lstate);
   nlua_pushref(lstate, ref);
   int nargs = (int)args.size;
   if (name != NULL) {
@@ -1616,12 +1620,12 @@ Object nlua_call_ref_ctx(bool fast, LuaRef ref, const char *name, Array args, Lu
   }
 
   if (fast) {
-    if (nlua_fast_cfpcall(lstate, nargs, 1, -1) < 0) {
+    if (nlua_fast_cfpcall(lstate, nargs, mode_ret(mode), -1) < 0) {
       // error is already scheduled, set anyways to convey failure.
       api_set_error(err, kErrorTypeException, "fast context failure");
       return NIL;
     }
-  } else if (nlua_pcall(lstate, nargs, 1)) {
+  } else if (nlua_pcall(lstate, nargs, mode_ret(mode))) {
     // if err is passed, the caller will deal with the error.
     if (err) {
       size_t len;
@@ -1633,16 +1637,18 @@ Object nlua_call_ref_ctx(bool fast, LuaRef ref, const char *name, Array args, Lu
     return NIL;
   }
 
-  return nlua_call_pop_retval(lstate, mode, arena, err);
+  return nlua_call_pop_retval(lstate, mode, arena, top, err);
 }
 
-static Object nlua_call_pop_retval(lua_State *lstate, LuaRetMode mode, Arena *arena, Error *err)
+static Object nlua_call_pop_retval(lua_State *lstate, LuaRetMode mode, Arena *arena, int pretop,
+                                   Error *err)
 {
-  if (lua_isnil(lstate, -1)) {
+  if (mode != kRetMulti && lua_isnil(lstate, -1)) {
     lua_pop(lstate, 1);
     return NIL;
   }
   Error dummy = ERROR_INIT;
+  Error *perr = err ? err : &dummy;
 
   switch (mode) {
   case kRetNilBool: {
@@ -1658,7 +1664,19 @@ static Object nlua_call_pop_retval(lua_State *lstate, LuaRetMode mode, Arena *ar
     return LUAREF_OBJ(ref);
   }
   case kRetObject:
-    return nlua_pop_Object(lstate, false, arena, err ? err : &dummy);
+    return nlua_pop_Object(lstate, false, arena, perr);
+  case kRetMulti:
+    ;
+    int nres = lua_gettop(lstate) - pretop;
+    Array res = arena_array(arena, (size_t)nres);
+    for (int i = 0; i < nres; i++) {
+      res.items[nres - i - 1] = nlua_pop_Object(lstate, false, arena, perr);
+      if (ERROR_SET(err)) {
+        return NIL;
+      }
+    }
+    res.size = (size_t)nres;
+    return ARRAY_OBJ(res);
   }
   UNREACHABLE;
 }

src/nvim/lua/executor.h

@@ -42,6 +42,7 @@ typedef enum {
   kRetNilBool,  ///< NIL preserved as such, other values return their booleanness
                 ///< Should also be used when return value is ignored, as it is allocation-free
   kRetLuaref,   ///< return value becomes a single Luaref, regardless of type (except NIL)
+  kRetMulti,    ///< like kRetObject but return muliple return values as an Array
 } LuaRetMode;
 
 /// Maximum number of errors in vim.ui_attach() and decor provider callbacks.