New helper to push function bytecode (to facilitate Lua 5.5 support)

3 files changed, 53 insertions, 38 deletions

diff --git a/src/intercopycontext.cpp b/src/intercopycontext.cpp
index 7be1326..653eeb6 100644
--- a/src/intercopycontext.cpp
+++ b/src/intercopycontext.cpp
@@ -36,23 +36,6 @@ THE SOFTWARE.
 
 // #################################################################################################
 
-// Lua 5.4.3 style of dumping (see lstrlib.c)
-// we have to do it that way because we can't unbalance the stack between buffer operations
-// namely, this means we can't push a function on top of the stack *after* we initialize the buffer!
-// luckily, this also works with earlier Lua versions
-[[nodiscard]]
-static int buf_writer(lua_State* L_, void const* b_, size_t size_, void* ud_)
-{
-    luaL_Buffer* const _B{ static_cast<luaL_Buffer*>(ud_) };
-    if (!_B->L) {
-        luaL_buffinit(L_, _B);
-    }
-    luaL_addlstring(_B, static_cast<char const*>(b_), size_);
-    return 0;
-}
-
-// #################################################################################################
-
 // function sentinel used to transfer native functions from/to keeper states
 [[nodiscard]]
 static int func_lookup_sentinel(lua_State* const L_)
@@ -211,19 +194,18 @@ void InterCopyContext::copyFunction() const
     // to the writer" (and we only return 0)
     // not sure this could ever fail but for memory shortage reasons
     // last argument is Lua 5.4-specific (no stripping)
