From 852d705c71c8e353f08de8f4a1f8fed72afe8be7 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Fri, 25 Sep 2020 20:14:57 +0300
Subject: [PATCH] FIR: optimize converting string expressions in raw FIR
 builder

Do get ASTNode for every string template entry as it is expensive operation
---
 .../kotlin/fir/builder/RawFirBuilder.kt        | 18 +++++++++++++++---
 .../kotlin/fir/builder/BaseFirBuilder.kt       |  3 ++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
index 02d780e999a..d60dc8ad81e 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
@@ -1449,9 +1449,21 @@ class RawFirBuilder(
             generateConstantExpressionByLiteral(expression)
 
         override fun visitStringTemplateExpression(expression: KtStringTemplateExpression, data: Unit): FirElement {
-            return expression.entries.toInterpolatingCall(expression) {
-                (this as KtStringTemplateEntryWithExpression).expression.toFirExpression(it)
-            }
+            return expression.entries.toInterpolatingCall(
+                expression,
+                getElementType = { element ->
+                    when (element) {
+                        is KtLiteralStringTemplateEntry -> KtNodeTypes.LITERAL_STRING_TEMPLATE_ENTRY
+                        is KtEscapeStringTemplateEntry -> KtNodeTypes.ESCAPE_STRING_TEMPLATE_ENTRY
+                        is KtSimpleNameStringTemplateEntry -> KtNodeTypes.SHORT_STRING_TEMPLATE_ENTRY
+                        is KtBlockStringTemplateEntry -> KtNodeTypes.LONG_STRING_TEMPLATE_ENTRY
+                        else -> error("invalid node type $element")
+                    }
+                },
+                convertTemplateEntry = {
+                    (this as KtStringTemplateEntryWithExpression).expression.toFirExpression(it)
+                },
+            )
         }
 
         override fun visitReturnExpression(expression: KtReturnExpression, data: Unit): FirElement {
diff --git a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt
index 7b9b87ca978..f6b2f03a083 100644
--- a/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt
+++ b/compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt
@@ -325,6 +325,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
 
     fun Array<out T?>.toInterpolatingCall(
         base: T,
+        getElementType: (T) -> IElementType = { it.elementType },
         convertTemplateEntry: T?.(String) -> FirExpression
     ): FirExpression {
         return buildStringConcatenationCall {
@@ -333,7 +334,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
             argumentList = buildArgumentList {
                 L@ for (entry in this@toInterpolatingCall) {
                     if (entry == null) continue
-                    arguments += when (entry.elementType) {
+                    arguments += when (getElementType(entry)) {
                         OPEN_QUOTE, CLOSING_QUOTE -> continue@L
                         LITERAL_STRING_TEMPLATE_ENTRY -> {
                             sb.append(entry.asText)