5 files changed, 84 insertions, 76 deletions

diff --git a/src/keeper.cpp b/src/keeper.cpp
index 11f1797..cdbfac9 100644
--- a/src/keeper.cpp
+++ b/src/keeper.cpp
@@ -876,7 +876,7 @@ void Keepers::initialize(Universe& U_, lua_State* L_, int const nbKeepers_, int
         // attempt to call on_state_create(), if we have one and it is a C function
         // (only support a C function because we can't transfer executable Lua code in keepers)
         // will raise an error in L_ in case of problem
-        state::CallOnStateCreate(U, _K, L, LookupMode::ToKeeper);
+        U->callOnStateCreate(_K, L, LookupMode::ToKeeper);
 
         // _R[kLindasRegKey] = {}
         kLindasRegKey.setValue(_K, [](lua_State* L_) { lua_newtable(L_); });
diff --git a/src/state.cpp b/src/state.cpp
index 1975148..0757203 100644
--- a/src/state.cpp
+++ b/src/state.cpp
@@ -41,10 +41,6 @@ THE SOFTWARE.
 #include "universe.h"
 
 // #################################################################################################
-
-static constexpr std::string_view kOnStateCreate{ "on_state_create" }; // update lanes.lua if the name changes!
-
-// #################################################################################################
 // #################################################################################################
 namespace {
     // #############################################################################################
@@ -157,42 +153,6 @@ namespace state {
     // #############################################################################################
     // #############################################################################################
 
-    void CallOnStateCreate(Universe* const U_, lua_State* const L_, lua_State* const from_, LookupMode const mode_)
-    {
-        if (U_->onStateCreateFunc == nullptr) {
-            return;
-        }
-
-        STACK_CHECK_START_REL(L_, 0);
-        DEBUGSPEW_CODE(DebugSpew(U_) << "calling on_state_create()" << std::endl);
-        if (U_->onStateCreateFunc != reinterpret_cast<lua_CFunction>(InitializeOnStateCreate)) {
-            // C function: recreate a closure in the new state, bypassing the lookup scheme
-            lua_pushcfunction(L_, U_->onStateCreateFunc);                                          // on_state_create()
-        } else { // Lua function located in the config table, copied when we opened "lanes.core"
-            if (mode_ != LookupMode::LaneBody) {
-                // if attempting to call in a keeper state, do nothing because the function doesn't exist there
-                // this doesn't count as an error though
-                STACK_CHECK(L_, 0);
-                return;
-            }
-            kConfigRegKey.pushValue(L_);                                                           // L_: {}
-            STACK_CHECK(L_, 1);
-            LuaType const _funcType{ luaG_getfield(L_, -1, kOnStateCreate) };                      // L_: {} on_state_create()
-            if (_funcType != LuaType::FUNCTION) {
-                raise_luaL_error(L_, "INTERNAL ERROR: %s is a %s, not a function", kOnStateCreate.data(), luaG_typename(L_, _funcType).data());
-            }
-            lua_remove(L_, -2);                                                                    // L_: on_state_create()
-        }
-        STACK_CHECK(L_, 1);
-        // capture error and raise it in caller state
-        std::string_view const _stateType{ mode_ == LookupMode::LaneBody ? "lane" : "keeper" };
-        luaG_pushstring(L_, _stateType);                                                           // L_: on_state_create() "<type>"
-        if (lua_pcall(L_, 1, 0, 0) != LUA_OK) {
-            raise_luaL_error(from_, "%s failed: \"%s\"", kOnStateCreate.data(), lua_isstring(L_, -1) ? luaG_tostring(L_, -1).data() : luaG_typename(L_, -1).data());
-        }
-        STACK_CHECK(L_, 0);
-    }
-
     // #############################################################################################
 
     lua_State* CreateState([[maybe_unused]] Universe* const U_, lua_State* const from_, std::string_view const& hint_)
@@ -219,33 +179,6 @@ namespace state {
 
     // #############################################################################################
 
-    void InitializeOnStateCreate(Universe* const U_, lua_State* const L_)
-    {
-        STACK_CHECK_START_REL(L_, 1);                                                              // L_: settings
-        if (luaG_getfield(L_, -1, kOnStateCreate) != LuaType::NIL) {                               // L_: settings on_state_create|nil
-            // store C function pointer in an internal variable
-            U_->onStateCreateFunc = lua_tocfunction(L_, -1);                                       // L_: settings on_state_create
-            if (U_->onStateCreateFunc != nullptr) {
-                // make sure the function doesn't have upvalues
-                char const* _upname{ lua_getupvalue(L_, -1, 1) };                                  // L_: settings on_state_create upval?
-                if (_upname != nullptr) { // should be "" for C functions with upvalues if any
-                    raise_luaL_error(L_, "%s shouldn't have upvalues", kOnStateCreate.data());
-                }
-                // remove this C function from the config table so that it doesn't cause problems
-                // when we transfer the config table in newly created Lua states
-                lua_pushnil(L_);                                                                   // L_: settings on_state_create nil
-                luaG_setfield(L_, -3, kOnStateCreate);                                             // L_: settings on_state_create
-            } else {
-                // optim: store marker saying we have such a function in the config table
-                U_->onStateCreateFunc = reinterpret_cast<lua_CFunction>(InitializeOnStateCreate);
-            }
-        }
-        lua_pop(L_, 1);                                                                            // L_: settings
-        STACK_CHECK(L_, 1);
-    }
-
-    // #############################################################################################
-
     /*
      * Like 'luaL_openlibs()' but allows the set of libraries be selected
      *
@@ -274,7 +207,7 @@ namespace state {
         STACK_CHECK(_L, 0);
 
         // neither libs (not even 'base') nor special init func: we are done
-        if (!libs_.has_value() && U_->onStateCreateFunc == nullptr) {
+        if (!libs_.has_value() && std::holds_alternative<std::nullptr_t>(U_->onStateCreateFunc)) {
             DEBUGSPEW_CODE(DebugSpew(U_) << "luaG_newstate(nullptr)" << std::endl);
             return _L;
         }
@@ -342,7 +275,7 @@ namespace state {
 
         // call this after the base libraries are loaded and GC is restarted
         // will raise an error in from_ in case of problem
-        CallOnStateCreate(U_, _L, from_, LookupMode::LaneBody);
+        U_->callOnStateCreate(_L, from_, LookupMode::LaneBody);
 
         STACK_CHECK(_L, 0);
         // after all this, register everything we find in our name<->function database
diff --git a/src/state.h b/src/state.h
index 98bb47b..f809d50 100644
--- a/src/state.h
+++ b/src/state.h
@@ -8,10 +8,6 @@ enum class LookupMode;
 class Universe;
 
 namespace state {
-
-    void CallOnStateCreate(Universe* U_, lua_State* L_, lua_State* from_, LookupMode mode_);
     [[nodiscard]] lua_State* CreateState(Universe* U_, lua_State* from_, std::string_view const& hint_);
-    void InitializeOnStateCreate(Universe* U_, lua_State* L_);
     [[nodiscard]] lua_State* NewLaneState(Universe* U_, SourceState from_, std::optional<std::string_view> const& libs_);
-
 } // namespace state
diff --git a/src/universe.cpp b/src/universe.cpp
index 097c642..b63008a 100644
--- a/src/universe.cpp
+++ b/src/universe.cpp
@@ -41,6 +41,10 @@ extern LUAG_FUNC(linda);
 
 // #################################################################################################
 
+static constexpr std::string_view kOnStateCreate{ "on_state_create" }; // update lanes.lua if the name changes!
+
+// #################################################################################################
+
 // xxh64 of string "kUniverseFullRegKey" generated at https://www.pelock.com/products/hash-calculator
 static constexpr RegistryUniqueKey kUniverseFullRegKey{ 0x1C2D76870DD9DD9Full };
 
@@ -74,6 +78,48 @@ Universe::Universe()
 
 // #################################################################################################
 
+void Universe::callOnStateCreate(lua_State* const L_, lua_State* const from_, LookupMode const mode_)
+{
+    if (std::holds_alternative<std::nullptr_t>(onStateCreateFunc)) {
+        return;
+    }
+
+    STACK_CHECK_START_REL(L_, 0);
+    DEBUGSPEW_CODE(DebugSpew(U_) << "calling on_state_create()" << std::endl);
+    if (std::holds_alternative<lua_CFunction>(onStateCreateFunc)) {
+            
+        // C function: recreate a closure in the new state, bypassing the lookup scheme
+        lua_pushcfunction(L_, std::get<lua_CFunction>(onStateCreateFunc));                         // on_state_create()
+    } else { // Lua function located in the config table, copied when we opened "lanes.core"
+        LUA_ASSERT(from_, std::holds_alternative<uintptr_t>(onStateCreateFunc));
+        if (mode_ != LookupMode::LaneBody) {
+            // if attempting to call in a keeper state, do nothing because the function doesn't exist there
+            // this doesn't count as an error though
+            STACK_CHECK(L_, 0);
+            return;
+        }
+        kConfigRegKey.pushValue(L_);                                                               // L_: config
+        STACK_CHECK(L_, 1);
+        LuaType const _funcType{ luaG_getfield(L_, -1, kOnStateCreate) };                          // L_: config on_state_create()
+        if (_funcType != LuaType::FUNCTION) {
+            raise_luaL_error(L_, "INTERNAL ERROR: %s is a %s, not a function", kOnStateCreate.data(), luaG_typename(L_, _funcType).data());
+        }
+        lua_remove(L_, -2);                                                                        // L_: on_state_create()
+    }
+    STACK_CHECK(L_, 1);
+    // capture error and raise it in caller state
+    std::string_view const _stateType{ mode_ == LookupMode::LaneBody ? "lane" : "keeper" };
+    luaG_pushstring(L_, _stateType);                                                               // L_: on_state_create() "<type>"
+    if (lua_pcall(L_, 1, 0, 0) != LUA_OK) {
+        raise_luaL_error(from_, "%s failed: \"%s\"", kOnStateCreate.data(), lua_isstring(L_, -1) ? luaG_tostring(L_, -1).data() : luaG_typename(L_, -1).data());
+    }
+    STACK_CHECK(L_, 0);
+}
+
+
+
+// #################################################################################################
+
 // only called from the master state
 [[nodiscard]] Universe* Universe::Create(lua_State* const L_)
 {
@@ -125,7 +171,7 @@ Universe::Universe()
     // Linked chains handling
     _U->selfdestructFirst = SELFDESTRUCT_END;
     _U->initializeAllocatorFunction(L_);
-    state::InitializeOnStateCreate(_U, L_);
+    _U->initializeOnStateCreate(L_);
     _U->keepers.initialize(*_U, L_, _nbUserKeepers, _keepers_gc_threshold);
     STACK_CHECK(L_, 0);
 
@@ -252,6 +298,37 @@ int Universe::InitializeFinalizer(lua_State* const L_)
 
 // #################################################################################################
 
+void Universe::initializeOnStateCreate(lua_State* const L_)
+{
+    STACK_CHECK_START_REL(L_, 1);                                                                  // L_: settings
+    if (luaG_getfield(L_, -1, kOnStateCreate) != LuaType::NIL) {                                   // L_: settings on_state_create|nil
+        LUA_ASSERT(L_, luaG_type(L_, -1) == LuaType::FUNCTION); // ensured by lanes.lua parameter validation
+        // make sure the function doesn't have upvalues
+        char const* _upname{ lua_getupvalue(L_, -1, 1) };                                          // L_: settings on_state_create upval?
+        if (_upname != nullptr) { // should be "" for C functions with upvalues if any
+            raise_luaL_error(L_, "%s with upvalues are forbidden", kOnStateCreate.data());
+        }
+        // store C function pointer in an internal variable
+        lua_CFunction const _func{ lua_tocfunction(L_, -1) };                                      // L_: settings on_state_create
+        if (_func) {
+            onStateCreateFunc.emplace<lua_CFunction>(_func);
+            // remove this C function from the config table so that it doesn't cause problems
+            // when we transfer the config table in newly created Lua states
+            lua_pushnil(L_);                                                                       // L_: settings on_state_create nil
+            luaG_setfield(L_, -3, kOnStateCreate);                                                 // L_: settings on_state_create
+        } else {
+            // the function is still in the config table
+            onStateCreateFunc.emplace<uintptr_t>(std::bit_cast<uintptr_t>(kOnStateCreate.data()));
+        }
+    } else {
+        LUA_ASSERT(L_, std::holds_alternative<std::nullptr_t>(onStateCreateFunc));
+    };
+    lua_pop(L_, 1);                                                                                // L_: settings
+    STACK_CHECK(L_, 1);
+}
+
+// #################################################################################################
+
 lanes::AllocatorDefinition Universe::resolveAllocator(lua_State* const L_, std::string_view const& hint_) const
 {
     lanes::AllocatorDefinition _ret{ protectedAllocator };
diff --git a/src/universe.h b/src/universe.h
index 6ccdaca..6374648 100644
--- a/src/universe.h
+++ b/src/universe.h
@@ -84,7 +84,7 @@ class Universe
     lua_CFunction provideAllocator{ nullptr };
 
     // after a state is created, this function will be called right after the bases libraries are loaded
-    lua_CFunction onStateCreateFunc{ nullptr };
+    std::variant<std::nullptr_t, uintptr_t, lua_CFunction> onStateCreateFunc;
 
     // if allocator="protected" is found in the configuration settings, a wrapper allocator will protect all allocator calls with a mutex
     // contains a mutex and the original allocator definition
@@ -131,10 +131,12 @@ class Universe
     Universe& operator=(Universe const&) = delete;
     Universe& operator=(Universe&&) = delete;
 
+    void callOnStateCreate(lua_State* const L_, lua_State* const from_, LookupMode const mode_);
     [[nodiscard]] static Universe* Create(lua_State* L_);
     [[nodiscard]] static inline Universe* Get(lua_State* L_);
     void initializeAllocatorFunction(lua_State* L_);
     static int InitializeFinalizer(lua_State* L_);
+    void initializeOnStateCreate(lua_State* const 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_);