From 46b206adaf397eae7c8e73ca13f914b03ea370a2 Mon Sep 17 00:00:00 2001
From: ike709 <ike709@users.noreply.github.com>
Date: Sat, 16 Nov 2024 20:21:07 -0600
Subject: [PATCH] String concat const folding (#2061)

Co-authored-by: ike709 <ike709@github.com>
Co-authored-by: wixoa <wixoag@gmail.com>
---
 DMCompiler/DM/Builders/DMASTFolder.cs         |  7 -----
 DMCompiler/Optimizer/AnnotatedBytecode.cs     |  4 +++
 DMCompiler/Optimizer/PeepholeOptimizations.cs | 27 +++++++++++++++++++
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/DMCompiler/DM/Builders/DMASTFolder.cs b/DMCompiler/DM/Builders/DMASTFolder.cs
index 89f01fd584..0bbe08b1ea 100644
--- a/DMCompiler/DM/Builders/DMASTFolder.cs
+++ b/DMCompiler/DM/Builders/DMASTFolder.cs
@@ -219,13 +219,6 @@ private DMASTExpression FoldExpression(DMASTExpression? expression) {
 
                 break;
             }
-            case DMASTAdd add: {
-                DMASTConstantString? lhsString = add.LHS as DMASTConstantString;
-                DMASTConstantString? rhsString = add.RHS as DMASTConstantString;
-                if (lhsString != null && rhsString != null) return new DMASTConstantString(expression.Location, lhsString.Value + rhsString.Value);
-
-                break;
-            }
 
             #endregion Math
 
diff --git a/DMCompiler/Optimizer/AnnotatedBytecode.cs b/DMCompiler/Optimizer/AnnotatedBytecode.cs
index 4231539f36..3bb38adb09 100644
--- a/DMCompiler/Optimizer/AnnotatedBytecode.cs
+++ b/DMCompiler/Optimizer/AnnotatedBytecode.cs
@@ -235,6 +235,10 @@ public void SetLocation(Location loc) {
     public Location GetLocation() {
         return Location;
     }
+
+    public string ResolveString(DMCompiler compiler) {
+        return compiler.DMObjectTree.StringTable[Id];
+    }
 }
 
 internal sealed class AnnotatedBytecodeArgumentType(DMCallArgumentsType value, Location location) : IAnnotatedBytecode {
diff --git a/DMCompiler/Optimizer/PeepholeOptimizations.cs b/DMCompiler/Optimizer/PeepholeOptimizations.cs
index 395e0dc38e..a80d48f9d3 100644
--- a/DMCompiler/Optimizer/PeepholeOptimizations.cs
+++ b/DMCompiler/Optimizer/PeepholeOptimizations.cs
@@ -767,6 +767,33 @@ public void Apply(DMCompiler compiler, List<IAnnotatedBytecode> input, int index
     }
 }
 
+// PushString [constant]
+// PushString [constant]
+// Add
+// -> PushString [result]
+internal sealed class ConstFoldAddStrings : IPeepholeOptimization {
+    public ReadOnlySpan<DreamProcOpcode> GetOpcodes() {
+        return [
+            DreamProcOpcode.PushString,
+            DreamProcOpcode.PushString,
+            DreamProcOpcode.Add,
+        ];
+    }
+
+    public void Apply(DMCompiler compiler, List<IAnnotatedBytecode> input, int index) {
+        var firstInstruction = (AnnotatedBytecodeInstruction)input[index];
+        var firstString = firstInstruction.GetArg<AnnotatedBytecodeString>(0);
+        var secondString = ((AnnotatedBytecodeInstruction)input[index+1]).GetArg<AnnotatedBytecodeString>(0);
+
+        var combinedId = compiler.DMObjectTree.AddString(firstString.ResolveString(compiler) + secondString.ResolveString(compiler)); // TODO: Currently doesn't handle removing strings from the string tree that have no other references
+
+        var args = new List<IAnnotatedBytecode>(1) {new AnnotatedBytecodeString(combinedId, firstInstruction.Location)};
+
+        IPeepholeOptimization.ReplaceInstructions(input, index, 3,
+            new AnnotatedBytecodeInstruction(DreamProcOpcode.PushString, 1, args));
+    }
+}
+
 // PushFloat [constant]
 // PushFloat [constant]
 // Subtract