feat(lua-api): avoid unnecessary allocations (#19877)

Lua makes (or reuses) an internal copy of strings, so we can safely push
buf pointers onto the stack.

11 files changed, 184 insertions, 96 deletions

diff --git a/scripts/gen_vimdoc.py b/scripts/gen_vimdoc.py
index a044c8c39d..e77d3ea286 100755
--- a/scripts/gen_vimdoc.py
+++ b/scripts/gen_vimdoc.py
@@ -796,7 +796,8 @@ def extract_from_xml(filename, target, width, fmt_vimhelp):
 
         prefix = '%s(' % name
         suffix = '%s)' % ', '.join('{%s}' % a[1] for a in params
-                                   if a[0] not in ('void', 'Error', 'Arena'))
+                                   if a[0] not in ('void', 'Error', 'Arena',
+                                                   'lua_State'))
 
         if not fmt_vimhelp:
             c_decl = '%s %s(%s);' % (return_type, name, ', '.join(c_args))
diff --git a/src/nvim/api/buffer.c b/src/nvim/api/buffer.c
index 51fedb302a..29c2ed6028 100644
--- a/src/nvim/api/buffer.c
+++ b/src/nvim/api/buffer.c
@@ -271,6 +271,7 @@ ArrayOf(String) nvim_buf_get_lines(uint64_t channel_id,
                                    Integer start,
                                    Integer end,
                                    Boolean strict_indexing,
+                                   lua_State *lstate,
                                    Error *err)
   FUNC_API_SINCE(1)
 {
@@ -300,21 +301,18 @@ ArrayOf(String) nvim_buf_get_lines(uint64_t channel_id,
     return rv;
   }
 
-  rv.size = (size_t)(end - start);
-  rv.items = xcalloc(rv.size, sizeof(Object));
+  size_t size = (size_t)(end - start);
 
-  if (!buf_collect_lines(buf, rv.size, start,
-                         (channel_id != VIML_INTERNAL_CALL), &rv, err)) {
+  init_line_array(lstate, &rv, size);
+
+  if (!buf_collect_lines(buf, size, (linenr_T)start, (channel_id != VIML_INTERNAL_CALL), &rv,
+                         lstate, err)) {
     goto end;
   }
 
 end:
   if (ERROR_SET(err)) {
-    for (size_t i = 0; i < rv.size; i++) {
-      xfree(rv.items[i].data.string.data);
-    }
-
-    xfree(rv.items);
+    api_free_array(rv);
     rv.items = NULL;
   }
 
@@ -790,7 +788,8 @@ early_end:
 ArrayOf(String) nvim_buf_get_text(uint64_t channel_id, Buffer buffer,
                                   Integer start_row, Integer start_col,
                                   Integer end_row, Integer end_col,
-                                  Dictionary opts, Error *err)
+                                  Dictionary opts, lua_State *lstate,
+                                  Error *err)
   FUNC_API_SINCE(9)
 {
   Array rv = ARRAY_DICT_INIT;
@@ -830,33 +829,38 @@ ArrayOf(String) nvim_buf_get_text(uint64_t channel_id, Buffer buffer,
 
   bool replace_nl = (channel_id != VIML_INTERNAL_CALL);
 
+  size_t size = (size_t)(end_row - start_row) + 1;
+
+  init_line_array(lstate, &rv, size);
+
   if (start_row == end_row) {
-    String line = buf_get_text(buf, start_row, start_col, end_col, replace_nl, err);
+    String line = buf_get_text(buf, start_row, start_col, end_col, err);
     if (ERROR_SET(err)) {
-      return rv;
+      goto end;
     }
-
-    ADD(rv, STRING_OBJ(line));
+    push_linestr(lstate, &rv, line.data, line.size, 0, replace_nl);
     return rv;
   }
 
-  rv.size = (size_t)(end_row - start_row) + 1;
-  rv.items = xcalloc(rv.size, sizeof(Object));
+  String str = buf_get_text(buf, start_row, start_col, MAXCOL - 1, err);
+
+  push_linestr(lstate, &rv, str.data, str.size, 0, replace_nl);
 
-  rv.items[0] = STRING_OBJ(buf_get_text(buf, start_row, start_col, MAXCOL - 1, replace_nl, err));
   if (ERROR_SET(err)) {
     goto end;
   }
 
-  if (rv.size > 2) {
+  if (size > 2) {
     Array tmp = ARRAY_DICT_INIT;
     tmp.items = &rv.items[1];
-    if (!buf_collect_lines(buf, rv.size - 2, start_row + 1, replace_nl, &tmp, err)) {
+    if (!buf_collect_lines(buf, size - 2, (linenr_T)start_row + 1, replace_nl, &tmp, lstate, err)) {
       goto end;
     }
   }
 
-  rv.items[rv.size - 1] = STRING_OBJ(buf_get_text(buf, end_row, 0, end_col, replace_nl, err));
+  str = buf_get_text(buf, end_row, 0, end_col, err);
+  push_linestr(lstate, &rv, str.data, str.size, (int)(size - 1), replace_nl);
+
   if (ERROR_SET(err)) {
     goto end;
   }
@@ -1390,3 +1394,90 @@ static int64_t normalize_index(buf_T *buf, int64_t index, bool end_exclusive, bo
   index++;
   return index;
 }
+
+/// Initialise a string array either:
+/// - on the Lua stack (as a table) (if lstate is not NULL)
+/// - as an API array object (if lstate is NULL).
+///
+/// @param lstate  Lua state. When NULL the Array is initialized instead.
+/// @param a       Array to initialize
+/// @param size    Size of array
+static inline void init_line_array(lua_State *lstate, Array *a, size_t size)
+{
+  if (lstate) {
+    lua_createtable(lstate, (int)size, 0);
+  } else {
+    a->size = size;
+    a->items = xcalloc(a->size, sizeof(Object));
+  }
+}
+
+/// Push a string onto either the Lua stack (as a table element) or an API array object.
+///
+/// For Lua, a table of the correct size must be created first.
+/// API array objects must be pre allocated.
+///
+/// @param lstate      Lua state. When NULL the Array is pushed to instead.
+/// @param a           Array to push onto when not using Lua
+/// @param s           String to push
+/// @param len         Size of string
+/// @param idx         0-based index to place s
+/// @param replace_nl  Replace newlines ('\n') with null ('\0')
+static void push_linestr(lua_State *lstate, Array *a, const char *s, size_t len, int idx,
+                         bool replace_nl)
+{
+  if (lstate) {
+    // Vim represents NULs as NLs
+    if (s && replace_nl && strchr(s, '\n')) {
+      char *tmp = xmemdupz(s, len);
+      strchrsub(tmp, '\n', '\0');
+      lua_pushlstring(lstate, tmp, len);
+      xfree(tmp);
+    } else {
+      lua_pushlstring(lstate, s, len);
+    }
+    lua_rawseti(lstate, -2, idx + 1);
+  } else {
+    String str = STRING_INIT;
+    if (s) {
+      str = cbuf_to_string(s, len);
+      if (replace_nl) {
+        // Vim represents NULs as NLs, but this may confuse clients.
+        strchrsub(str.data, '\n', '\0');
+      }
+    }
+
+    a->items[idx] = STRING_OBJ(str);
+  }
+}
+
+/// Collects `n` buffer lines into array `l` and/or lua_State `lstate`, optionally replacing
+/// newlines with NUL.
+///
+/// @param buf Buffer to get lines from
+/// @param n Number of lines to collect
+/// @param replace_nl Replace newlines ("\n") with NUL
+/// @param start Line number to start from
+/// @param[out] l If not NULL, Lines are copied here
+/// @param[out] lstate If not NULL, Lines are pushed into a table onto the stack
+/// @param err[out] Error, if any
+/// @return true unless `err` was set
+bool buf_collect_lines(buf_T *buf, size_t n, linenr_T start, bool replace_nl, Array *l,
+                       lua_State *lstate, Error *err)
+{
+  for (size_t i = 0; i < n; i++) {
+    linenr_T lnum = start + (linenr_T)i;
+
+    if (lnum >= MAXLNUM) {
+      if (err != NULL) {
+        api_set_error(err, kErrorTypeValidation, "Line index is too high");
+      }
+      return false;
+    }
+
+    char *bufstr = ml_get_buf(buf, lnum, false);
+    push_linestr(lstate, l, bufstr, strlen(bufstr), (int)i, replace_nl);
+  }
+
+  return true;
+}
diff --git a/src/nvim/api/buffer.h b/src/nvim/api/buffer.h
index 1c4a93a587..0814da63cd 100644
--- a/src/nvim/api/buffer.h
+++ b/src/nvim/api/buffer.h
@@ -1,7 +1,10 @@
 #ifndef NVIM_API_BUFFER_H
 #define NVIM_API_BUFFER_H
 
+#include <lauxlib.h>
+
 #include "nvim/api/private/defs.h"
+#include "nvim/buffer_defs.h"
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "api/buffer.h.generated.h"
diff --git a/src/nvim/api/deprecated.c b/src/nvim/api/deprecated.c
index abaac07755..8e1a615bbb 100644
--- a/src/nvim/api/deprecated.c
+++ b/src/nvim/api/deprecated.c
@@ -190,7 +190,7 @@ String buffer_get_line(Buffer buffer, Integer index, Error *err)
   String rv = { .size = 0 };
 
   index = convert_index(index);
-  Array slice = nvim_buf_get_lines(0, buffer, index, index + 1, true, err);
+  Array slice = nvim_buf_get_lines(0, buffer, index, index + 1, true, NULL, err);
 
   if (!ERROR_SET(err) && slice.size) {
     rv = slice.items[0].data.string;
@@ -263,7 +263,7 @@ ArrayOf(String) buffer_get_line_slice(Buffer buffer,
 {
   start = convert_index(start) + !include_start;
   end = convert_index(end) + include_end;
-  return nvim_buf_get_lines(0, buffer, start, end, false, err);
+  return nvim_buf_get_lines(0, buffer, start, end, false, NULL, err);
 }
 
 /// Replaces a line range on the buffer
diff --git a/src/nvim/api/options.c b/src/nvim/api/options.c
index 92f20fff48..1b04392d47 100644
--- a/src/nvim/api/options.c
+++ b/src/nvim/api/options.c
@@ -256,7 +256,7 @@ void nvim_set_option(uint64_t channel_id, String name, Object value, Error *err)
 /// @param name     Option name
 /// @param[out] err Error details, if any
 /// @return         Option value (global)
-Object nvim_get_option(String name, Error *err)
+Object nvim_get_option(String name, Arena *arena, Error *err)
   FUNC_API_SINCE(1)
 {
   return get_option_from(NULL, SREQ_GLOBAL, name, err);
@@ -268,7 +268,7 @@ Object nvim_get_option(String name, Error *err)
 /// @param name       Option name
 /// @param[out] err   Error details, if any
 /// @return Option value
-Object nvim_buf_get_option(Buffer buffer, String name, Error *err)
+Object nvim_buf_get_option(Buffer buffer, String name, Arena *arena, Error *err)
   FUNC_API_SINCE(1)
 {
   buf_T *buf = find_buffer_by_handle(buffer, err);
@@ -306,7 +306,7 @@ void nvim_buf_set_option(uint64_t channel_id, Buffer buffer, String name, Object
 /// @param name     Option name
 /// @param[out] err Error details, if any
 /// @return Option value
-Object nvim_win_get_option(Window window, String name, Error *err)
+Object nvim_win_get_option(Window window, String name, Arena *arena, Error *err)
   FUNC_API_SINCE(1)
 {
   win_T *win = find_window_by_handle(window, err);
@@ -346,7 +346,7 @@ void nvim_win_set_option(uint64_t channel_id, Window window, String name, Object
 /// @param name The option name
 /// @param[out] err Details of an error that may have occurred
 /// @return the option value
-Object get_option_from(void *from, int type, String name, Error *err)
+static Object get_option_from(void *from, int type, String name, Error *err)
 {
   Object rv = OBJECT_INIT;
 
@@ -358,8 +358,7 @@ Object get_option_from(void *from, int type, String name, Error *err)
   // Return values
   int64_t numval;
   char *stringval = NULL;
-  int flags = get_option_value_strict(name.data, &numval, &stringval,
-                                      type, from);
+  int flags = get_option_value_strict(name.data, &numval, &stringval, type, from);
 
   if (!flags) {
     api_set_error(err, kErrorTypeValidation, "Invalid option name: '%s'",
diff --git a/src/nvim/api/private/helpers.c b/src/nvim/api/private/helpers.c
index 73b5489d5c..d10d17c88d 100644
--- a/src/nvim/api/private/helpers.c
+++ b/src/nvim/api/private/helpers.c
@@ -394,6 +394,12 @@ String cstrn_to_string(const char *str, size_t maxsize)
   return cbuf_to_string(str, STRNLEN(str, maxsize));
 }
 
+String cstrn_as_string(char *str, size_t maxsize)
+  FUNC_ATTR_NONNULL_ALL
+{
+  return cbuf_as_string(str, STRNLEN(str, maxsize));
+}
+
 /// Creates a String using the given C string. Unlike
 /// cstr_to_string this function DOES NOT copy the C string.
 ///
@@ -462,53 +468,15 @@ Array string_to_array(const String input, bool crlf)
   return ret;
 }
 
-/// Collects `n` buffer lines into array `l`, optionally replacing newlines
-/// with NUL.
-///
-/// @param buf Buffer to get lines from
-/// @param n Number of lines to collect
-/// @param replace_nl Replace newlines ("\n") with NUL
-/// @param start Line number to start from
-/// @param[out] l Lines are copied here
-/// @param err[out] Error, if any
-/// @return true unless `err` was set
-bool buf_collect_lines(buf_T *buf, size_t n, int64_t start, bool replace_nl, Array *l, Error *err)
-{
-  for (size_t i = 0; i < n; i++) {
-    int64_t lnum = start + (int64_t)i;
-
-    if (lnum >= MAXLNUM) {
-      if (err != NULL) {
-        api_set_error(err, kErrorTypeValidation, "Line index is too high");
-      }
-      return false;
-    }
-
-    const char *bufstr = ml_get_buf(buf, (linenr_T)lnum, false);
-    Object str = STRING_OBJ(cstr_to_string(bufstr));
-
-    if (replace_nl) {
-      // Vim represents NULs as NLs, but this may confuse clients.
-      strchrsub(str.data.string.data, '\n', '\0');
-    }
-
-    l->items[i] = str;
-  }
-
-  return true;
-}
-
 /// Returns a substring of a buffer line
 ///
 /// @param buf          Buffer handle
 /// @param lnum         Line number (1-based)
 /// @param start_col    Starting byte offset into line (0-based)
 /// @param end_col      Ending byte offset into line (0-based, exclusive)
-/// @param replace_nl   Replace newlines ('\n') with null ('\0')
 /// @param err          Error object
 /// @return The text between start_col and end_col on line lnum of buffer buf
-String buf_get_text(buf_T *buf, int64_t lnum, int64_t start_col, int64_t end_col, bool replace_nl,
-                    Error *err)
+String buf_get_text(buf_T *buf, int64_t lnum, int64_t start_col, int64_t end_col, Error *err)
 {
   String rv = STRING_INIT;
 
@@ -517,7 +485,7 @@ String buf_get_text(buf_T *buf, int64_t lnum, int64_t start_col, int64_t end_col
     return rv;
   }
 
-  const char *bufstr = ml_get_buf(buf, (linenr_T)lnum, false);
+  char *bufstr = ml_get_buf(buf, (linenr_T)lnum, false);
   size_t line_length = strlen(bufstr);
 
   start_col = start_col < 0 ? (int64_t)line_length + start_col + 1 : start_col;
@@ -537,12 +505,7 @@ String buf_get_text(buf_T *buf, int64_t lnum, int64_t start_col, int64_t end_col
     return rv;
   }
 
-  rv = cstrn_to_string(&bufstr[start_col], (size_t)(end_col - start_col));
-  if (replace_nl) {
-    strchrsub(rv.data, '\n', '\0');
-  }
-
-  return rv;
+  return cstrn_as_string((char *)&bufstr[start_col], (size_t)(end_col - start_col));
 }
 
 void api_free_string(String value)
diff --git a/src/nvim/buffer_updates.c b/src/nvim/buffer_updates.c
index 1b3c0bc28f..681d5df047 100644
--- a/src/nvim/buffer_updates.c
+++ b/src/nvim/buffer_updates.c
@@ -1,6 +1,7 @@
 // This is an open source non-commercial project. Dear PVS-Studio, please check
 // it. PVS-Studio Static Code Analyzer for C, C++ and C#: http://www.viva64.com
 
+#include "nvim/api/buffer.h"
 #include "nvim/api/private/helpers.h"
 #include "nvim/assert.h"
 #include "nvim/buffer.h"
@@ -34,12 +35,10 @@ bool buf_updates_register(buf_T *buf, uint64_t channel_id, BufUpdateCallbacks cb
 
   // count how many channels are currently watching the buffer
   size_t size = kv_size(buf->update_channels);
-  if (size) {
-    for (size_t i = 0; i < size; i++) {
-      if (kv_A(buf->update_channels, i) == channel_id) {
-        // buffer is already registered ... nothing to do
-        return true;
-      }
+  for (size_t i = 0; i < size; i++) {
+    if (kv_A(buf->update_channels, i) == channel_id) {
+      // buffer is already registered ... nothing to do
+      return true;
     }
   }
 
@@ -69,7 +68,7 @@ bool buf_updates_register(buf_T *buf, uint64_t channel_id, BufUpdateCallbacks cb
       linedata.size = line_count;
       linedata.items = xcalloc(line_count, sizeof(Object));
 
-      buf_collect_lines(buf, line_count, 1, true, &linedata, NULL);
+      buf_collect_lines(buf, line_count, 1, true, &linedata, NULL, NULL);
     }
 
     args.items[4] = ARRAY_OBJ(linedata);
@@ -231,7 +230,7 @@ void buf_updates_send_changes(buf_T *buf, linenr_T firstline, int64_t num_added,
       linedata.size = (size_t)num_added;
       linedata.items = xcalloc((size_t)num_added, sizeof(Object));
       buf_collect_lines(buf, (size_t)num_added, firstline, true, &linedata,
-                        NULL);
+                        NULL, NULL);
     }
     args.items[4] = ARRAY_OBJ(linedata);
     args.items[5] = BOOLEAN_OBJ(false);
diff --git a/src/nvim/generators/c_grammar.lua b/src/nvim/generators/c_grammar.lua
index d872ffd6a9..3ea5904338 100644
--- a/src/nvim/generators/c_grammar.lua
+++ b/src/nvim/generators/c_grammar.lua
@@ -27,6 +27,7 @@ local c_void = P('void')
 local c_param_type = (
   ((P('Error') * fill * P('*') * fill) * Cc('error')) +
   ((P('Arena') * fill * P('*') * fill) * Cc('arena')) +
+  ((P('lua_State') * fill * P('*') * fill) * Cc('lstate')) +
   C((P('const ') ^ -1) * (c_id) * (ws ^ 1) * P('*')) +
   (C(c_id) * (ws ^ 1))
   )
diff --git a/src/nvim/generators/gen_api_dispatch.lua b/src/nvim/generators/gen_api_dispatch.lua
index 67b8f5f0f5..6894b7a6df 100644
--- a/src/nvim/generators/gen_api_dispatch.lua
+++ b/src/nvim/generators/gen_api_dispatch.lua
@@ -78,6 +78,10 @@ for i = 6, #arg do
         fn.arena_return = true
         fn.parameters[#fn.parameters] = nil
       end
+      if #fn.parameters ~= 0 and fn.parameters[#fn.parameters][1] == 'lstate' then
+        fn.has_lua_imp = true
+        fn.parameters[#fn.parameters] = nil
+      end
     end
   end
   input:close()
@@ -329,6 +333,14 @@ for i = 1, #functions do
         output:write(', arena')
     end
 
+    if fn.has_lua_imp then
+      if #args > 0 then
+        output:write(', NULL')
+      else
+        output:write('NULL')
+      end
+    end
+
     if fn.can_fail then
       -- if the function can fail, also pass a pointer to the local error object
       if #args > 0 then
@@ -497,6 +509,10 @@ local function process_function(fn)
     ]])
   end
 
+  if fn.has_lua_imp then
+    cparams = cparams .. 'lstate, '
+  end
+
   if fn.can_fail then
     cparams = cparams .. '&err'
   else
@@ -539,13 +555,27 @@ local function process_function(fn)
     end
     write_shifted_output(output, string.format([[
     const %s ret = %s(%s);
+    ]], fn.return_type, fn.name, cparams))
+
+    if fn.has_lua_imp then
+      -- only push onto the Lua stack if we haven't already
+      write_shifted_output(output, string.format([[
+    if (lua_gettop(lstate) == 0) {
+      nlua_push_%s(lstate, ret, true);
+    }
+      ]], return_type))
+    else
+      write_shifted_output(output, string.format([[
     nlua_push_%s(lstate, ret, true);
+      ]], return_type))
+    end
+
+    write_shifted_output(output, string.format([[
   %s
   %s
   %s
     return 1;
-    ]], fn.return_type, fn.name, cparams, return_type,
-        free_retval, free_at_exit_code, err_throw_code))
+    ]], free_retval, free_at_exit_code, err_throw_code))
   else
     write_shifted_output(output, string.format([[
     %s(%s);
diff --git a/src/nvim/option.c b/src/nvim/option.c
index 23c4a1ccf3..330900a9d6 100644
--- a/src/nvim/option.c
+++ b/src/nvim/option.c
@@ -3012,7 +3012,7 @@ int get_option_value_strict(char *name, int64_t *numval, char **stringval, int o
 
   if (varp != NULL) {
     if (p->flags & P_STRING) {
-      *stringval = xstrdup(*(char **)(varp));
+      *stringval = *(char **)(varp);
     } else if (p->flags & P_NUM) {
       *numval = *(long *)varp;
     } else {
diff --git a/test/functional/api/buffer_spec.lua b/test/functional/api/buffer_spec.lua
index 598bfeb754..5980d99f97 100644
--- a/test/functional/api/buffer_spec.lua
+++ b/test/functional/api/buffer_spec.lua
@@ -186,12 +186,13 @@ describe('api/buf', function()
     end)
   end)
 
-  describe('nvim_buf_get_lines, nvim_buf_set_lines', function()
-    local get_lines, set_lines = curbufmeths.get_lines, curbufmeths.set_lines
-    local line_count = curbufmeths.line_count
+  describe_lua_and_rpc('nvim_buf_get_lines, nvim_buf_set_lines', function(api)
+    local get_lines = api.curbufmeths.get_lines
+    local set_lines = api.curbufmeths.set_lines
+    local line_count = api.curbufmeths.line_count
 
     it('fails correctly when input is not valid', function()
-      eq(1, curbufmeths.get_number())
+      eq(1, api.curbufmeths.get_number())
       eq([[String cannot contain newlines]],
         pcall_err(bufmeths.set_lines, 1, 1, 2, false, {'b\na'}))
     end)
@@ -199,7 +200,7 @@ describe('api/buf', function()
     it("fails if 'nomodifiable'", function()
       command('set nomodifiable')
       eq([[Buffer is not 'modifiable']],
-        pcall_err(bufmeths.set_lines, 1, 1, 2, false, {'a','b'}))
+        pcall_err(api.bufmeths.set_lines, 1, 1, 2, false, {'a','b'}))
     end)
 
     it('has correct line_count when inserting and deleting', function()
@@ -355,7 +356,7 @@ describe('api/buf', function()
         Who would win?
         A real window
         with proper text]])
-      local buf = meths.create_buf(false,true)
+      local buf = api.meths.create_buf(false,true)
       screen:expect([[
         Who would win?      |
         A real window       |
@@ -364,7 +365,7 @@ describe('api/buf', function()
                             |
       ]])
 
-      meths.buf_set_lines(buf, 0, -1, true, {'or some', 'scratchy text'})
+      api.meths.buf_set_lines(buf, 0, -1, true, {'or some', 'scratchy text'})
       feed('i') -- provoke redraw
       screen:expect([[
         Who would win?      |
@@ -380,15 +381,15 @@ describe('api/buf', function()
         visible buffer line 1
         line 2
       ]])
-      local hiddenbuf = meths.create_buf(false,true)
+      local hiddenbuf = api.meths.create_buf(false,true)
       command('vsplit')
       command('vsplit')
       feed('<c-w>l<c-w>l<c-w>l')
       eq(3, funcs.winnr())
       feed('<c-w>h')
       eq(2, funcs.winnr())
-      meths.buf_set_lines(hiddenbuf, 0, -1, true,
-                          {'hidden buffer line 1', 'line 2'})
+      api.meths.buf_set_lines(hiddenbuf, 0, -1, true,
+                              {'hidden buffer line 1', 'line 2'})
       feed('<c-w>p')
       eq(3, funcs.winnr())
     end)