InterpolatedStringInjectorProcessor in functional-style

idea/idea-jvm/src/org/jetbrains/kotlin/idea/injection/InterpolatedStringInjectorProcessor.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.idea.injection
 
 import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.util.Trinity
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiLanguageInjectionHost
 import org.intellij.plugins.intelliLang.inject.InjectedLanguage
 import org.intellij.plugins.intelliLang.inject.InjectorUtils
@@ -27,87 +28,71 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
-import java.util.*
 
-data class InjectionSplitResult(val isUnparsable: Boolean, val ranges: List<Trinity<PsiLanguageInjectionHost, InjectedLanguage, TextRange>>)
+typealias Injection = Trinity<PsiLanguageInjectionHost, InjectedLanguage, TextRange>
+
+data class InjectionSplitResult(val isUnparsable: Boolean, val ranges: List<Injection>)
 
 fun splitLiteralToInjectionParts(injection: BaseInjection, literal: KtStringTemplateExpression): InjectionSplitResult? {
     InjectorUtils.getLanguage(injection) ?: return null
 
-    val children = literal.children.toList()
+    fun injectionRange(range: TextRange, prefix: String, suffix: String): Injection {
-    val len = children.size
-
-    if (children.isEmpty()) return null
-
-    val result = ArrayList<Trinity<PsiLanguageInjectionHost, InjectedLanguage, TextRange>>()
-
-    fun addInjectionRange(range: TextRange, prefix: String, suffix: String) {
         TextRange.assertProperRange(range, injection)
         val injectedLanguage = InjectedLanguage.create(injection.injectedLanguageId, prefix, suffix, true)!!
-        result.add(Trinity.create(literal, injectedLanguage, range))
+        return Trinity.create(literal, injectedLanguage, range)
     }
 
-    fun suffix(i: Int) = if (i == len - 1) injection.suffix else ""
+    tailrec fun collectInjections(
-
+        children: List<PsiElement>,
-    var unparsable = false
+        pendingPrefix: String,
-
+        unparseable: Boolean,
-    var prefix = injection.prefix
+        collected: MutableList<Injection>
-
+    ): InjectionSplitResult {
-    var i = 0
+        val child = children.firstOrNull() ?: return InjectionSplitResult(unparseable, collected)
-    while (i < len) {
+        val tail = children.subList(1, children.size)
-        val child = children[i]
         val partOffsetInParent = child.startOffsetInParent
 
-        if (child is KtSimpleNameStringTemplateEntry || child is KtBlockStringTemplateEntry) {
+        if (child is KtLiteralStringTemplateEntry || child is KtEscapeStringTemplateEntry) {
-            if (!prefix.isEmpty() || i == 0) {
+            val consequentStringsCount = tail.asSequence()
-                // Store part with prefix before replacing it
+                .takeWhile { it is KtLiteralStringTemplateEntry || it is KtEscapeStringTemplateEntry }
-                addInjectionRange(TextRange.from(partOffsetInParent, 0), prefix, suffix(i))
+                .count()
-            }
 
-            prefix = when (child) {
+            val lastChild = children[consequentStringsCount]
-                is KtSimpleNameStringTemplateEntry -> {
+            val remaining = tail.subList(consequentStringsCount, tail.size)
-                    tryEvaluateConstant(child.expression) ?: run {
-                        unparsable = true
-                        child.expression?.text ?: NO_VALUE_NAME
-                    }
-                }
-                is KtBlockStringTemplateEntry -> {
-                    tryEvaluateConstant(child.expression) ?: run {
-                        unparsable = true
-                        NO_VALUE_NAME
-                    }
-                }
-                else -> {
-                    error("Child type should be KtSimpleNameStringTemplateEntry or KtBlockStringTemplateEntry")
-                }
-            }
 
-            if (i == len - 1 && !prefix.isEmpty()) {
+            collected += injectionRange(
-                // There won't be more elements, so create part with prefix right away
+                TextRange.create(partOffsetInParent, lastChild.startOffsetInParent + lastChild.textLength),
-                addInjectionRange(TextRange.from(partOffsetInParent + child.textLength, 0), prefix, suffix(i))
+                pendingPrefix,
-            }
+                if (remaining.isEmpty()) injection.suffix else ""
-        } else {
-            if (child is KtLiteralStringTemplateEntry || child is KtEscapeStringTemplateEntry) {
-                // Merge all consecutive string literals into one part
-                i += children.countFromWhile(i + 1) { it is KtLiteralStringTemplateEntry || it is KtEscapeStringTemplateEntry }
-            } else {
-                unparsable = true
-            }
-
-            val lastInPart = children[i]
-            addInjectionRange(
-                TextRange.create(partOffsetInParent, lastInPart.startOffsetInParent + lastInPart.textLength),
-                prefix,
-                suffix(i)
             )
+            return collectInjections(remaining, "", unparseable, collected)
+        } else {
+            if (pendingPrefix.isNotEmpty() || collected.isEmpty()) {
+                // Store pending prefix before creating a new one,
+                // or if it is a first part then create a dummy injection to distinguish "inner" prefixes
+                collected += injectionRange(TextRange.from(partOffsetInParent, 0), pendingPrefix, "")
+            }
 
-            prefix = ""
+            val (prefix, myUnparseable) = makePlaceholder(child)
+
+            if (tail.isEmpty()) {
+                // There won't be more elements, so create part with prefix right away
+                collected += injectionRange(TextRange.from(partOffsetInParent + child.textLength, 0), prefix, injection.suffix)
+            }
+            return collectInjections(tail, prefix, unparseable || myUnparseable, collected)
         }
-
-        i++
     }
 
-    return InjectionSplitResult(unparsable, result)
+    return collectInjections(literal.children.toList(), injection.prefix, false, ArrayList())
+}
+
+private fun makePlaceholder(child: PsiElement): Pair<String, Boolean> = when (child) {
+    is KtSimpleNameStringTemplateEntry ->
+        tryEvaluateConstant(child.expression)?.let { it to false } ?: ((child.expression?.text ?: NO_VALUE_NAME) to true)
+    is KtBlockStringTemplateEntry ->
+        tryEvaluateConstant(child.expression)?.let { it to false } ?: (NO_VALUE_NAME to true)
+    else ->
+        ((child.text ?: NO_VALUE_NAME) to true)
 }
 
 private fun tryEvaluateConstant(ktExpression: KtExpression?) =
@@ -119,15 +104,3 @@ private fun tryEvaluateConstant(ktExpression: KtExpression?) =
     }
 
 private const val NO_VALUE_NAME = "missingValue"
-
-private inline fun <T : Any> List<T>.countFromWhile(from: Int, predicate: (T) -> Boolean): Int {
-    var i = from
-    while (i < size) {
-        if (!predicate(get(i))) {
-            break
-        }
-        i++
-    }
-
-    return i - from
-}