Allow injection in strings with interpolation (KT-6610)

#KT-6610 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/KtStringTemplateExpression.java

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.psi;
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.ElementManipulators;
 import com.intellij.psi.LiteralTextEscaper;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.psi.tree.TokenSet;
 import org.jetbrains.annotations.NotNull;
@@ -58,4 +59,14 @@ public class KtStringTemplateExpression extends KtExpressionImpl implements PsiL
     public LiteralTextEscaper<? extends PsiLanguageInjectionHost> createLiteralTextEscaper() {
         return new KotlinStringLiteralTextEscaper(this);
     }
+
+    public boolean hasInterpolation() {
+        for (PsiElement child : getChildren()) {
+            if (child instanceof KtSimpleNameStringTemplateEntry || child instanceof KtBlockStringTemplateEntry) {
+                return true;
+            }
+        }
+
+        return false;
+    }
 }

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

@@ -0,0 +1,100 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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
+import org.intellij.plugins.intelliLang.inject.config.BaseInjection
+import org.jetbrains.kotlin.psi.*
+import java.util.*
+
+data class InjectionSplitResult(val isUnparsable: Boolean, val ranges: List<Trinity<PsiLanguageInjectionHost, InjectedLanguage, TextRange>>)
+
+fun splitLiteralToInjectionParts(injection: BaseInjection, literal: KtStringTemplateExpression): InjectionSplitResult? {
+    InjectorUtils.getLanguage(injection) ?: return null
+
+    val children = literal.children.toList()
+    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))
+    }
+
+    var unparsable = false
+
+    var prefix = injection.prefix
+    val lastChild = children.lastOrNull()
+
+    var i = 0
+    while (i < len) {
+        val child = children[i]
+        val partOffsetInParent = child.startOffsetInParent
+
+        val part = when (child) {
+            is KtLiteralStringTemplateEntry, is KtEscapeStringTemplateEntry -> {
+                val partSize = children.subList(i, len).asSequence()
+                        .takeWhile { it is KtLiteralStringTemplateEntry || it is KtEscapeStringTemplateEntry }
+                        .count()
+                i += partSize - 1
+                children[i]
+            }
+            is KtSimpleNameStringTemplateEntry -> {
+                unparsable = true
+                child.expression?.text ?: NO_VALUE_NAME
+            }
+            is KtBlockStringTemplateEntry -> {
+                unparsable = true
+                NO_VALUE_NAME
+            }
+            else -> {
+                unparsable = true
+                child
+            }
+        }
+
+        val suffix = if (child == lastChild) injection.suffix else ""
+
+        if (part is PsiElement) {
+            addInjectionRange(TextRange.create(partOffsetInParent, part.startOffsetInParent + part.textLength), prefix, suffix)
+        }
+        else if (!prefix.isEmpty() || i == 0) {
+            addInjectionRange(TextRange.from(partOffsetInParent, 0), prefix, suffix)
+        }
+
+        prefix = part as? String ?: ""
+        i++
+    }
+
+    if (lastChild != null && !prefix.isEmpty()) {
+        // Last element was interpolated part, need to add a range after it
+        addInjectionRange(TextRange.from(lastChild.startOffsetInParent + lastChild.textLength, 0), prefix, injection.suffix)
+    }
+
+    return InjectionSplitResult(unparsable, result)
+}
+
+private val NO_VALUE_NAME = "missingValue"

idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjectionSupport.kt

