Allow injection in strings with interpolation (KT-6610)

#KT-6610 In Progress

compiler/frontend/src/org/jetbrains/kotlin/psi/KotlinStringLiteralTextEscaper.kt

@@ -38,13 +38,6 @@ class KotlinStringLiteralTextEscaper(host: KtStringTemplateExpression): LiteralT
                 continue
             }
             when (child) {
-                is KtLiteralStringTemplateEntry -> {
-                    val textRange = rangeInsideHost.intersection(childRange)!!.shiftRight(-childRange.startOffset)
-                    outChars.append(child.getText(), textRange.startOffset, textRange.endOffset)
-                    repeat(textRange.length) {
-                        sourceOffsetsList.add(sourceOffset++)
-                    }
-                }
                 is KtEscapeStringTemplateEntry -> {
                     if (!rangeInsideHost.contains(childRange)) {
                         //don't allow injection if its range starts or ends inside escaped sequence
@@ -57,7 +50,13 @@ class KotlinStringLiteralTextEscaper(host: KtStringTemplateExpression): LiteralT
                     }
                     sourceOffset += child.getTextLength()
                 }
-                else -> return false
+                else -> {
+                    val textRange = rangeInsideHost.intersection(childRange)!!.shiftRight(-childRange.startOffset)
+                    outChars.append(child.text, textRange.startOffset, textRange.endOffset)
+                    repeat(textRange.length) {
+                        sourceOffsetsList.add(sourceOffset++)
+                    }
+                }
             }
         }
         sourceOffsetsList.add(sourceOffset)

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

@@ -22,11 +22,10 @@ import com.intellij.psi.LiteralTextEscaper;
 import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.psi.tree.TokenSet;
 import org.jetbrains.annotations.NotNull;
-import org.jetbrains.kotlin.KtNodeTypes;
 import org.jetbrains.kotlin.lexer.KtTokens;
 
 public class KtStringTemplateExpression extends KtExpressionImpl implements PsiLanguageInjectionHost {
-    private static final TokenSet TOKENS_SUITABLE_FOR_INJECTION = TokenSet.create(KtNodeTypes.LITERAL_STRING_TEMPLATE_ENTRY, KtNodeTypes.ESCAPE_STRING_TEMPLATE_ENTRY);
+    private static final TokenSet CLOSE_QUOTE_TOKEN_SET = TokenSet.create(KtTokens.CLOSING_QUOTE);
 
     public KtStringTemplateExpression(@NotNull ASTNode node) {
         super(node);
@@ -44,14 +43,7 @@ public class KtStringTemplateExpression extends KtExpressionImpl implements PsiL
 
     @Override
     public boolean isValidHost() {
-        ASTNode node = getNode();
-        ASTNode child = node.getFirstChildNode().getTreeNext();
-        while (child != null) {
-            if (child.getElementType() == KtTokens.CLOSING_QUOTE) return true;
-            if (!TOKENS_SUITABLE_FOR_INJECTION.contains(child.getElementType())) return false;
-            child = child.getTreeNext();
-        }
-        return false;
+        return getNode().getChildren(CLOSE_QUOTE_TOKEN_SET).length != 0;
     }
 
     @Override

idea/tests/org/jetbrains/kotlin/psi/injection/StringInjectionHostTest.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.test.ConfigurationKind
 import org.jetbrains.kotlin.test.KotlinTestWithEnvironment
+import org.jetbrains.kotlin.utils.keysToMap
 import java.util.*
 
 class StringInjectionHostTest : KotlinTestWithEnvironment() {
@@ -41,6 +42,11 @@ class StringInjectionHostTest : KotlinTestWithEnvironment() {
         }
     }
 
+    fun testInterpolation1(): Unit = checkAllRanges("a \$b c")
+    fun testInterpolation2(): Unit = checkAllRanges("a \${b} c")
+    fun testInterpolation3(): Unit = checkAllRanges("a\${b}c")
+    fun testInterpolation4(): Unit = checkAllRanges("a \${b.foo()} c")
+
     fun testUnclosedSimpleLiteral() {
         assertFalse(stringExpression("\"").isValidHost)
         assertFalse(stringExpression("\"a").isValidHost)
@@ -124,6 +130,13 @@ class StringInjectionHostTest : KotlinTestWithEnvironment() {
         assertFalse(createLiteralTextEscaper().isOneLine)
     }
 
+    private fun checkAllRanges(str: String) {
+        with (quoted(str)) {
+            checkInjection(str, (0..str.length).keysToMap { it + 1 })
+            assertOneLine()
+        }
+    }
+
     private fun KtStringTemplateExpression.checkInjection(
             decoded: String, targetToSourceOffsets: Map<Int, Int>, rangeInHost: TextRange? = null
     ) {
@@ -132,8 +145,10 @@ class StringInjectionHostTest : KotlinTestWithEnvironment() {
             val escaper = createLiteralTextEscaper()
             val chars = StringBuilder(prefix)
             val range = rangeInHost ?: escaper.relevantTextRange
+
             assertTrue(escaper.decode(range, chars))
             assertEquals(decoded, chars.substring(prefix.length))
+
             val extendedOffsets = HashMap(targetToSourceOffsets)
             val beforeStart = targetToSourceOffsets.keys.min()!! - 1
             if (beforeStart >= 0) {