CFG Part 0: Max Stack Size & Peephole Optimizer Fixes/Improvements

DMCompiler/Bytecode/DreamProcOpcode.cs

@@ -136,7 +136,7 @@ public enum DreamProcOpcode : byte {
     EnumerateAssoc = 0x43,
     [OpcodeMetadata(-2)]
     Link = 0x44,
-    [OpcodeMetadata(-3, OpcodeArgType.TypeId)]
+    [OpcodeMetadata(-4, OpcodeArgType.TypeId)]
     Prompt = 0x45,
     [OpcodeMetadata(-3)]
     Ftp = 0x46,
@@ -301,7 +301,7 @@ public enum DreamProcOpcode : byte {
     ReturnFloat = 0x98,
     [OpcodeMetadata(1, OpcodeArgType.Reference, OpcodeArgType.String)]
     IndexRefWithString = 0x99,
-    [OpcodeMetadata(2, OpcodeArgType.Float, OpcodeArgType.Reference)]
+    [OpcodeMetadata(0, OpcodeArgType.Float, OpcodeArgType.Reference)]
     PushFloatAssign = 0x9A,
     [OpcodeMetadata(true, 0, OpcodeArgType.Int)]
     NPushFloatAssign = 0x9B,

DMCompiler/DM/Builders/DMProcBuilder.cs

@@ -153,8 +153,10 @@ private void ProcessStatementSpawn(DMASTProcStatementSpawn statementSpawn) {
             ProcessBlockInner(statementSpawn.Body);
 
             //Prevent the new thread from executing outside its own code
-            proc.PushNull();
-            proc.Return();
+            if (!proc.LastInstructionTransfersControl()) {
+                proc.PushNull();
+                proc.Return();
+            }
         }
         proc.EndScope();
 
@@ -242,7 +244,8 @@ private void ProcessStatementIf(DMASTProcStatementIf statement) {
             proc.StartScope();
             ProcessBlockInner(statement.Body);
             proc.EndScope();
-            proc.Jump(endLabel);
+            if (!proc.LastInstructionTransfersControl())
+                proc.Jump(endLabel);
 
             proc.AddLabel(elseLabel);
             proc.StartScope();
@@ -758,7 +761,9 @@ Constant CoerceBound(Constant bound, bool upperRange) {
             proc.EndScope();
         }
 
-        proc.Jump(endLabel);
+        // don't Jump after a Return (or similar)
+        if (!proc.LastInstructionTransfersControl())
+            proc.Jump(endLabel);
 
         foreach ((string CaseLabel, DMASTProcBlockInner CaseBody) valueCase in valueCases) {
             proc.AddLabel(valueCase.CaseLabel);
@@ -767,7 +772,10 @@ Constant CoerceBound(Constant bound, bool upperRange) {
                 ProcessBlockInner(valueCase.CaseBody);
             }
             proc.EndScope();
-            proc.Jump(endLabel);
+
+            // don't Jump after a Return (or similar)
+            if (!proc.LastInstructionTransfersControl())
+                proc.Jump(endLabel);
         }
 
         proc.AddLabel(endLabel);

DMCompiler/DM/DMProc.cs

@@ -200,7 +200,9 @@ public void ValidateReturnType(DMExpression expr) {
     public ProcDefinitionJson GetJsonRepresentation() {
         var serializer = new AnnotatedBytecodeSerializer(_compiler);
 
-        _compiler.BytecodeOptimizer.Optimize(AnnotatedBytecode.GetAnnotatedBytecode());
+        List<IAnnotatedBytecode> annotatedBytecode = AnnotatedBytecode.GetAnnotatedBytecode();
+        _compiler.BytecodeOptimizer.Optimize(annotatedBytecode);
+        int maxStackSize = AnnotatedBytecode.RecalculateMaxStackSize();
 
         List<ProcArgumentJson>? arguments = null;
         if (_parameters.Count > 0) {
@@ -228,8 +230,8 @@ public ProcDefinitionJson GetJsonRepresentation() {
             OwningTypeId = _dmObject.Id,
             Name = Name,
             Attributes = Attributes,
-            MaxStackSize = AnnotatedBytecode.GetMaxStackSize(),
-            Bytecode = serializer.Serialize(AnnotatedBytecode.GetAnnotatedBytecode()),
+            MaxStackSize = maxStackSize,
+            Bytecode = serializer.Serialize(annotatedBytecode),
             Arguments = arguments,
             SourceInfo = serializer.SourceInfo,
             Locals = (_localVariableNames.Count > 0) ? serializer.GetLocalVariablesJson() : null,
@@ -848,6 +850,57 @@ public void Jump(string label) {
         WriteLabel(label);
     }
 
+    public bool LastInstructionTransfersControl() {
+        List<IAnnotatedBytecode> bytecode = AnnotatedBytecode.GetAnnotatedBytecode();
+        HashSet<string>? referencedLabels = null;
+        // Man sometimes it'd be nice if we had a list of just instructions and didn't need loops like this just to grab the last actual instruction
+        for (int i = bytecode.Count - 1; i >= 0; i--) {
+            switch (bytecode[i]) {
+                case AnnotatedBytecodeVariable:
+                    continue;
+                case AnnotatedBytecodeLabel label:
+                    referencedLabels ??= GetReferencedLabels(bytecode);
+                    if (referencedLabels is not null && referencedLabels.Contains(label.LabelName))
+                        return false;
+
+                    continue;
+                case AnnotatedBytecodeInstruction instruction:
+                    return InstructionTransfersControl(instruction.Opcode);
+                default:
+                    return false;
+            }
+        }
+
+        return false;
+    }
+
+    private HashSet<string>? GetReferencedLabels(List<IAnnotatedBytecode> bytecode) {
+        HashSet<string>? referencedLabels = null;
+
+        foreach (IAnnotatedBytecode item in bytecode) {
+            if (item is not AnnotatedBytecodeInstruction instruction)
+                continue;
+
+            foreach (IAnnotatedBytecode arg in instruction.GetArgs()) {
+                if (arg is AnnotatedBytecodeLabel label) {
+                    referencedLabels ??= new HashSet<string>(1);
+                    referencedLabels.Add(label.LabelName);
+                }
+            }
+        }
+
+        return referencedLabels;
+    }
+
+    // TODO: Once we have a CFG we'll likely be storing this info in opcode metadata and this method's hardcoded list can be removed
+    private bool InstructionTransfersControl(DreamProcOpcode opcode) {
+        return opcode is DreamProcOpcode.Jump or
+            DreamProcOpcode.Return or
+            DreamProcOpcode.ReturnReferenceValue or
+            DreamProcOpcode.ReturnFloat or
+            DreamProcOpcode.Throw;
+    }
+
     public void JumpIfFalse(string label) {
         WriteOpcode(DreamProcOpcode.JumpIfFalse);
         WriteLabel(label);
@@ -901,9 +954,8 @@ public void Call(DMReference reference, DMCallArgumentsType argumentsType, int a
         WriteStackDelta(argumentStackSize);
     }
 
-    public void CallStatement(DMCallArgumentsType argumentsType, int argumentStackSize) {
-        //Shrinks the stack by argumentStackSize. Could also shrink it by argumentStackSize+1, but assume not.
-        ResizeStack(-argumentStackSize);
+    public void CallStatement(DMCallArgumentsType argumentsType, int argumentStackSize, bool hasProcName) {
+        ResizeStack(-(argumentStackSize + (hasProcName ? 1 : 0)));
         WriteOpcode(DreamProcOpcode.CallStatement);
         WriteArgumentType(argumentsType);
         WriteStackDelta(argumentStackSize);
@@ -937,7 +989,7 @@ public void AssignInto(DMReference reference) {
     }
 
     public void CreateObject(DMCallArgumentsType argumentsType, int argumentStackSize) {
-        ResizeStack(-argumentStackSize); // Pops type and arguments, pushes new object
+        ResizeStack(-(argumentStackSize + 1)); // Pops overrides, type, and arguments, pushes new object
         WriteOpcode(DreamProcOpcode.CreateObject);
         WriteArgumentType(argumentsType);
         WriteStackDelta(argumentStackSize);
@@ -1267,7 +1319,7 @@ public void Animate(DMCallArgumentsType argumentsType, int argumentStackSize) {
     }
 
     public void PickWeighted(int count) {
-        ResizeStack(-(count - 1));
+        ResizeStack(-(count * 2 - 1));
         WriteOpcode(DreamProcOpcode.PickWeighted);
         WritePickCount(count);
     }

DMCompiler/DM/Expressions/Builtins.cs

@@ -674,7 +674,7 @@ public override void EmitPushValue(ExpressionContext ctx) {
 
         _b?.EmitPushValue(ctx);
         _a.EmitPushValue(ctx);
-        ctx.Proc.CallStatement(argumentInfo.Type, argumentInfo.StackSize);
+        ctx.Proc.CallStatement(argumentInfo.Type, argumentInfo.StackSize, _b is not null);
     }
 }
 

DMCompiler/Optimizer/AnnotatedByteCodeWriter.cs

@@ -18,6 +18,7 @@ private readonly List<IAnnotatedBytecode>
     private Location _location;
     private int _maxStackSize;
     private bool _negativeStackSizeError;
+    private int _pendingInstructionStackDelta;
     private int _requiredArgIdx;
     private OpcodeMetadata? _currentMetadata;
     private Dictionary<string, long> _labels = new();
@@ -49,9 +50,10 @@ public void WriteOpcode(DreamProcOpcode opcode, Location location) {
 
         // Goal here is to maintain correspondence between the raw bytecode and the annotated bytecode such that
         // the annotated bytecode can be used to generate the raw bytecode again.
-        _annotatedBytecode.Add(new AnnotatedBytecodeInstruction(opcode, metadata.StackDelta, location));
+        _annotatedBytecode.Add(new AnnotatedBytecodeInstruction(opcode, metadata.StackDelta + _pendingInstructionStackDelta, location));
+        _pendingInstructionStackDelta = 0;
 
-        ResizeStack(metadata.StackDelta);
+        ResizeStackOnly(metadata.StackDelta);
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -215,6 +217,19 @@ public void ResolveCodeLabelReferences(Stack<DMProc.CodeLabelReference> pendingL
     /// </summary>
     /// <param name="sizeDelta">The net change in stack size caused by an operation</param>
     public void ResizeStack(int sizeDelta) {
+        _pendingInstructionStackDelta += sizeDelta;
+        ResizeStackOnly(sizeDelta);
+    }
+
+    private void ResizeCurrentInstructionStack(int sizeDelta) {
+        if (_annotatedBytecode.Count > 0 && _annotatedBytecode[^1] is AnnotatedBytecodeInstruction instruction) {
+            instruction.StackSizeDelta += sizeDelta;
+        }
+
+        ResizeStackOnly(sizeDelta);
+    }
+
+    private void ResizeStackOnly(int sizeDelta) {
         _currentStackSize += sizeDelta;
         _maxStackSize = Math.Max(_currentStackSize, _maxStackSize);
         if (_currentStackSize < 0 && !_negativeStackSizeError) {
@@ -230,6 +245,27 @@ public int GetMaxStackSize() {
         return _maxStackSize;
     }
 
+    /// <summary>
+    /// Recomputes the maximum possible stack size from the current annotated bytecode.
+    /// Used after optimization because peephole rewrites can change the max stack size.
+    /// </summary>
+    public int RecalculateMaxStackSize() {
+        _currentStackSize = 0;
+        _maxStackSize = 0;
+        _negativeStackSizeError = false;
+        _pendingInstructionStackDelta = 0;
+
+        foreach (IAnnotatedBytecode bytecode in _annotatedBytecode) {
+            if (bytecode is not AnnotatedBytecodeInstruction instruction)
+                continue;
+
+            _location = instruction.GetLocation();
+            ResizeStackOnly(instruction.StackSizeDelta);
+        }
+
+        return _maxStackSize;
+    }
+
     public void WriteResource(string value, Location location) {
         _location = location;
         ValidateArgument(location, OpcodeArgType.Resource);
@@ -294,7 +330,7 @@ public void WriteReference(DMReference reference, Location location, bool affect
                 int fieldId = compiler.DMObjectTree.AddString(reference.Name);
                 _annotatedBytecode[^1]
                     .AddArg(compiler, new AnnotatedBytecodeReference(reference.RefType, fieldId, location));
-                ResizeStack(affectStack ? -1 : 0);
+                ResizeCurrentInstructionStack(affectStack ? -1 : 0);
                 break;
 
             case DMReference.Type.SrcProc:
@@ -306,7 +342,7 @@ public void WriteReference(DMReference reference, Location location, bool affect
 
             case DMReference.Type.ListIndex:
                 _annotatedBytecode[^1].AddArg(compiler, new AnnotatedBytecodeReference(reference.RefType, location));
-                ResizeStack(affectStack ? -2 : 0);
+                ResizeCurrentInstructionStack(affectStack ? -2 : 0);
                 break;
 
             case DMReference.Type.SuperProc:

DMCompiler/Optimizer/AnnotatedBytecode.cs

@@ -37,20 +37,41 @@ public AnnotatedBytecodeInstruction(AnnotatedBytecodeInstruction instruction, Li
     public AnnotatedBytecodeInstruction(DreamProcOpcode op, List<IAnnotatedBytecode> args) {
         Opcode = op;
         OpcodeMetadata metadata = OpcodeMetadataCache.GetMetadata(op);
-        StackSizeDelta = metadata.StackDelta;
+        StackSizeDelta = metadata.StackDelta + GetArgsStackSizeDelta(args);
         Location = new Location("Internal", null, null);
         ValidateArgs(metadata, args);
         _args = args;
     }
 
     public AnnotatedBytecodeInstruction(DreamProcOpcode opcode, int stackSizeDelta, List<IAnnotatedBytecode> args) {
         Opcode = opcode;
-        StackSizeDelta = stackSizeDelta;
+        StackSizeDelta = stackSizeDelta + GetArgsStackSizeDelta(args);
         Location = new Location("Internal", null, null);
         ValidateArgs(OpcodeMetadataCache.GetMetadata(opcode), args);
         _args = args;
     }
 
+    private int GetArgsStackSizeDelta(List<IAnnotatedBytecode> args) {
+        int delta = 0;
+
+        foreach (IAnnotatedBytecode arg in args) {
+            if (arg is not AnnotatedBytecodeReference reference)
+                continue;
+
+            delta += GetReferenceStackSizeDelta(reference.RefType);
+        }
+
+        return delta;
+    }
+
+    private int GetReferenceStackSizeDelta(DMReference.Type refType) {
+        return refType switch {
+            DMReference.Type.Field => -1,
+            DMReference.Type.ListIndex => -2,
+            _ => 0
+        };
+    }
+
     private void ValidateArgs(OpcodeMetadata metadata, List<IAnnotatedBytecode> args) {
         if (metadata.VariableArgs) {
             if (args[0] is not AnnotatedBytecodeInteger) {