@@ -57,7 +57,9 @@ class KotlinLanguageInjectionSupport : AbstractLanguageInjectionSupport() {
         val configuration = Configuration.getProjectInstance(host.project).advancedConfiguration
         if (!configuration.isSourceModificationAllowed) {
             // It's not allowed to modify code without explicit permission. Postpone adding @Inject or comment till it granted.
-            host.putUserData(InjectLanguageAction.FIX_KEY, Processor { host -> addInjectionInstructionInCode(language, host) })
+            host.putUserData(InjectLanguageAction.FIX_KEY, Processor { fixHost ->
+                addInjectionInstructionInCode(language, fixHost)
+            })
             return false
         }
 

idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt

@@ -25,11 +25,15 @@ import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiReference
+import com.intellij.psi.impl.source.tree.injected.InjectedLanguageUtil
 import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.psi.util.PsiTreeUtil
 import org.intellij.plugins.intelliLang.Configuration
+import org.intellij.plugins.intelliLang.inject.InjectedLanguage
 import org.intellij.plugins.intelliLang.inject.InjectorUtils
+import org.intellij.plugins.intelliLang.inject.LanguageInjectionSupport
+import org.intellij.plugins.intelliLang.inject.TemporaryPlacesRegistry
 import org.intellij.plugins.intelliLang.inject.config.BaseInjection
 import org.intellij.plugins.intelliLang.inject.java.JavaLanguageInjectionSupport
 import org.intellij.plugins.intelliLang.util.AnnotationUtilEx
@@ -44,7 +48,10 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import java.util.*
 import kotlin.collections.ArrayList
 
-class KotlinLanguageInjector : MultiHostInjector {
+class KotlinLanguageInjector(
+        val configuration: Configuration,
+        val project: Project,
+        val temporaryPlacesRegistry: TemporaryPlacesRegistry) : MultiHostInjector {
     companion object {
         private val STRING_LITERALS_REGEXP = "\"([^\"]*)\"".toRegex()
     }
@@ -61,10 +68,38 @@ class KotlinLanguageInjector : MultiHostInjector {
 
         if (!ProjectRootsUtil.isInProjectOrLibSource(ktHost)) return
 
-        val injectionInfo = findInjectionInfo(context) ?: return
+        val containingFile = ktHost.containingFile
-        InjectorUtils.getLanguageByString(injectionInfo.languageId) ?: return
+        val tempInjectedLanguage: InjectedLanguage? = temporaryPlacesRegistry.getLanguageFor(ktHost, containingFile)
 
-        InjectorUtils.registerInjectionSimple(ktHost, injectionInfo.toBaseInjection(support)!!, support, registrar)
+        val baseInjection: BaseInjection = if (tempInjectedLanguage == null) {
+            val injectionInfo = findInjectionInfo(context) ?: return
+            injectionInfo.toBaseInjection(support)
+        }
+        else {
+            InjectorUtils.putInjectedFileUserData(registrar, LanguageInjectionSupport.TEMPORARY_INJECTED_LANGUAGE, tempInjectedLanguage)
+            BaseInjection(support.id).apply {
+                injectedLanguageId = tempInjectedLanguage.id
+                prefix = tempInjectedLanguage.prefix
+                suffix = tempInjectedLanguage.suffix
+            }
+        } ?: return
+
+        val language = InjectorUtils.getLanguageByString(baseInjection.injectedLanguageId) ?: return
+
+        if (ktHost.hasInterpolation()) {
+            val file = ktHost.containingKtFile
+            val parts = splitLiteralToInjectionParts(baseInjection, ktHost) ?: return
+
+            if (parts.ranges.isEmpty()) return
+
+            InjectorUtils.registerInjection(language, parts.ranges, file, registrar)
+            InjectorUtils.registerSupport(support, false, registrar)
+            InjectorUtils.putInjectedFileUserData(registrar, InjectedLanguageUtil.FRANKENSTEIN_INJECTION,
+                                                  if (parts.isUnparsable) java.lang.Boolean.TRUE else null)
+        }
+        else {
+            InjectorUtils.registerInjectionSimple(ktHost, baseInjection, support, registrar)
+        }
     }
 
     override fun elementsToInjectIn(): List<Class<out PsiElement>> {

idea/tests/org/jetbrains/kotlin/psi/AbstractInjectionTest.kt

@@ -16,7 +16,9 @@
 
 package org.jetbrains.kotlin.psi
 
+import com.intellij.injected.editor.DocumentWindowImpl
 import com.intellij.injected.editor.EditorWindow
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.injection.Injectable
 import com.intellij.testFramework.LightProjectDescriptor
 import junit.framework.TestCase
@@ -40,9 +42,17 @@ abstract class AbstractInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
         }
     }
 
+    data class ShredInfo(
+            val range: TextRange,
+            val hostRange: TextRange,
+            val prefix: String = "",
+            val suffix: String = "") {
+    }
+
     protected fun doInjectionPresentTest(
             @Language("kotlin") text: String, @Language("Java") javaText: String? = null,
-            languageId: String? = null, unInjectShouldBePresent: Boolean = true) {
+            languageId: String? = null, unInjectShouldBePresent: Boolean = true,
+            shreds: List<ShredInfo>? = null) {
         if (javaText != null) {
             myFixture.configureByText("${getTestName(true)}.java", javaText.trimIndent())
         }
@@ -50,6 +60,16 @@ abstract class AbstractInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
         myFixture.configureByText("${getTestName(true)}.kt", text.trimIndent())
 
         assertInjectionPresent(languageId, unInjectShouldBePresent)
+
+        if (shreds != null) {
+            val actualShreds = (editor.document as DocumentWindowImpl).shreds.map {
+                ShredInfo(it.range, it.rangeInsideHost, it.prefix, it.suffix)
+            }
+
+            assertOrderedEquals(
+                    actualShreds.sortedBy { it.range.startOffset },
+                    shreds.sortedBy { it.range.startOffset })
+        }
     }
 
     protected fun assertInjectionPresent(languageId: String?, unInjectShouldBePresent: Boolean) {
@@ -103,4 +123,6 @@ abstract class AbstractInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
             configuration.isSourceModificationAllowed = allowed
         }
     }
-}
+
+    fun range(start: Int, end: Int) = TextRange.create(start, end)
+}

