Merge pull request #53 from Akeit0/fix-closure-jmp

nuskey8 · web-flow · commit 3fa08d6654e4 · 2024-12-29T18:05:14.000+09:00
Fix: closure upvalues
diff --git a/src/Lua/CodeAnalysis/Compilation/LuaCompiler.cs b/src/Lua/CodeAnalysis/Compilation/LuaCompiler.cs
@@ -255,23 +255,26 @@ public bool VisitTableConstructorExpressionNode(TableConstructorExpressionNode n
                         // For the last element, we need to take into account variable arguments and multiple return values.
                         if (listItem == lastField)
                         {
+                            bool isFixedItems = true;
                             switch (listItem.Expression)
                             {
                                 case CallFunctionExpressionNode call:
                                     CompileCallFunctionExpression(call, context, false, -1);
+                                    isFixedItems = false;
                                     break;
                                 case CallTableMethodExpressionNode method:
                                     CompileTableMethod(method, context, false, -1);
                                     break;
                                 case VariableArgumentsExpressionNode varArg:
                                     CompileVariableArgumentsExpression(varArg, context, -1);
+                                    isFixedItems = false;
                                     break;
                                 default:
                                     listItem.Expression.Accept(this, context);
                                     break;
                             }
 
-                            context.PushInstruction(Instruction.SetList(tableRegisterIndex, 0, arrayBlock), listItem.Position);
+                            context.PushInstruction(Instruction.SetList(tableRegisterIndex, (ushort)(isFixedItems ? context.StackTopPosition - tableRegisterIndex : 0), arrayBlock), listItem.Position);
                             currentArrayChunkSize = 0;
                         }
                         else
@@ -635,7 +638,7 @@ public bool VisitTableMethodDeclarationStatementNode(TableMethodDeclarationState
         // assign global variable
         var first = node.MemberPath[0];
         var tableIndex = GetOrLoadIdentifier(first.Name, context, first.Position, true);
-        
+
         for (int i = 1; i < node.MemberPath.Length - 1; i++)
         {
             var member = node.MemberPath[i];
@@ -736,7 +739,7 @@ public bool VisitIfStatementNode(IfStatementNode node, ScopeCompilationContext c
     {
         using var endJumpIndexList = new PooledList<int>(8);
         var hasElse = node.ElseNodes.Length > 0;
-
+        var stackPositionToClose = (byte)(context.StackPosition + 1);
         // if
         using (var scopeContext = context.CreateChildScope())
         {
@@ -750,15 +753,15 @@ public bool VisitIfStatementNode(IfStatementNode node, ScopeCompilationContext c
                 childNode.Accept(this, scopeContext);
             }
 
+            stackPositionToClose = scopeContext.HasCapturedLocalVariables ? stackPositionToClose : (byte)0;
             if (hasElse)
             {
                 endJumpIndexList.Add(scopeContext.Function.Instructions.Length);
-                var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;
-                scopeContext.PushInstruction(Instruction.Jmp(a, 0), node.Position, true);
+                scopeContext.PushInstruction(Instruction.Jmp(stackPositionToClose, 0), node.Position, true);
             }
             else
             {
-                scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);
+                scopeContext.TryPushCloseUpValue(stackPositionToClose, node.Position);
             }
 
             scopeContext.Function.Instructions[ifPosition].SBx = scopeContext.Function.Instructions.Length - 1 - ifPosition;
@@ -779,16 +782,16 @@ public bool VisitIfStatementNode(IfStatementNode node, ScopeCompilationContext c
                 childNode.Accept(this, scopeContext);
             }
 
+            stackPositionToClose = scopeContext.HasCapturedLocalVariables ? stackPositionToClose : (byte)0;
             // skip if node doesn't have else statements
             if (hasElse)
             {
                 endJumpIndexList.Add(scopeContext.Function.Instructions.Length);
-                var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;
-                scopeContext.PushInstruction(Instruction.Jmp(a, 0), node.Position);
+                scopeContext.PushInstruction(Instruction.Jmp(stackPositionToClose, 0), node.Position);
             }
             else
             {
-                scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);
+                scopeContext.TryPushCloseUpValue(stackPositionToClose, node.Position);
             }
 
             scopeContext.Function.Instructions[elseifPosition].SBx = scopeContext.Function.Instructions.Length - 1 - elseifPosition;
@@ -821,14 +824,14 @@ public bool VisitRepeatStatementNode(RepeatStatementNode node, ScopeCompilationC
         context.Function.LoopLevel++;
 
         using var scopeContext = context.CreateChildScope();
-
+        var stackPosition = scopeContext.StackPosition;
         foreach (var childNode in node.Nodes)
         {
             childNode.Accept(this, scopeContext);
         }
 
         CompileConditionNode(node.ConditionNode, scopeContext, true);
-        var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;
+        var a = scopeContext.HasCapturedLocalVariables ? (byte)(stackPosition + 1) : (byte)0;
         scopeContext.PushInstruction(Instruction.Jmp(a, startIndex - scopeContext.Function.Instructions.Length - 1), node.Position);
         scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);
 
@@ -848,6 +851,7 @@ public bool VisitWhileStatementNode(WhileStatementNode node, ScopeCompilationCon
         context.Function.LoopLevel++;
 
         using var scopeContext = context.CreateChildScope();
+        var stackPosition = scopeContext.StackPosition;
 
         foreach (var childNode in node.Nodes)
         {
@@ -860,7 +864,7 @@ public bool VisitWhileStatementNode(WhileStatementNode node, ScopeCompilationCon
         scopeContext.Function.Instructions[conditionIndex].SBx = scopeContext.Function.Instructions.Length - 1 - conditionIndex;
 
         CompileConditionNode(node.ConditionNode, scopeContext, false);
-        var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;
+        var a = scopeContext.HasCapturedLocalVariables ? (byte)(1 + stackPosition) : (byte)0;
         scopeContext.PushInstruction(Instruction.Jmp(a, conditionIndex - context.Function.Instructions.Length), node.Position);
         scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);
 
diff --git a/src/Lua/Runtime/Closure.cs b/src/Lua/Runtime/Closure.cs
@@ -1,3 +1,4 @@
+using System.Runtime.CompilerServices;
 using Lua.Internal;
 
 namespace Lua.Runtime;
@@ -16,30 +17,41 @@ public Closure(LuaState state, Chunk proto, LuaTable? environment = null)
         for (int i = 0; i < proto.UpValues.Length; i++)
         {
             var description = proto.UpValues[i];
-            var upValue = GetUpValueFromDescription(state, environment == null ? state.EnvUpValue : UpValue.Closed(environment), proto, description, 1);
+            var upValue = GetUpValueFromDescription(state, state.CurrentThread, environment == null ? state.EnvUpValue : UpValue.Closed(environment), description);
             upValues.Add(upValue);
         }
     }
 
     public Chunk Proto => proto;
     public ReadOnlySpan<UpValue> UpValues => upValues.AsSpan();
 
-    static UpValue GetUpValueFromDescription(LuaState state, UpValue envUpValue, Chunk proto, UpValueInfo description, int depth)
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    internal LuaValue GetUpValue(int index)
+    {
+        return upValues[index].GetValue();
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    internal void SetUpValue(int index, LuaValue value)
+    {
+        upValues[index].SetValue(value);
+    }
+
+    static UpValue GetUpValueFromDescription(LuaState state, LuaThread thread, UpValue envUpValue, UpValueInfo description)
     {
         if (description.IsInRegister)
         {
-            var thread = state.CurrentThread;
-            var callStack = thread.GetCallStackFrames();
-            var frame = callStack[^depth];
-            return state.GetOrAddUpValue(thread, frame.Base + description.Index);
+            return state.GetOrAddUpValue(thread, thread.GetCallStackFrames()[^1].Base + description.Index);
         }
-        else if (description.Index == -1) // -1 is global environment
+        if (description.Index == -1) // -1 is global environment
         {
             return envUpValue;
         }
-        else
+        if (thread.GetCallStackFrames()[^1].Function is Closure parentClosure)
         {
-            return GetUpValueFromDescription(state, envUpValue, proto.Parent!, proto.Parent!.UpValues[description.Index], depth + 1);
+            return parentClosure.UpValues[description.Index];
         }
+
+        throw new Exception();
     }
 }
diff --git a/src/Lua/Runtime/LuaVirtualMachine.cs b/src/Lua/Runtime/LuaVirtualMachine.cs
@@ -532,7 +532,7 @@ internal async static ValueTask<int> ExecuteClosureAsync(LuaState state, CallSta
                         pc += instruction.SBx;
                         if (instruction.A != 0)
                         {
-                            state.CloseUpValues(thread, instruction.A - 1);
+                            state.CloseUpValues(thread, frame.Base + instruction.A - 1);
                         }
                         break;
                     case OpCode.Eq:

Original file line number	Diff line number	Diff line change
`@@ -255,23 +255,26 @@ public bool VisitTableConstructorExpressionNode(TableConstructorExpressionNode n`
`255`	`255`	`// For the last element, we need to take into account variable arguments and multiple return values.`
`256`	`256`	`if (listItem == lastField)`
`257`	`257`	`{`
	`258`	`+ bool isFixedItems = true;`
`258`	`259`	`switch (listItem.Expression)`
`259`	`260`	`{`
`260`	`261`	`case CallFunctionExpressionNode call:`
`261`	`262`	`CompileCallFunctionExpression(call, context, false, -1);`
	`263`	`+ isFixedItems = false;`
`262`	`264`	`break;`
`263`	`265`	`case CallTableMethodExpressionNode method:`
`264`	`266`	`CompileTableMethod(method, context, false, -1);`
`265`	`267`	`break;`
`266`	`268`	`case VariableArgumentsExpressionNode varArg:`
`267`	`269`	`CompileVariableArgumentsExpression(varArg, context, -1);`
	`270`	`+ isFixedItems = false;`
`268`	`271`	`break;`
`269`	`272`	`default:`
`270`	`273`	`listItem.Expression.Accept(this, context);`
`271`	`274`	`break;`
`272`	`275`	`}`
`273`	`276`
`274`		`- context.PushInstruction(Instruction.SetList(tableRegisterIndex, 0, arrayBlock), listItem.Position);`
	`277`	`+ context.PushInstruction(Instruction.SetList(tableRegisterIndex, (ushort)(isFixedItems ? context.StackTopPosition - tableRegisterIndex : 0), arrayBlock), listItem.Position);`
`275`	`278`	`currentArrayChunkSize = 0;`
`276`	`279`	`}`
`277`	`280`	`else`
`@@ -635,7 +638,7 @@ public bool VisitTableMethodDeclarationStatementNode(TableMethodDeclarationState`
`635`	`638`	`// assign global variable`
`636`	`639`	`var first = node.MemberPath[0];`
`637`	`640`	`var tableIndex = GetOrLoadIdentifier(first.Name, context, first.Position, true);`
`638`		`-`
	`641`	`+`
`639`	`642`	`for (int i = 1; i < node.MemberPath.Length - 1; i++)`
`640`	`643`	`{`
`641`	`644`	`var member = node.MemberPath[i];`
`@@ -736,7 +739,7 @@ public bool VisitIfStatementNode(IfStatementNode node, ScopeCompilationContext c`
`736`	`739`	`{`
`737`	`740`	`using var endJumpIndexList = new PooledList<int>(8);`
`738`	`741`	`var hasElse = node.ElseNodes.Length > 0;`
`739`		`-`
	`742`	`+ var stackPositionToClose = (byte)(context.StackPosition + 1);`
`740`	`743`	`// if`
`741`	`744`	`using (var scopeContext = context.CreateChildScope())`
`742`	`745`	`{`
`@@ -750,15 +753,15 @@ public bool VisitIfStatementNode(IfStatementNode node, ScopeCompilationContext c`
`750`	`753`	`childNode.Accept(this, scopeContext);`
`751`	`754`	`}`
`752`	`755`
	`756`	`+ stackPositionToClose = scopeContext.HasCapturedLocalVariables ? stackPositionToClose : (byte)0;`
`753`	`757`	`if (hasElse)`
`754`	`758`	`{`
`755`	`759`	`endJumpIndexList.Add(scopeContext.Function.Instructions.Length);`
`756`		`- var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;`
`757`		`- scopeContext.PushInstruction(Instruction.Jmp(a, 0), node.Position, true);`
	`760`	`+ scopeContext.PushInstruction(Instruction.Jmp(stackPositionToClose, 0), node.Position, true);`
`758`	`761`	`}`
`759`	`762`	`else`
`760`	`763`	`{`
`761`		`- scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);`
	`764`	`+ scopeContext.TryPushCloseUpValue(stackPositionToClose, node.Position);`
`762`	`765`	`}`
`763`	`766`
`764`	`767`	`scopeContext.Function.Instructions[ifPosition].SBx = scopeContext.Function.Instructions.Length - 1 - ifPosition;`
`@@ -779,16 +782,16 @@ public bool VisitIfStatementNode(IfStatementNode node, ScopeCompilationContext c`
`779`	`782`	`childNode.Accept(this, scopeContext);`
`780`	`783`	`}`
`781`	`784`
	`785`	`+ stackPositionToClose = scopeContext.HasCapturedLocalVariables ? stackPositionToClose : (byte)0;`
`782`	`786`	`// skip if node doesn't have else statements`
`783`	`787`	`if (hasElse)`
`784`	`788`	`{`
`785`	`789`	`endJumpIndexList.Add(scopeContext.Function.Instructions.Length);`
`786`		`- var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;`
`787`		`- scopeContext.PushInstruction(Instruction.Jmp(a, 0), node.Position);`
	`790`	`+ scopeContext.PushInstruction(Instruction.Jmp(stackPositionToClose, 0), node.Position);`
`788`	`791`	`}`
`789`	`792`	`else`
`790`	`793`	`{`
`791`		`- scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);`
	`794`	`+ scopeContext.TryPushCloseUpValue(stackPositionToClose, node.Position);`
`792`	`795`	`}`
`793`	`796`
`794`	`797`	`scopeContext.Function.Instructions[elseifPosition].SBx = scopeContext.Function.Instructions.Length - 1 - elseifPosition;`
`@@ -821,14 +824,14 @@ public bool VisitRepeatStatementNode(RepeatStatementNode node, ScopeCompilationC`
`821`	`824`	`context.Function.LoopLevel++;`
`822`	`825`
`823`	`826`	`using var scopeContext = context.CreateChildScope();`
`824`		`-`
	`827`	`+ var stackPosition = scopeContext.StackPosition;`
`825`	`828`	`foreach (var childNode in node.Nodes)`
`826`	`829`	`{`
`827`	`830`	`childNode.Accept(this, scopeContext);`
`828`	`831`	`}`
`829`	`832`
`830`	`833`	`CompileConditionNode(node.ConditionNode, scopeContext, true);`
`831`		`- var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;`
	`834`	`+ var a = scopeContext.HasCapturedLocalVariables ? (byte)(stackPosition + 1) : (byte)0;`
`832`	`835`	`scopeContext.PushInstruction(Instruction.Jmp(a, startIndex - scopeContext.Function.Instructions.Length - 1), node.Position);`
`833`	`836`	`scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);`
`834`	`837`
`@@ -848,6 +851,7 @@ public bool VisitWhileStatementNode(WhileStatementNode node, ScopeCompilationCon`
`848`	`851`	`context.Function.LoopLevel++;`
`849`	`852`
`850`	`853`	`using var scopeContext = context.CreateChildScope();`
	`854`	`+ var stackPosition = scopeContext.StackPosition;`
`851`	`855`
`852`	`856`	`foreach (var childNode in node.Nodes)`
`853`	`857`	`{`
`@@ -860,7 +864,7 @@ public bool VisitWhileStatementNode(WhileStatementNode node, ScopeCompilationCon`
`860`	`864`	`scopeContext.Function.Instructions[conditionIndex].SBx = scopeContext.Function.Instructions.Length - 1 - conditionIndex;`
`861`	`865`
`862`	`866`	`CompileConditionNode(node.ConditionNode, scopeContext, false);`
`863`		`- var a = scopeContext.HasCapturedLocalVariables ? scopeContext.StackPosition : (byte)0;`
	`867`	`+ var a = scopeContext.HasCapturedLocalVariables ? (byte)(1 + stackPosition) : (byte)0;`
`864`	`868`	`scopeContext.PushInstruction(Instruction.Jmp(a, conditionIndex - context.Function.Instructions.Length), node.Position);`
`865`	`869`	`scopeContext.TryPushCloseUpValue(scopeContext.StackPosition, node.Position);`
`866`	`870`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+using System.Runtime.CompilerServices;`
`1`	`2`	`using Lua.Internal;`
`2`	`3`
`3`	`4`	`namespace Lua.Runtime;`
`@@ -16,30 +17,41 @@ public Closure(LuaState state, Chunk proto, LuaTable? environment = null)`
`16`	`17`	`for (int i = 0; i < proto.UpValues.Length; i++)`
`17`	`18`	`{`
`18`	`19`	`var description = proto.UpValues[i];`
`19`		`- var upValue = GetUpValueFromDescription(state, environment == null ? state.EnvUpValue : UpValue.Closed(environment), proto, description, 1);`
	`20`	`+ var upValue = GetUpValueFromDescription(state, state.CurrentThread, environment == null ? state.EnvUpValue : UpValue.Closed(environment), description);`
`20`	`21`	`upValues.Add(upValue);`
`21`	`22`	`}`
`22`	`23`	`}`
`23`	`24`
`24`	`25`	`public Chunk Proto => proto;`
`25`	`26`	`public ReadOnlySpan<UpValue> UpValues => upValues.AsSpan();`
`26`	`27`
`27`		`- static UpValue GetUpValueFromDescription(LuaState state, UpValue envUpValue, Chunk proto, UpValueInfo description, int depth)`
	`28`	`+ [MethodImpl(MethodImplOptions.AggressiveInlining)]`
	`29`	`+ internal LuaValue GetUpValue(int index)`
	`30`	`+ {`
	`31`	`+ return upValues[index].GetValue();`
	`32`	`+ }`
	`33`	`+`
	`34`	`+ [MethodImpl(MethodImplOptions.AggressiveInlining)]`
	`35`	`+ internal void SetUpValue(int index, LuaValue value)`
	`36`	`+ {`
	`37`	`+ upValues[index].SetValue(value);`
	`38`	`+ }`
	`39`	`+`
	`40`	`+ static UpValue GetUpValueFromDescription(LuaState state, LuaThread thread, UpValue envUpValue, UpValueInfo description)`
`28`	`41`	`{`
`29`	`42`	`if (description.IsInRegister)`
`30`	`43`	`{`
`31`		`- var thread = state.CurrentThread;`
`32`		`- var callStack = thread.GetCallStackFrames();`
`33`		`- var frame = callStack[^depth];`
`34`		`- return state.GetOrAddUpValue(thread, frame.Base + description.Index);`
	`44`	`+ return state.GetOrAddUpValue(thread, thread.GetCallStackFrames()[^1].Base + description.Index);`
`35`	`45`	`}`
`36`		`- else if (description.Index == -1) // -1 is global environment`
	`46`	`+ if (description.Index == -1) // -1 is global environment`
`37`	`47`	`{`
`38`	`48`	`return envUpValue;`
`39`	`49`	`}`
`40`		`- else`
	`50`	`+ if (thread.GetCallStackFrames()[^1].Function is Closure parentClosure)`
`41`	`51`	`{`
`42`		`- return GetUpValueFromDescription(state, envUpValue, proto.Parent!, proto.Parent!.UpValues[description.Index], depth + 1);`
	`52`	`+ return parentClosure.UpValues[description.Index];`
`43`	`53`	`}`
	`54`	`+`
	`55`	`+ throw new Exception();`
`44`	`56`	`}`
`45`	`57`	`}`
Original file line number	Diff line number	Diff line change
`@@ -532,7 +532,7 @@ internal async static ValueTask<int> ExecuteClosureAsync(LuaState state, CallSta`
`532`	`532`	`pc += instruction.SBx;`
`533`	`533`	`if (instruction.A != 0)`
`534`	`534`	`{`
`535`		`- state.CloseUpValues(thread, instruction.A - 1);`
	`535`	`+ state.CloseUpValues(thread, frame.Base + instruction.A - 1);`
`536`	`536`	`}`
`537`	`537`	`break;`
`538`	`538`	`case OpCode.Eq:`