Cleaner protocol between 'lua_dump' and writer function

'lua_dump' signals to the writer function the end of a dump, so that
is has more freedom when using the stack.

6 files changed, 66 insertions, 46 deletions

diff --git a/lapi.c b/lapi.c
index 374b3872..fd9ec4e4 100644
--- a/lapi.c
+++ b/lapi.c
@@ -1116,36 +1116,18 @@ LUA_API int lua_load (lua_State *L, lua_Reader reader, void *data,
 
 
 /*
-** Dump a function, calling 'writer' to write its parts. Because the
-** writer can use the stack in unkown ways, this function should not
-** push things on the stack, but it must anchor an auxiliary table
-** used by 'luaU_dump'. To do so, it creates the table, anchors the
-** function that is on the stack in the table, and substitutes the
-** table for the function in the stack.
+** Dump a Lua function, calling 'writer' to write its parts. Ensure
+** the stack returns with its original size.
 */
-
 LUA_API int lua_dump (lua_State *L, lua_Writer writer, void *data, int strip) {
   int status;
-  StkId fstk;  /* pointer to function */
-  TValue *o;
+  ptrdiff_t otop = savestack(L, L->top.p);  /* original top */
+  TValue *f = s2v(L->top.p - 1);  /* function to be dumped */
   lua_lock(L);
   api_checknelems(L, 1);
-  fstk = L->top.p - 1;
-  o = s2v(fstk);
-  if (!isLfunction(o))
-    status = 1;
-  else {
-    LClosure *f = clLvalue(o);
-    ptrdiff_t fidx = savestack(L, fstk);  /* function index */
-    Table *h = luaH_new(L);  /* auxiliary table used by 'luaU_dump' */
-    sethvalue2s(L, L->top.p, h);  /* anchor it (luaH_set may call GC) */
-    L->top.p++;  /* (assume extra slot) */
-    luaH_set(L, h, o, o);  /* anchor function into table */
-    setobjs2s(L, fstk, L->top.p - 1);  /* move table over function */
-    L->top.p--;  /* stack back to initial size */
-    status = luaU_dump(L, f->p, writer, data, strip, h);
-    setclLvalue2s(L, restorestack(L, fidx), f);  /* put function back */
-  }
+  api_check(L, isLfunction(f), "Lua function expected");
+  status = luaU_dump(L, clLvalue(f)->p, writer, data, strip);
+  L->top.p = restorestack(L, otop);  /* restore top */
   lua_unlock(L);
   return status;
 }
diff --git a/ldump.c b/ldump.c
index 9c315cdd..6cd5671f 100644
--- a/ldump.c
+++ b/ldump.c
@@ -43,8 +43,13 @@ typedef struct {
 #define dumpLiteral(D, s)       dumpBlock(D,s,sizeof(s) - sizeof(char))
 
 
+/*
+** Dump the block of memory pointed by 'b' with given 'size'.
+** 'b' should not be NULL, except for the last call signaling the end
+** of the dump.
+*/
 static void dumpBlock (DumpState *D, const void *b, size_t size) {
-  if (D->status == 0 && size > 0) {
+  if (D->status == 0) {  /* do not write anything after an error */
     lua_unlock(D->L);
     D->status = (*D->writer)(D->L, b, size, D->data);
     lua_lock(D->L);
@@ -53,13 +58,18 @@ static void dumpBlock (DumpState *D, const void *b, size_t size) {
 }
 
 
+/*
+** Dump enough zeros to ensure that current position is a multiple of
+** 'align'.
+*/
 static void dumpAlign (DumpState *D, int align) {
   int padding = align - (D->offset % align);
-  if (padding < align) {  /* apd == align means no padding */
+  if (padding < align) {  /* padding == align means no padding */
     static lua_Integer paddingContent = 0;
+    lua_assert(cast_uint(align) <= sizeof(lua_Integer));
     dumpBlock(D, &paddingContent, padding);
-    lua_assert(D->offset % align == 0);
   }
+  lua_assert(D->offset % align == 0);
 }
 
 
@@ -91,6 +101,7 @@ static void dumpSize (DumpState *D, size_t x) {
 
 
 static void dumpInt (DumpState *D, int x) {
+  lua_assert(x >= 0);
   dumpSize(D, x);
 }
 
@@ -140,6 +151,7 @@ static void dumpString (DumpState *D, TString *ts) {
 static void dumpCode (DumpState *D, const Proto *f) {
   dumpInt(D, f->sizecode);
   dumpAlign(D, sizeof(f->code[0]));
+  lua_assert(f->code != NULL);
   dumpVector(D, f->code, f->sizecode);
 }
 
@@ -196,7 +208,8 @@ static void dumpDebug (DumpState *D, const Proto *f) {
   int i, n;
   n = (D->strip) ? 0 : f->sizelineinfo;
   dumpInt(D, n);
-  dumpVector(D, f->lineinfo, n);
+  if (f->lineinfo != NULL)
+    dumpVector(D, f->lineinfo, n);
   n = (D->strip) ? 0 : f->sizeabslineinfo;
   dumpInt(D, n);
   for (i = 0; i < n; i++) {
@@ -248,20 +261,23 @@ static void dumpHeader (DumpState *D) {
 /*
 ** dump Lua function as precompiled chunk
 */
-int luaU_dump(lua_State *L, const Proto *f, lua_Writer w, void *data,
-              int strip, Table *h) {
+int luaU_dump (lua_State *L, const Proto *f, lua_Writer w, void *data,
+               int strip) {
   DumpState D;
+  D.h = luaH_new(L);  /* aux. table to keep strings already dumped */
+  sethvalue2s(L, L->top.p, D.h);  /* anchor it */
+  L->top.p++;
   D.L = L;
   D.writer = w;
   D.offset = 0;
   D.data = data;
   D.strip = strip;
   D.status = 0;
-  D.h = h;
   D.nstr = 0;
   dumpHeader(&D);
   dumpByte(&D, f->sizeupvalues);
   dumpFunction(&D, f);
+  dumpBlock(&D, NULL, 0);  /* signal end of dump */
   return D.status;
 }
 
diff --git a/lstrlib.c b/lstrlib.c
index e29c09b9..a90c4fd1 100644
--- a/lstrlib.c
+++ b/lstrlib.c
@@ -225,7 +225,12 @@ static int writer (lua_State *L, const void *b, size_t size, void *ud) {
     state->init = 1;
     luaL_buffinit(L, &state->B);
   }
-  luaL_addlstring(&state->B, (const char *)b, size);
+  if (b == NULL) {  /* finishing dump? */
+    luaL_pushresult(&state->B);  /* push result */
+    lua_replace(L, 1);  /* move it to reserved slot */
+  }
+  else
+    luaL_addlstring(&state->B, (const char *)b, size);
   return 0;
 }
 
@@ -233,13 +238,13 @@ static int writer (lua_State *L, const void *b, size_t size, void *ud) {
 static int str_dump (lua_State *L) {
   struct str_Writer state;
   int strip = lua_toboolean(L, 2);
-  luaL_checktype(L, 1, LUA_TFUNCTION);
-  lua_settop(L, 1);  /* ensure function is on the top of the stack */
+  luaL_argcheck(L, lua_type(L, 1) == LUA_TFUNCTION && !lua_iscfunction(L, 1),
+                   1, "Lua function expected");
+  /* ensure function is on the top of the stack and vacate slot 1 */
+  lua_pushvalue(L, 1);
   state.init = 0;
-  if (l_unlikely(lua_dump(L, writer, &state, strip) != 0))
-    return luaL_error(L, "unable to dump given function");
-  luaL_pushresult(&state.B);
-  lua_assert(lua_isfunction(L, 1));  /* lua_dump kept that value */
+  lua_dump(L, writer, &state, strip);
+  lua_settop(L, 1);  /* leave final result on top */
   return 1;
 }
 
diff --git a/lundump.c b/lundump.c
index 3f18343a..f850dc4a 100644
--- a/lundump.c
+++ b/lundump.c
@@ -152,7 +152,7 @@ static void loadString (LoadState *S, Proto *p, TString **sl) {
     luaH_getint(S->h, idx, &stv);
     *sl = ts = tsvalue(&stv);
     luaC_objbarrier(L, p, ts);
-    return;
+    return;  /* do not save it again */
   }
   else if ((size -= 2) <= LUAI_MAXSHORTLEN) {  /* short string? */
     char buff[LUAI_MAXSHORTLEN + 1];  /* extra space for '\0' */
@@ -168,10 +168,10 @@ static void loadString (LoadState *S, Proto *p, TString **sl) {
   else {  /* create internal copy */
     *sl = ts = luaS_createlngstrobj(L, size);  /* create string */
     luaC_objbarrier(L, p, ts);
-    loadVector(S, getlngstr(ts), size);  /* load directly in final place */
-    loadByte(S);  /* skip ending '\0' */
+    loadVector(S, getlngstr(ts), size + 1);  /* load directly in final place */
   }
-  S->nstr++;  /* add string to list of saved strings */
+  /* add string to list of saved strings */
+  S->nstr++;
   setsvalue(L, &sv, ts);
   luaH_setint(L, S->h, S->nstr, &sv);
   luaC_objbarrierback(L, obj2gco(S->h), ts);
diff --git a/lundump.h b/lundump.h
index 05ac7f85..b10307e4 100644
--- a/lundump.h
+++ b/lundump.h
@@ -31,6 +31,6 @@ LUAI_FUNC LClosure* luaU_undump (lua_State* L, ZIO* Z, const char* name,
 
 /* dump one chunk; from ldump.c */
 LUAI_FUNC int luaU_dump (lua_State* L, const Proto* f, lua_Writer w,
-                         void* data, int strip, Table *h);
+                         void* data, int strip);
 
 #endif
diff --git a/manual/manual.of b/manual/manual.of
index 8607e57d..ef1bdfd2 100644
--- a/manual/manual.of
+++ b/manual/manual.of
@@ -3266,6 +3266,13 @@ As it produces parts of the chunk,
 with the given @id{data}
 to write them.
 
+The function @Lid{lua_dump} fully preserves the Lua stack
+through the calls to the writer function,
+except that it may push some values for internal use
+before the first call,
+and it restores the stack size to its original size
+after the last call.
+
 If @id{strip} is true,
 the binary representation may not include all debug information
 about the function,
@@ -3275,8 +3282,6 @@ The value returned is the error code returned by the last
 call to the writer;
 @N{0 means} no errors.
 
-This function does not pop the Lua function from the stack.
-
 }
 
 @APIEntry{int lua_error (lua_State *L);|
@@ -4688,6 +4693,10 @@ passing along the buffer to be written (@id{p}),
 its size (@id{sz}),
 and the @id{ud} parameter supplied to @Lid{lua_dump}.
 
+After @Lid{lua_dump} writes its last piece,
+it will signal that by calling the writer function one more time,
+with a @id{NULL} buffer (and size 0).
+
 The writer returns an error code:
 @N{0 means} no errors;
 any other value means an error and stops @Lid{lua_dump} from
@@ -9260,6 +9269,14 @@ it is equivalent to @Lid{lua_closethread} with
 }
 
 @item{
+The function @Lid{lua_dump} changed the way it keeps the stack
+through the calls to the writer function.
+(That was not specified in previous versions.)
+Also, it calls the writer function one extra time,
+to signal the end of the dump.
+}
+
+@item{
 There were several changes in the parameters
 for the options @Lid{LUA_GCINC} and @Lid{LUA_GCGEN}
 of the function @Lid{lua_gc}.