idea/tests/org/jetbrains/kotlin/psi/KotlinInjectionTest.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.psi
 
 import com.intellij.lang.html.HTMLLanguage
+import com.intellij.openapi.fileTypes.PlainTextLanguage
 import org.intellij.lang.regexp.RegExpLanguage
 import org.intellij.plugins.intelliLang.Configuration
 import org.intellij.plugins.intelliLang.inject.config.BaseInjection
@@ -324,4 +325,89 @@ class KotlinInjectionTest : AbstractInjectionTest() {
             """,
             languageId = HTMLLanguage.INSTANCE.id, unInjectShouldBePresent = false
     )
+
+    fun testInjectionOnInterpolationWithAnnotation() = doInjectionPresentTest(
+            """
+            val b = 2
+
+            @org.intellij.lang.annotations.Language("HTML")
+            val test = "<caret>simple${'$'}{b}.kt"
+            """,
+            unInjectShouldBePresent = false,
+            shreds = listOf(
+                    ShredInfo(range(0, 6), hostRange=range(1, 7)),
+                    ShredInfo(range(6, 21), hostRange=range(11, 14), prefix="missingValue")
+            )
+    )
+
+    fun testInjectionOnInterpolatedStringWithComment() = doInjectionPresentTest(
+            """
+            val some = 42
+            // language=HTML
+            val test = "<ht<caret>ml>${'$'}some</html>"
+            """,
+            languageId = HTMLLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+            shreds = listOf(
+                    ShredInfo(range(0, 6), hostRange = range(1, 7)),
+                    ShredInfo(range(6, 17), hostRange = range(12, 19), prefix="some"))
+    )
+
+    fun testEditorShortShreadsInInterpolatedInjection() = doInjectionPresentTest(
+            """
+            val s = 42
+            // language=TEXT
+            val test = "${'$'}s <caret>text ${'$'}s${'$'}{s}${'$'}s text ${'$'}s"
+            """,
+            languageId = PlainTextLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+            shreds = listOf(
+                    ShredInfo(range(0, 0), hostRange=range(1, 1)),
+                    ShredInfo(range(0, 7), hostRange=range(3, 9), prefix="s"),
+                    ShredInfo(range(7, 8), hostRange=range(11, 11), prefix="s"),
+                    ShredInfo(range(8, 20), hostRange=range(15, 15), prefix="missingValue"),
+                    ShredInfo(range(20, 27), hostRange=range(17, 23), prefix="s"),
+                    ShredInfo(range(27, 28), hostRange=range(25, 25), prefix="s")
+            )
+    )
+
+    fun testEditorLongShreadsInInterpolatedInjection() = doInjectionPresentTest(
+            """
+            val s = 42
+            // language=TEXT
+            val test = "${'$'}{s} <caret>text ${'$'}{s}${'$'}s${'$'}{s} text ${'$'}{s}"
+            """,
+            languageId = PlainTextLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+            shreds = listOf(
+                    ShredInfo(range(0, 0), hostRange=range(1, 1)),
+                    ShredInfo(range(0, 18), hostRange=range(5, 11), prefix="missingValue"),
+                    ShredInfo(range(18, 30), hostRange=range(15, 15), prefix="missingValue"),
+                    ShredInfo(range(30, 31), hostRange=range(17, 17), prefix="s"),
+                    ShredInfo(range(31, 49), hostRange=range(21, 27), prefix="missingValue"),
+                    ShredInfo(range(49, 61), hostRange=range(31, 31), prefix="missingValue")
+            )
+    )
+
+    fun testEditorShreadsWithEscapingInjection() = doInjectionPresentTest(
+            """
+            // language=TEXT
+            val test = "\rte<caret>xt\ttext\n\t"
+            """,
+            languageId = PlainTextLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+            shreds = listOf(
+                    ShredInfo(range(0, 12), hostRange=range(1, 17))
+            )
+    )
+
+    fun testEditorShreadsInInterpolatedWithEscapingInjection() = doInjectionPresentTest(
+            """
+            val s = 1
+            // language=TEXT
+            val test = "\r${'$'}s te<caret>xt${'$'}s\ttext\n\t"
+            """,
+            languageId = PlainTextLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+            shreds = listOf(
+                    ShredInfo(range(0, 1), hostRange=range(1, 3)),
+                    ShredInfo(range(1, 7), hostRange=range(5, 10), prefix="s"),
+                    ShredInfo(range(7, 15), hostRange=range(12, 22), prefix="s")
+            )
+    )
 }