Merge branch 'main' into optimize

src/Lua/CodeAnalysis/Compilation/FunctionCompilationContext.cs

@@ -115,6 +115,150 @@ public class FunctionCompilationContext : IDisposable
         instructionPositions.Add(position);
     }
 
+    /// <summary>
+    /// Push or merge the new instruction.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void PushOrMergeInstruction(int lastLocal, in Instruction instruction, in SourcePosition position, ref bool incrementStackPosition)
+    {
+        if (instructions.Length == 0)
+        {
+            instructions.Add(instruction);
+            instructionPositions.Add(position);
+            return;
+        }
+        ref var lastInstruction = ref instructions.AsSpan()[^1];
+        var opcode = instruction.OpCode;
+        switch (opcode)
+        {
+            case OpCode.Move:
+                // last A is not local variable
+                if (lastInstruction.A != lastLocal &&
+                    // available to merge
+                    lastInstruction.A == instruction.B &&
+                    // not already merged
+                    lastInstruction.A != lastInstruction.B)
+                {
+                    switch (lastInstruction.OpCode)
+                    {
+                        case OpCode.GetTable:
+                        case OpCode.Add:
+                        case OpCode.Sub:
+                        case OpCode.Mul:
+                        case OpCode.Div:
+                        case OpCode.Mod:
+                        case OpCode.Pow:
+                        case OpCode.Concat:
+                            {
+                                lastInstruction.A = instruction.A;
+                                incrementStackPosition = false;
+                                return;
+                            }
+                    }
+                }
+                break;
+            case OpCode.GetTable:
+                {
+                    // Merge MOVE GetTable
+                    if (lastInstruction.OpCode == OpCode.Move && lastLocal != lastInstruction.A)
+                    {
+                        if (lastInstruction.A == instruction.B)
+                        {
+                            lastInstruction = Instruction.GetTable(instruction.A, lastInstruction.B, instruction.C);
+                            instructionPositions[^1] = position;
+                            incrementStackPosition = false;
+                            return;
+                        }
+
+                    }
+                    break;
+                }
+            case OpCode.SetTable:
+                {
+                    // Merge MOVE SETTABLE
+                    if (lastInstruction.OpCode == OpCode.Move && lastLocal != lastInstruction.A)
+                    {
+                        var lastB = lastInstruction.B;
+                        var lastA = lastInstruction.A;
+                        if (lastB < 255 && lastA == instruction.A)
+                        {
+                            // Merge MOVE MOVE SETTABLE
+                            if (instructions.Length > 2)
+                            {
+                                ref var last2Instruction = ref instructions.AsSpan()[^2];
+                                var last2A = last2Instruction.A;
+                                if (last2Instruction.OpCode == OpCode.Move && lastLocal != last2A && instruction.C == last2A)
+                                {
+                                    last2Instruction = Instruction.SetTable((byte)(lastB), instruction.B, last2Instruction.B);
+                                    instructions.RemoveAtSwapback(instructions.Length - 1);
+                                    instructionPositions.RemoveAtSwapback(instructionPositions.Length - 1);
+                                    instructionPositions[^1] = position;
+                                    incrementStackPosition = false;
+                                    return;
+                                }
+                            }
+                            lastInstruction = Instruction.SetTable((byte)(lastB), instruction.B, instruction.C);
+                            instructionPositions[^1] = position;
+                            incrementStackPosition = false;
+                            return;
+                        }
+
+                        if (lastA == instruction.C)
+                        {
+                            lastInstruction = Instruction.SetTable(instruction.A, instruction.B, lastB);
+                            instructionPositions[^1] = position;
+                            incrementStackPosition = false;
+                            return;
+                        }
+                    }
+                    else if (lastInstruction.OpCode == OpCode.GetTabUp && instructions.Length >= 2)
+                    {
+                        ref var last2Instruction = ref instructions[^2];
+                        var last2OpCode = last2Instruction.OpCode;
+                        if (last2OpCode is OpCode.LoadK or OpCode.Move)
+                        {
+
+                            var last2A = last2Instruction.A;
+                            if (last2A != lastLocal && instruction.C == last2A)
+                            {
+                                var c = last2OpCode == OpCode.LoadK ? last2Instruction.Bx + 256 : last2Instruction.B;
+                                last2Instruction = lastInstruction;
+                                lastInstruction = instruction with { C = (ushort)c };
+                                instructionPositions[^2] = instructionPositions[^1];
+                                instructionPositions[^1] = position;
+                                incrementStackPosition = false;
+                                return;
+                            }
+                        }
+                    }
+                    break;
+                }
+            case OpCode.Unm:
+            case OpCode.Not:
+            case OpCode.Len:
+                if (lastInstruction.OpCode == OpCode.Move && lastLocal != lastInstruction.A && lastInstruction.A == instruction.B)
+                {
+                    lastInstruction = instruction with { B = lastInstruction.B }; ;
+                    instructionPositions[^1] = position;
+                    incrementStackPosition = false;
+                    return;
+                }
+                break;
+            case OpCode.Return:
+                if (lastInstruction.OpCode == OpCode.Move && instruction.B == 2 && lastInstruction.B < 256)
+                {
+                    lastInstruction = instruction with { A = (byte)lastInstruction.B };
+                    instructionPositions[^1] = position;
+                    incrementStackPosition = false;
+                    return;
+                }
+                break;
+        }
+
+        instructions.Add(instruction);
+        instructionPositions.Add(position);
+    }
+
     /// <summary>
     /// Gets the index of the constant from the value, or if the constant is not registered it is added and its index is returned.
     /// </summary>

