Flatten nested string concatenation expressions into a single IrStringConcatenation, in a separate lowering phase.

Consolidating these into IrStringConcatenations allows the backend to produce efficient code for string concatenations (e.g., using StringBuilder for JVM).

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/FlattenStringConcatenationLowering.kt

@@ -0,0 +1,143 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.common.lower
+
+import org.jetbrains.kotlin.backend.common.CommonBackendContext
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrStringConcatenation
+import org.jetbrains.kotlin.ir.expressions.impl.IrStringConcatenationImpl
+import org.jetbrains.kotlin.ir.types.isString
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.name.Name
+
+val flattenStringConcatenationPhase = makeIrFilePhase(
+    ::FlattenStringConcatenationLowering,
+    name = "FlattenStringConcatenationLowering",
+    description = "Flatten nested string concatenation expressions into a single IrStringConcatenation"
+)
+
+/**
+ * Flattens nested string concatenation expressions into a single [IrStringConcatenation]. Consolidating these into IrStringConcatenations
+ * allows the backend to produce efficient code for string concatenations (e.g., using StringBuilder for JVM).
+ *
+ * Example expression:
+ *
+ *   val s = "1" + 2 + ("s1: '$s1'" + 3.0 + null)
+ *
+ * IR before lowering:
+ *
+ *   VAR name:s type:kotlin.String flags:val
+ *     CALL 'plus(Any?): String' type=kotlin.String origin=PLUS
+ *       $this: CALL 'plus(Any?): String' type=kotlin.String origin=PLUS
+ *         $this: CONST String type=kotlin.String value="1"
+ *         other: CONST Int type=kotlin.Int value=2
+ *       other: CALL 'plus(Any?): String' type=kotlin.String origin=PLUS
+ *         $this: CALL 'plus(Any?): String' type=kotlin.String origin=PLUS
+ *           $this: STRING_CONCATENATION type=kotlin.String
+ *             CONST String type=kotlin.String value="s1: '"
+ *             GET_VAR 's1: String' type=kotlin.String origin=null
+ *             CONST String type=kotlin.String value="'"
+ *           other: CONST Double type=kotlin.Double value=3.0
+ *         other: CONST Null type=kotlin.Nothing? value=null
+ *
+ * IR after lowering:
+ *
+ *   VAR name:s type:kotlin.String flags:val
+ *     STRING_CONCATENATION type=kotlin.String
+ *       CONST String type=kotlin.String value="1"
+ *       CONST Int type=kotlin.Int value=2
+ *       CONST String type=kotlin.String value="s1: '"
+ *       GET_VAR 's1: String' type=kotlin.String origin=null
+ *       CONST String type=kotlin.String value="'"
+ *       CONST Double type=kotlin.Double value=3.0
+ *       CONST Null type=kotlin.Nothing? value=null
+ */
+class FlattenStringConcatenationLowering(val context: CommonBackendContext) : FileLoweringPass, IrElementTransformerVoid() {
+
+    companion object {
+        private val PLUS_NAME = Name.identifier("plus")
+
+        /** @return true if the given expression is a [IrStringConcatenation] or [String.plus] [IrCall]. */
+        private fun isStringConcatenationExpression(expression: IrExpression): Boolean {
+            return when (expression) {
+                is IrStringConcatenation -> true
+                is IrCall -> {
+                    val dispatchReceiver = expression.dispatchReceiver
+                    dispatchReceiver != null &&
+                            dispatchReceiver.type.isString() &&
+                            expression.symbol.owner.name == PLUS_NAME &&
+                            expression.type.isString() &&
+                            expression.valueArgumentsCount == 1
+                }
+                else -> false
+            }
+        }
+
+        /** Recursively collects string concatenation arguments from the given expression. */
+        private fun collectStringConcatenationArguments(expression: IrExpression): List<IrExpression> {
+            val arguments = mutableListOf<IrExpression>()
+            expression.acceptChildrenVoid(object : IrElementVisitorVoid {
+
+                override fun visitElement(element: IrElement) {
+                    // Theoretically this is unreachable code since all descendants of IrExpressions are IrExpressions.
+                    element.acceptChildrenVoid(this)
+                }
+
+                override fun visitCall(expression: IrCall) {
+                    if (isStringConcatenationExpression(expression)) {
+                        // Recursively collect from call arguments.
+                        expression.acceptChildrenVoid(this)
+                    } else {
+                        // Add call itself as an argument.
+                        arguments.add(expression)
+                    }
+                }
+
+                override fun visitStringConcatenation(expression: IrStringConcatenation) {
+                    // Recursively collect from concatenation arguments.
+                    expression.acceptChildrenVoid(this)
+                }
+
+                override fun visitExpression(expression: IrExpression) {
+                    // These IrExpressions are neither IrCalls nor IrStringConcatenations and should be added as an argument.
+                    arguments.add(expression)
+                }
+            })
+
+            return arguments
+        }
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun visitExpression(expression: IrExpression): IrExpression {
+        // Only modify/flatten string concatenation expressions.
+        val transformedExpression =
+            if (isStringConcatenationExpression(expression))
+                expression.run {
+                    IrStringConcatenationImpl(
+                        startOffset,
+                        endOffset,
+                        type,
+                        collectStringConcatenationArguments(this)
+                    )
+                }
+            else expression
+
+        transformedExpression.transformChildrenVoid(this)
+        return transformedExpression
+    }
+}

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -82,6 +82,7 @@ internal val jvmPhases = namedIrFilePhase(
             tailrecPhase then
             toArrayPhase then
             jvmTypeOperatorLoweringPhase then
+            flattenStringConcatenationPhase then
             jvmBuiltinOptimizationLoweringPhase then
 
             makePatchParentsPhase(3)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -858,14 +858,26 @@ class ExpressionCodegen(
 
     override fun visitStringConcatenation(expression: IrStringConcatenation, data: BlockInfo): StackValue {
         expression.markLineNumber(startOffset = true)
-        AsmUtil.genStringBuilderConstructor(mv)
+        return when (expression.arguments.size) {
-        expression.arguments.forEach {
+            0 -> StackValue.constant("", expression.asmType)
-            val stackValue = gen(it, data)
+            1 -> {
-            AsmUtil.genInvokeAppendMethod(mv, stackValue.type, stackValue.kotlinType)
+                // Convert single arg to string.
+                val arg = expression.arguments[0]
+                val argStackValue = gen(arg, arg.asmType, data)
+                AsmUtil.genToString(argStackValue, argStackValue.type, argStackValue.kotlinType, typeMapper).put(expression.asmType, mv)
+                expression.onStack
+            }
+            else -> {
+                // Use StringBuilder to concatenate.
+                AsmUtil.genStringBuilderConstructor(mv)
+                expression.arguments.forEach {
+                    val stackValue = gen(it, data)
+                    AsmUtil.genInvokeAppendMethod(mv, stackValue.type, stackValue.kotlinType)
+                }
+                mv.invokevirtual("java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false)
+                expression.onStack
+            }
         }
-
-        mv.invokevirtual("java/lang/StringBuilder", "toString", "()Ljava/lang/String;", false)
-        return expression.onStack
     }
 
     override fun visitWhileLoop(loop: IrWhileLoop, data: BlockInfo): StackValue {

compiler/testData/codegen/box/strings/nestedConcat.kt

@@ -8,14 +8,34 @@ fun test1(s1: String, s2: String, s3: String) =
 fun test2(s1: String, s2: String, s3: String) =
         s1 + (s2 + s3)
 
-fun test3(s1: String, s2: String, s3: String) =
+fun test3(s1: String, s2: String, s3: String, s4: String) =
+        ((s1 + s2) + ((s3 + s4)))
+
+fun test4(s1: String, s2: String, s3: String) =
         "s1: $s1; " +
         "s2: $s2; " +
         "s3: $s3"
 
+fun test5(s1: String, s2: String, s3: String) =
+        "${"s1:" + "${" " + s1};"} " +
+        "${"s2:" + "${" " + s2};"} " +
+        "${"s3:" + "${" " + s3}"}"
+
+fun test6(s1: String, s2: String, s3: String) =
+    "${"s1:" + "${" " + s1};"} ${"s2:" + "${" " + s2};"} ${"s3:" + "${" " + s3}"}"
+
+fun test7(s1: String, s2: String, s3: String): String {
+    fun foo(s: String) = s
+    return "foo: " + foo(s1 + s2 + " ${foo("\${s3.length} = ${s3.length}")}")
+}
+
 fun box(): String {
     assertEquals("123", test1("1", "2", "3"))
     assertEquals("123", test2("1", "2", "3"))
-    assertEquals("s1: 1; s2: 2; s3: 3", test3("1", "2", "3"))
+    assertEquals("1234", test3("1", "2", "3", "4"))
+    assertEquals("s1: 1; s2: 2; s3: 3", test4("1", "2", "3"))
+    assertEquals("s1: 1; s2: 2; s3: 3", test5("1", "2", "3"))
+    assertEquals("s1: 1; s2: 2; s3: 3", test6("1", "2", "3"))
+    assertEquals("foo: 12 \${s3.length} = 1", test7("1", "2", "3"))
     return "OK"
 }

compiler/testData/codegen/bytecodeText/stringOperations/concat.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 class A() {
 
     override fun toString(): String {

compiler/testData/codegen/bytecodeText/stringOperations/nestedConcat.kt

@@ -1,16 +1,41 @@
-// IGNORE_BACKEND: JVM_IR
+// Uses 1 StringBuilder
 fun test1(s1: String, s2: String, s3: String) =
         (s1 + s2) + s3
 
+// Uses 1 StringBuilder
 fun test2(s1: String, s2: String, s3: String) =
         s1 + (s2 + s3)
 
+// Uses 1 StringBuilder
 fun test3(s1: String, s2: String, s3: String, s4: String) =
         ((s1 + s2) + ((s3 + s4)))
 
+// Combination of String.plus and string literal
+// Uses 1 StringBuilder
 fun test4(s1: String, s2: String, s3: String) =
         "s1: $s1; " +
         "s2: $s2; " +
         "s3: $s3"
 
-// 4 NEW java/lang/StringBuilder
+// Combination of String.plus and nested string literal
+// Uses 1 StringBuilder
+fun test5(s1: String, s2: String, s3: String) =
+        "${"s1:" + "${" " + s1};"}; " +
+        "${"s2:" + "${" " + s2};"}; " +
+        "${"s3:" + "${" " + s3}"}"
+
+// Top-level string concatenation element is a string literal
+// Uses 1 StringBuilder
+fun test6(s1: String, s2: String, s3: String) =
+    "${"s1:" + "${" " + s1};"} ${"s2:" + "${" " + s2};"} ${"s3:" + "${" " + s3}"}"
+
+// Uses 3 StringBuilders:
+// - In return expression
+// - In argument to 1st call to foo()
+// - In argument to 2nd call to foo() inside string literal
+fun test7(s1: String, s2: String, s3: String): String {
+    fun foo(s: String) = s
+    return "foo: " + foo(s1 + s2 + " ${foo("\${s3.length} = ${s3.length}")}")
+}
+
+// 9 NEW java/lang/StringBuilder

compiler/testData/codegen/bytecodeText/stringOperations/plusAssign.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 class A() {
 
     override fun toString(): String {

compiler/testData/codegen/bytecodeText/stringOperations/primitivesAsStringTemplates.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 fun test(b: Byte, s: Short, i: Int, l: Long) {
     "$b"
     "$s"