Give a chance to config.allocator to provide a specific allocator for internal stuff

4 files changed, 33 insertions, 24 deletions

diff --git a/src/allocator.h b/src/allocator.h
index e09ab38..0505251 100644
--- a/src/allocator.h
+++ b/src/allocator.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include "compat.h"
+
 // #################################################################################################
 
 namespace lanes {
diff --git a/src/state.cpp b/src/state.cpp
index 6d3910b..1975148 100644
--- a/src/state.cpp
+++ b/src/state.cpp
@@ -204,23 +204,8 @@ namespace state {
                         // for some reason, LuaJIT 64 bits does not support creating a state with lua_newstate...
                         return luaL_newstate();
                     } else {
-                        if (U->provideAllocator != nullptr) { // we have a function we can call to obtain an allocator
-                            STACK_CHECK_START_REL(from, 0);
-                            lua_pushcclosure(from, U->provideAllocator, 0);                        // L: provideAllocator()
-                            luaG_pushstring(from, hint);                                           // L: provideAllocator() "<hint>"
-                            lua_call(from, 1, 1);                                                  // L: result
-                            lanes::AllocatorDefinition* const _def{ luaG_tofulluserdata<lanes::AllocatorDefinition>(from, -1) };
-                            if (!_def || _def->version != lanes::AllocatorDefinition::kAllocatorVersion) {
-                                raise_luaL_error(from, "Bad config.allocator function, must provide a valid AllocatorDefinition");
-                            }
-                            lua_State* const _L{ lua_newstate(_def->allocF, _def->allocUD) };
-                            lua_pop(from, 1);                                                      // L:
-                            STACK_CHECK(from, 0);
-                            return _L;
-                        } else {
-                            // reuse the allocator provided when the master state was created
-                            return lua_newstate(U->protectedAllocator.allocF, U->protectedAllocator.allocUD);
-                        }
+                        lanes::AllocatorDefinition const _def{ U->resolveAllocator(from, hint) };
+                        return lua_newstate(_def.allocF, _def.allocUD);
                     }
                 }
             )
diff --git a/src/universe.cpp b/src/universe.cpp
index 770fccf..097c642 100644
--- a/src/universe.cpp
+++ b/src/universe.cpp
@@ -176,7 +176,7 @@ Universe::Universe()
 // Do I need to disable this when compiling for LuaJIT to prevent issues?
 void Universe::initializeAllocatorFunction(lua_State* const L_)
 {
-    // start by just grabbing whatever allocator was provided to lua_newstate
+    // start by just grabbing whatever allocator was provided to the master state
     protectedAllocator.initFrom(L_);
     STACK_CHECK_START_REL(L_, 1);                                                                  // L_: settings
     switch (luaG_getfield(L_, -1, "allocator")) {                                                  // L_: settings allocator|nil|"protected"
@@ -212,19 +212,17 @@ void Universe::initializeAllocatorFunction(lua_State* const L_)
     default: // should be filtered out in lanes.lua
         raise_luaL_error(L_, "Bad config.allocator type %s", luaG_typename(L_, -1).data());
     }
-    lua_pop(L_, 1); // L_: settings
+    lua_pop(L_, 1);                                                                                // L_: settings
     STACK_CHECK(L_, 1);
 
     std::ignore = luaG_getfield(L_, -1, "internal_allocator");                                     // L_: settings "libc"|"allocator"
+    LUA_ASSERT(L_, lua_isstring(L_, -1)); // should be the case due to lanes.lua parameter validation
     std::string_view const _allocator{ luaG_tostring(L_, -1) };
     if (_allocator == "libc") {
         internalAllocator = lanes::AllocatorDefinition{ lanes::AllocatorDefinition::kAllocatorVersion, libc_lua_Alloc, nullptr };
-    } else if (provideAllocator == luaG_provide_protected_allocator) {
-        // user wants mutex protection on the state's allocator. Use protection for our own allocations too, just in case.
-        internalAllocator = protectedAllocator.makeDefinition();
     } else {
-        // no protection required, just use whatever we have as-is.
-        internalAllocator = protectedAllocator;
+        // use whatever the provider provides
+        internalAllocator = resolveAllocator(L_, "internal");
     }
     lua_pop(L_, 1);                                                                                // L_: settings
     STACK_CHECK(L_, 1);
@@ -254,6 +252,29 @@ int Universe::InitializeFinalizer(lua_State* const L_)
 
 // #################################################################################################
 
+lanes::AllocatorDefinition Universe::resolveAllocator(lua_State* const L_, std::string_view const& hint_) const
+{
+    lanes::AllocatorDefinition _ret{ protectedAllocator };
+    if (provideAllocator == nullptr) {
+        return _ret;
+    }
+
+    STACK_CHECK_START_REL(L_, 0); // here, we have a function we can call to obtain an allocator
+    lua_pushcclosure(L_, provideAllocator, 0);                                                     // L_: provideAllocator()
+    luaG_pushstring(L_, hint_);                                                                    // L_: provideAllocator() "<hint>"
+    lua_call(L_, 1, 1);                                                                            // L_: result
+    lanes::AllocatorDefinition* const _def{ luaG_tofulluserdata<lanes::AllocatorDefinition>(L_, -1) };
+    if (!_def || _def->version != lanes::AllocatorDefinition::kAllocatorVersion) {
+        raise_luaL_error(L_, "Bad config.allocator function, must provide a valid AllocatorDefinition");
+    }
+    _ret = *_def;
+    lua_pop(L_, 1);                                                                                // L_:
+    STACK_CHECK(L_, 0);
+    return _ret;
+}
+
+// #################################################################################################
+
 void Universe::terminateFreeRunningLanes(lua_State* const L_, lua_Duration const shutdownTimeout_, CancelOp const op_)
 {
     if (selfdestructFirst != SELFDESTRUCT_END) {
diff --git a/src/universe.h b/src/universe.h
index d890e5c..6ccdaca 100644
--- a/src/universe.h
+++ b/src/universe.h
@@ -135,6 +135,7 @@ class Universe
     [[nodiscard]] static inline Universe* Get(lua_State* L_);
     void initializeAllocatorFunction(lua_State* L_);
     static int InitializeFinalizer(lua_State* L_);
+    lanes::AllocatorDefinition resolveAllocator(lua_State* const L_, std::string_view const& hint_) const;
     static inline void Store(lua_State* L_, Universe* U_);
     void terminateFreeRunningLanes(lua_State* L_, lua_Duration shutdownTimeout_, CancelOp op_);
 };