src/Lua/CodeAnalysis/Compilation/LuaCompiler.cs

@@ -255,23 +255,26 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
                         // 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
@@ -744,7 +747,7 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
     {
         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())
         {
@@ -758,15 +761,15 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
                 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;
@@ -787,16 +790,16 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
                 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;
@@ -829,14 +832,14 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
         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);
 
@@ -856,6 +859,7 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
         context.Function.LoopLevel++;
 
         using var scopeContext = context.CreateChildScope();
+        var stackPosition = scopeContext.StackPosition;
 
         foreach (var childNode in node.Nodes)
         {
@@ -868,7 +872,7 @@ public sealed class LuaCompiler : ISyntaxNodeVisitor<ScopeCompilationContext, bo
         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);
 

src/Lua/CodeAnalysis/Compilation/ScopeCompilationContext.cs

@@ -33,6 +33,8 @@ public class ScopeCompilationContext : IDisposable
     readonly Dictionary<ReadOnlyMemory<char>, LocalVariableDescription> localVariables = new(256, Utf16StringMemoryComparer.Default);
     readonly Dictionary<ReadOnlyMemory<char>, LabelDescription> labels = new(32, Utf16StringMemoryComparer.Default);
 
+    byte lastLocalVariableIndex;
+
     public byte StackStartPosition { get; private set; }
     public byte StackPosition { get; set; }
 
@@ -72,8 +74,11 @@ public class ScopeCompilationContext : IDisposable
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public void PushInstruction(in Instruction instruction, SourcePosition position, bool incrementStackPosition = false)
     {
-        Function.PushInstruction(instruction, position);
-        if (incrementStackPosition) StackPosition++;
+        Function.PushOrMergeInstruction(lastLocalVariableIndex, instruction, position, ref incrementStackPosition);
+        if (incrementStackPosition)
+        {
+            StackPosition++;
+        }
         Function.MaxStackPosition = Math.Max(Function.MaxStackPosition, StackPosition);
     }
 
@@ -90,11 +95,13 @@ public class ScopeCompilationContext : IDisposable
     /// Add new local variable.
     /// </summary>
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public void AddLocalVariable(ReadOnlyMemory<char> name, LocalVariableDescription description)
+    public void AddLocalVariable(ReadOnlyMemory<char> name, LocalVariableDescription description, bool markAsLastLocalVariable = true)
     {
         localVariables[name] = description;
+        lastLocalVariableIndex = description.RegisterIndex;
     }
 
+
     /// <summary>
     /// Gets the local variable in scope.
     /// </summary>
@@ -158,6 +165,7 @@ public class ScopeCompilationContext : IDisposable
         HasCapturedLocalVariables = false;
         localVariables.Clear();
         labels.Clear();
+        lastLocalVariableIndex = 0;
     }
 
     /// <summary>

src/Lua/Runtime/Closure.cs

@@ -45,19 +45,17 @@ public sealed class Closure : LuaFunction
         {
             return state.GetOrAddUpValue(thread, thread.GetCallStackFrames()[^1].Base + description.Index);
         }
-
+        
         if (description.Index == -1) // -1 is global environment
         {
             return envUpValue;
         }
-
+        
+        if (thread.GetCallStackFrames()[^1].Function is Closure parentClosure)
         {
-            if (thread.GetCallStackFrames()[^1].Function is Closure parentClosure)
-            {
-                return parentClosure.UpValues[description.Index];
-            }
-
-            throw new InvalidOperationException("invalid upvalue description.");
+             return parentClosure.UpValues[description.Index];
         }
+        
+        throw new Exception();
     }
 }