Reuse LuaState instances, copy hook and return allocator from global state

Author: Quahu

src/Laylua/Library/Lua.cs

@@ -12,8 +12,14 @@ namespace Laylua;
 /// <remarks>
 ///     This type is not thread-safe; operations on it are not thread-safe.
 /// </remarks>
-public unsafe partial class Lua : LuaThread, ISpanFormattable
+public sealed unsafe partial class Lua : LuaThread, ISpanFormattable
 {
+    /// <inheritdoc/>
+    public override LuaState State { get; }
+
+    /// <inheritdoc/>
+    public override LuaStack Stack { get; }
+
     /// <inheritdoc/>
     public override LuaThread MainThread { get; }
 
@@ -48,8 +54,11 @@ public Lua(LuaAllocator allocator)
     public Lua(LuaAllocator allocator, LuaMarshaler marshaler)
     {
         Stack = new LuaStack(this);
-        State = new LuaState(allocator);
-        State.State = this;
+        State = new LuaMainState(allocator)
+        {
+            State = this
+        };
+
         Reference = LuaRegistry.Indices.MainThread;
         MainThread = new LuaMainThread(this);
         Globals = LuaTable.CreateGlobalsTable(this);

src/Laylua/Library/LuaChildThread.cs

@@ -7,6 +7,12 @@ namespace Laylua;
 
 internal sealed unsafe class LuaChildThread : LuaThread
 {
+    /// <inheritdoc />
+    public override LuaState State => _state;
+
+    /// <inheritdoc />
+    public override LuaStack Stack { get; }
+
     /// <inheritdoc/>
     public override LuaThread MainThread => Lua.FromExtraSpace(State.L).MainThread;
 
@@ -23,10 +29,25 @@ protected override Lua? LuaCore
         set => throw new InvalidOperationException();
     }
 
-    internal void Initialize(lua_State* L, int reference)
+    private readonly LuaChildState _state;
+
+    public LuaChildThread()
     {
         Stack = new LuaStack(this);
-        State = new LuaState(L);
+        _state = new LuaChildState();
+    }
+
+    internal void Initialize(lua_State* L, int reference)
+    {
         Reference = reference;
+
+        _state.Initialize(Lua.FromExtraSpace(L).State, L);
+    }
+
+    public override void Dispose()
+    {
+        base.Dispose();
+
+        _state.Close();
     }
 }

src/Laylua/Library/LuaMainThread.cs

@@ -7,6 +7,12 @@ namespace Laylua;
 
 internal sealed class LuaMainThread : LuaThread
 {
+    /// <inheritdoc />
+    public override LuaState State => _lua.State;
+
+    /// <inheritdoc />
+    public override LuaStack Stack => _lua.Stack;
+
     public override LuaThread MainThread => this;
 
     public override LuaTable Globals => _lua.Globals;
@@ -27,8 +33,6 @@ protected override Lua? LuaCore
     public LuaMainThread(Lua lua)
     {
         _lua = lua;
-        Stack = lua.Stack;
-        State = lua.State;
         Reference = LuaRegistry.Indices.MainThread;
     }
 }

src/Laylua/Library/LuaThread.cs

@@ -20,15 +20,15 @@ public abstract unsafe partial class LuaThread : LuaReference
     /// <summary>
     ///     Gets the low-level state of this thread.
     /// </summary>
-    public LuaState State { get; protected set; } = null!;
+    public abstract LuaState State { get; }
 
     /// <summary>
     ///     Gets the Lua stack of this thread.
     /// </summary>
     /// <remarks>
     ///     <inheritdoc cref="LuaStack"/>
     /// </remarks>
-    public LuaStack Stack { get; protected set; } = null!;
+    public abstract LuaStack Stack { get; }
 
     /// <summary>
     ///     Gets the main Lua thread, i.e. the parent of this thread.

src/Laylua/Moon/LuaChildState.cs

@@ -0,0 +1,27 @@
+namespace Laylua.Moon;
+
+internal sealed unsafe class LuaChildState : LuaState
+{
+    public override LuaAllocator? Allocator => _mainState.Allocator;
+
+    public override LuaGC GC => _mainState.GC;
+
+    protected override lua_State* LCore => _L;
+
+    private lua_State* _L;
+    private LuaState _mainState = null!;
+
+    public void Initialize(LuaState mainState, lua_State* childState)
+    {
+        _mainState = mainState;
+        _L = childState;
+
+        CopyHook(mainState);
+    }
+
+    internal override void Close()
+    {
+        // TODO: reset/close thread? Adjust LuaChildThread if so
+        _L = null;
+    }
+}

src/Laylua/Moon/LuaMainState.cs

@@ -0,0 +1,92 @@
+using System;
+using System.Globalization;
+using Qommon;
+
+namespace Laylua.Moon;
+
+internal sealed unsafe class LuaMainState : LuaState
+{
+    public override LuaAllocator? Allocator { get; }
+
+    public override LuaGC GC { get; }
+
+    protected override lua_State* LCore => _L;
+
+    internal object? State
+    {
+        get => _layluaState.State;
+        set => _layluaState.State = value;
+    }
+
+    private LuaAllocFunction? _safeAllocFunction;
+
+    private LayluaState _layluaState = null!;
+    private lua_State* _L;
+
+    internal LuaMainState(LuaAllocator? allocator)
+    {
+        if (allocator == null)
+        {
+            _L = luaL_newstate();
+        }
+        else
+        {
+            _L = lua_newstate(WrapAllocFunction(allocator), null);
+            Allocator = allocator;
+        }
+
+        ValidateNewState();
+        InitializeLayluaState();
+
+        GC = new(this);
+    }
+
+    private void ValidateNewState()
+    {
+        if (_L == null)
+        {
+            Throw.InvalidOperationException("Failed to create a new Lua state.");
+        }
+
+        var version = lua_version(_L).ToString(CultureInfo.InvariantCulture);
+        if (version.Length != 3 || version[0] != LuaVersionMajor[0] || version[2] != LuaVersionMinor[0])
+        {
+            Throw.InvalidOperationException($"Invalid Lua version loaded. Expected '{LuaVersion}'.");
+        }
+    }
+
+    private void InitializeLayluaState()
+    {
+        if (!LayluaState.TryFromExtraSpace(_L, out var state))
+        {
+            var panicPtr = LayluaNative.InitializePanic();
+            state = LayluaState.Create(panicPtr);
+            *(IntPtr*) lua_getextraspace(_L) = state.GCPtr;
+            lua_atpanic(_L, (delegate* unmanaged[Cdecl]<lua_State*, int>) panicPtr);
+        }
+
+        _layluaState = state;
+    }
+
+    internal LuaAllocFunction WrapAllocFunction(LuaAllocator allocator)
+    {
+        Guard.IsNotNull(allocator);
+
+        _safeAllocFunction = (_, ptr, osize, nsize) =>
+        {
+            return allocator.Allocate(ptr, osize, nsize);
+        };
+
+        return _safeAllocFunction;
+    }
+
+    internal override void Close()
+    {
+        if (_L == null)
+            return;
+
+        lua_close(_L);
+        _layluaState.Dispose();
+        _L = null;
+    }
+}