-    luaL_Buffer B{};
-    if (luaW_dump(L1, buf_writer, &B, U->stripFunctions) != 0) {
+    if (tools::PushFunctionBytecode(L1, U->stripFunctions) != 0) {                                 // L1: ... f "<bytecode>"
         raise_luaL_error(getErrL(), "internal error: function dump failed.");
     }
 
-    // pushes dumped string on 'L1'
-    luaL_pushresult(&B);                                                                           // L1: ... f b
-
-    // if not pushed, no need to pop
+    // if pushed, we need to pop
     if (_needToPush) {
-        lua_remove(L1, -2);                                                                        // L1: ... b
+        lua_remove(L1, -2);                                                                        // L1: ... "<bytecode>"
     }
 
+    // When we are done, the stack should be the original one, with the bytecode string added on top
+    STACK_CHECK(L1, 1);
+
     // transfer the bytecode, then the upvalues, to create a similar closure
     {
         char const* _fname{};
@@ -231,16 +213,16 @@ void InterCopyContext::copyFunction() const
         if constexpr (LOG_FUNC_INFO)
         {
             lua_Debug _ar;
-            lua_pushvalue(L1, L1_i);                                                               // L1: ... b f
+            lua_pushvalue(L1, L1_i);                                                               // L1: ... "<bytecode>" f
             // "To get information about a function you push it onto the stack and start the what string with the character '>'."
             // fills 'fname' 'namewhat' and 'linedefined', pops function
-            lua_getinfo(L1, ">nS", &_ar);                                                          // L1: ... b
+            lua_getinfo(L1, ">nS", &_ar);                                                          // L1: ... "<bytecode>"
             _fname = _ar.namewhat;
             DEBUGSPEW_CODE(DebugSpew(U) << "FNAME: " << _ar.short_src << " @ " << _ar.linedefined << std::endl);
         }
 
         {
-            std::string_view const _bytecode{ luaW_tostring(L1, kIdxTop) };                        // L1: ... b
+            std::string_view const _bytecode{ luaW_tostring(L1, kIdxTop) };                        // L1: ... "<bytecode>"
             LUA_ASSERT(L1, !_bytecode.empty());
             STACK_GROW(L2, 2);
             // Note: Line numbers seem to be taken precisely from the
@@ -257,7 +239,7 @@ void InterCopyContext::copyFunction() const
                 raise_luaL_error(getErrL(), "%s: %s", _fname, lua_tostring(L2, -1));
             }
             // remove the dumped string
-            lua_pop(L1, 1); // ...
+            lua_pop(L1, 1);                                                                        // L1: ...
             // now set the cache as soon as we can.
             // this is necessary if one of the function's upvalues references it indirectly
             // we need to find it in the cache even if it isn't fully transfered yet
diff --git a/src/tools.cpp b/src/tools.cpp
index 25949e4..aafb9e8 100644
--- a/src/tools.cpp
+++ b/src/tools.cpp
@@ -45,13 +45,33 @@ static constexpr RegistryUniqueKey kLookupCacheRegKey{ 0x9BF75F84E54B691Bull };
 
 // #################################################################################################
 
-static constexpr int kWriterReturnCode{ 666 };
-[[nodiscard]]
-static int dummy_writer([[maybe_unused]] lua_State* const L_, [[maybe_unused]] void const* p_, [[maybe_unused]] size_t sz_, [[maybe_unused]] void* ud_)
-{
-    // always fail with this code
-    return kWriterReturnCode;
-}
+namespace {
+    namespace local {
+        static int buf_writer(lua_State* L_, void const* b_, size_t size_, void* ud_)
+        {
+            auto* const _B{ static_cast<luaL_Buffer*>(ud_) };
+            if (!_B->L) {
+                luaL_buffinit(L_, _B);
+            }
+            // starting with Lua 5.5, the writer is called one last time with nullptr, 0 at the end
+            if (b_ && size_) {
+                luaL_addlstring(_B, static_cast<char const*>(b_), size_);
+            }
+            return 0;
+        }
+
+        static constexpr int kWriterReturnCode{ 666 };
+        [[nodiscard]]
+        static int dummy_writer([[maybe_unused]] lua_State* const L_, [[maybe_unused]] void const* p_, [[maybe_unused]] size_t sz_, [[maybe_unused]] void* ud_)
+        {
+            // always fail with this code
+            return kWriterReturnCode;
+        }
+
+    } // namespace local
+} // namespace
+
+// #################################################################################################
 
 /*
  * differentiation between C, bytecode and JIT-fast functions
@@ -77,11 +97,11 @@ FuncSubType luaW_getfuncsubtype(lua_State* const L_, StackIndex const i_)
     // luaW_dump expects the function at the top of the stack
     int const _popCount{ (luaW_absindex(L_, i_) == lua_gettop(L_)) ? 0 : (lua_pushvalue(L_, i_), 1) };
     // here we either have a Lua bytecode or a LuaJIT-compiled function
-    int const _dumpres{ luaW_dump(L_, dummy_writer, nullptr, 0) };
+    int const _dumpres{ luaW_dump(L_, local::dummy_writer, nullptr, 0) };
     if (_popCount > 0) {
         lua_pop(L_, _popCount);
     }
-    if (_dumpres == kWriterReturnCode) {
+    if (_dumpres == local::kWriterReturnCode) {
         // anytime we get kWriterReturnCode, this means that luaW_dump() attempted a dump
         return FuncSubType::Bytecode;
     }
@@ -92,7 +112,6 @@ FuncSubType luaW_getfuncsubtype(lua_State* const L_, StackIndex const i_)
 // #################################################################################################
 
 namespace tools {
-
     // inspired from tconcat() in ltablib.c
     [[nodiscard]]
     std::string_view PushFQN(lua_State* const L_, StackIndex const t_)
@@ -118,6 +137,18 @@ namespace tools {
         return luaW_tostring(L_, kIdxTop);
     }
 
+    // #############################################################################################
+
+    [[nodiscard]]
+    int PushFunctionBytecode(lua_State* const L_, int const strip_)
+    {
+        luaL_Buffer B{};
+        int const result_{ luaW_dump(L_, local::buf_writer, &B, strip_) };
+        if (result_ == 0) { // documentation says it should always be the case (because our writer only ever returns 0), but better safe than sorry
+            luaL_pushresult(&B);
+        }
+        return result_;
+    }
 } // namespace tools
 
 // #################################################################################################
diff --git a/src/tools.hpp b/src/tools.hpp
index 14f9855..51cbb9b 100644
--- a/src/tools.hpp
+++ b/src/tools.hpp
@@ -37,5 +37,7 @@ namespace tools {
     void PopulateFuncLookupTable(lua_State* L_, StackIndex i_, std::string_view const& name_);
     [[nodiscard]]
     std::string_view PushFQN(lua_State* L_, StackIndex t_);
+    [[nodiscard]]
+    int PushFunctionBytecode(lua_State* L_, int strip_);
     void SerializeRequire(lua_State* L_);
 } // namespace tools