Improved cancellation process and avoidance of infinite loops

src/Lua/Exceptions.cs

@@ -2,6 +2,7 @@
 using Lua.CodeAnalysis.Syntax;
 using Lua.Internal;
 using Lua.Runtime;
+using System.Diagnostics;
 using System.Runtime.CompilerServices;
 
 namespace Lua;
@@ -67,9 +68,14 @@ static string GetMessageWithNearToken(string message, string? nearToken)
 
 public class LuaUnDumpException(string message) : Exception(message);
 
-public class LuaRuntimeException : Exception
+internal interface ILuaTracebackBuildable
 {
-    public LuaRuntimeException(LuaThread? thread, Exception innerException) : base(innerException.Message,innerException)
+    Traceback? BuildOrGet();
+}
+
+public class LuaRuntimeException : Exception, ILuaTracebackBuildable
+{
+    public LuaRuntimeException(LuaThread? thread, Exception innerException) : base(innerException.Message, innerException)
     {
         Thread = thread;
     }
@@ -78,7 +84,7 @@ public LuaRuntimeException(LuaThread? thread, LuaValue errorObject)
     {
         if (thread != null)
         {
-            thread.CurrentException?.Build();
+            thread.CurrentException?.BuildOrGet();
             thread.ExceptionTrace.Clear();
             thread.CurrentException = this;
         }
@@ -96,7 +102,7 @@ public Traceback? LuaTraceback
         {
             if (luaTraceback == null)
             {
-                Build();
+                ((ILuaTracebackBuildable)this).BuildOrGet();
             }
 
             return luaTraceback;
@@ -185,7 +191,7 @@ static string CreateMessage(Traceback traceback, LuaValue errorObject)
 
 
     [MethodImpl(MethodImplOptions.NoInlining)]
-    internal Traceback? Build()
+    Traceback? ILuaTracebackBuildable.BuildOrGet()
     {
         if (luaTraceback != null) return luaTraceback;
         if (Thread != null)
@@ -247,4 +253,50 @@ public override string ToString()
 
 public class LuaAssertionException(LuaThread? traceback, string message) : LuaRuntimeException(traceback, message);
 
-public class LuaModuleNotFoundException(string moduleName) : Exception($"module '{moduleName}' not found");
+public class LuaModuleNotFoundException(string moduleName) : Exception($"module '{moduleName}' not found");
+
+public sealed class LuaCancelledException : OperationCanceledException, ILuaTracebackBuildable
+{
+    Traceback? luaTraceback;
+
+    public Traceback? LuaTraceback
+    {
+        get
+        {
+            if (luaTraceback == null)
+            {
+                ((ILuaTracebackBuildable)this).BuildOrGet();
+            }
+
+            return luaTraceback;
+        }
+    }
+
+    internal LuaThread? Thread { get; private set; }
+
+    internal LuaCancelledException(LuaThread thread, CancellationToken cancellationToken, Exception? innerException = null) : base("The operation was cancelled during execution on Lua.", innerException, cancellationToken)
+    {
+        thread.CurrentException?.BuildOrGet();
+        thread.ExceptionTrace.Clear();
+        thread.CurrentException = this;
+        Thread = thread;
+    }
+
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    Traceback? ILuaTracebackBuildable.BuildOrGet()
+    {
+        if (luaTraceback != null) return luaTraceback;
+
+        if (Thread != null)
+        {
+            var callStack = Thread.ExceptionTrace.AsSpan();
+            if (callStack.IsEmpty) return null;
+            luaTraceback = new Traceback(Thread.State, callStack);
+            Thread.ExceptionTrace.Clear();
+            Thread = null!;
+        }
+
+        return luaTraceback;
+    }
+}

src/Lua/LuaCoroutine.cs

@@ -186,9 +186,9 @@ async ValueTask<int> ResumeAsyncCore(LuaStack stack, int argCount, int returnBas
             {
                 if (IsProtectedMode)
                 {
-                    if (ex is LuaRuntimeException luaRuntimeException)
+                    if (ex is ILuaTracebackBuildable tracebackBuildable)
                     {
-                        traceback = luaRuntimeException.Build();
+                        traceback = tracebackBuildable.BuildOrGet();
                     }
 
                     Volatile.Write(ref status, (byte)LuaThreadStatus.Dead);

src/Lua/LuaThread.cs

@@ -70,8 +70,11 @@ public void Release()
     internal int LastVersion;
     internal int CurrentVersion;
 
-    internal LuaRuntimeException? CurrentException;
+    internal ILuaTracebackBuildable? CurrentException;
     internal readonly ReversedStack<CallStackFrame> ExceptionTrace = new();
+
+    // internal bool CancelRequested;
+    // internal CancellationToken CancellationToken;
 
     public bool IsRunning => CallStackFrameCount != 0;
     internal LuaFunction? Hook { get; set; }
@@ -131,7 +134,7 @@ void UpdateCurrentVersion(ref FastStackCore<CallStackFrame> callStack)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     internal LuaThreadAccess PushCallStackFrame(in CallStackFrame frame)
     {
-        CurrentException?.Build();
+        CurrentException?.BuildOrGet();
         CurrentException = null;
         ref var callStack = ref CoreData!.CallStack;
         callStack.Push(frame);

src/Lua/LuaThreadExtensions.cs

@@ -13,4 +13,19 @@ public static CoroutineLease RentCoroutine(this LuaThread thread, LuaFunction fu
     {
         return new(LuaCoroutine.Create(thread, function, isProtectedMode));
     }
+
+    internal static void ThrowIfCancellationRequested(this LuaThread thread, CancellationToken cancellationToken)
+    {
+        if (cancellationToken.IsCancellationRequested)
+        {
+            Throw(thread, cancellationToken);
+        }
+
+        return;
+
+        static void Throw(LuaThread thread, CancellationToken cancellationToken)
+        {
+            throw new LuaCancelledException(thread, cancellationToken);
+        }
+    }
 }

src/Lua/Runtime/LuaStack.cs

@@ -32,7 +32,7 @@ static void Resize(ref LuaValue[] array, int newSize)
 
             if (1000000 < size)
             {
-                throw new LuaException("Lua Stack overflow");
+                throw new ("Lua Stack overflow");
             }
 
             Array.Resize(ref array, size);

src/Lua/Runtime/LuaThreadAccess.cs

@@ -53,6 +53,7 @@ public async ValueTask<int> RunAsync(LuaFunction function, int argumentCount, in
             throw new ArgumentNullException(nameof(function));
         }
 
+        Thread.ThrowIfCancellationRequested(cancellationToken);
         var thread = Thread;
         var varArgumentCount = function.GetVariableArgumentCount(argumentCount);
         if (varArgumentCount != 0)
@@ -78,7 +79,7 @@ public async ValueTask<int> RunAsync(LuaFunction function, int argumentCount, in
 
         var access = thread.PushCallStackFrame(frame);
         LuaFunctionExecutionContext context = new() { Access = access, ArgumentCount = argumentCount, ReturnFrameBase = returnBase, };
-
+        var callStackTop = thread.CallStackFrameCount;
         try
         {
             if (this.Thread.CallOrReturnHookMask.Value != 0 && !this.Thread.IsInHook)
@@ -90,7 +91,7 @@ public async ValueTask<int> RunAsync(LuaFunction function, int argumentCount, in
         }
         finally
         {
-            this.Thread.PopCallStackFrame();
+            this.Thread.PopCallStackFrameUntil(callStackTop-1);
         }
     }
 

src/Lua/Runtime/LuaVirtualMachine.Debug.cs

@@ -47,7 +47,7 @@ static async ValueTask<int> Impl(VirtualMachineExecutionContext context)
                 countHookIsDone = true;
             }
 
-
+            context.ThrowIfCancellationRequested();
             if (context.Thread.IsLineHookEnabled)
             {
                 var sourcePositions = prototype.LineInfo;
@@ -126,6 +126,7 @@ internal static async ValueTask<int> ExecuteCallHook(LuaFunctionExecutionContext
                 context.Thread.PopCallStackFrameWithStackPop();
             }
         }
+        context.Thread.ThrowIfCancellationRequested(cancellationToken);
 
         {
             var frame = context.Thread.GetCurrentFrame();
@@ -135,7 +136,7 @@ internal static async ValueTask<int> ExecuteCallHook(LuaFunctionExecutionContext
             {
                 return r;
             }
-
+            context.Thread.ThrowIfCancellationRequested(cancellationToken);
             var top = stack.Count;
             stack.Push("return");
             stack.Push(LuaValue.Nil);