Support prefix and suffix in Kotlin injection with comments and annotations (KT-23005)

#KT-23005 Fixed

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2000-2018 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.idea.injection
+
+import org.intellij.plugins.intelliLang.inject.config.BaseInjection
+
+internal class InjectionInfo(val languageId: String?, val prefix: String?, val suffix: String?) {
+    fun toBaseInjection(injectionSupport: KotlinLanguageInjectionSupport): BaseInjection? {
+        if (languageId == null) return null
+
+        val baseInjection = BaseInjection(injectionSupport.id)
+        baseInjection.injectedLanguageId = languageId
+
+        if (prefix != null) {
+            baseInjection.prefix = prefix
+        }
+
+        if (suffix != null) {
+            baseInjection.suffix = suffix
+        }
+
+        return baseInjection
+    }
+
+    companion object {
+        fun fromBaseInjection(baseInjection: BaseInjection?): InjectionInfo? {
+            if (baseInjection == null) {
+                return null
+            }
+
+            return InjectionInfo(
+                baseInjection.injectedLanguageId,
+                baseInjection.prefix,
+                baseInjection.suffix
+            )
+        }
+    }
+}

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

@@ -95,27 +95,39 @@ class KotlinLanguageInjectionSupport : AbstractLanguageInjectionSupport() {
         return null
     }
 
-    fun findInjectionCommentLanguageId(host: KtElement): String? {
+    fun findCommentInjection(host: KtElement): BaseInjection? {
-        return InjectorUtils.findCommentInjection(host, "", null)?.injectedLanguageId
+        return InjectorUtils.findCommentInjection(host, "", null)
     }
 
-    fun findInjectionComment(host: KtElement): PsiComment? {
+    private fun findInjectionComment(host: KtElement): PsiComment? {
         val commentRef = Ref.create<PsiElement>(null)
         InjectorUtils.findCommentInjection(host, "", commentRef) ?: return null
 
         return commentRef.get() as? PsiComment
     }
 
-    fun findAnnotationInjectionLanguageId(host: KtElement): String? {
+    internal fun findAnnotationInjectionLanguageId(host: KtElement): InjectionInfo? {
         val annotationEntry = findAnnotationInjection(host) ?: return null
-        return extractLanguageFromInjectAnnotation(annotationEntry)
+        val extractLanguageFromInjectAnnotation = extractLanguageFromInjectAnnotation(annotationEntry) ?: return null
+        val prefix = extractStringArgumentByName(annotationEntry, "prefix")
+        val suffix = extractStringArgumentByName(annotationEntry, "suffix")
+        return InjectionInfo(extractLanguageFromInjectAnnotation, prefix, suffix)
     }
 }
 
-fun extractLanguageFromInjectAnnotation(annotationEntry: KtAnnotationEntry): String? {
+private fun extractStringArgumentByName(annotationEntry: KtAnnotationEntry, name: String): String? {
-    val firstArgument: ValueArgument = annotationEntry.valueArguments.firstOrNull() ?: return null
+    val namedArgument: ValueArgument = annotationEntry.valueArguments.firstOrNull { it.getArgumentName()?.asName?.asString() == name } ?: return null
+    return extractStringValue(namedArgument)
+}
 
-    val firstStringArgument = firstArgument.getArgumentExpression() as? KtStringTemplateExpression ?: return null
+
+private fun extractLanguageFromInjectAnnotation(annotationEntry: KtAnnotationEntry): String? {
+    val firstArgument: ValueArgument = annotationEntry.valueArguments.firstOrNull() ?: return null
+    return extractStringValue(firstArgument)
+}
+
+private fun extractStringValue(valueArgument: ValueArgument): String? {
+    val firstStringArgument = valueArgument.getArgumentExpression() as? KtStringTemplateExpression ?: return null
     val firstStringEntry = firstStringArgument.entries.singleOrNull() ?: return null
 
     return firstStringEntry.text

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

@@ -176,11 +176,7 @@ class KotlinLanguageInjector(
 
     private fun injectWithExplicitCodeInstruction(host: KtElement): InjectionInfo? {
         val support = kotlinSupport ?: return null
-        val languageId =
+        return InjectionInfo.fromBaseInjection(support.findCommentInjection(host)) ?: support.findAnnotationInjectionLanguageId(host)
-                support.findInjectionCommentLanguageId(host) ?:
-                support.findAnnotationInjectionLanguageId(host) ?:
-                return null
-        return InjectionInfo(languageId, null, null)
     }
 
     private fun injectWithReceiver(host: KtElement): InjectionInfo? {
@@ -354,25 +350,6 @@ class KotlinLanguageInjector(
         return configuration.advancedConfiguration.dfaOption == Configuration.DfaOption.OFF
     }
 
-    private class InjectionInfo(val languageId: String?, val prefix: String?, val suffix: String?) {
-        fun toBaseInjection(injectionSupport: KotlinLanguageInjectionSupport): BaseInjection? {
-            if (languageId == null) return null
-
-            val baseInjection = BaseInjection(injectionSupport.id)
-            baseInjection.injectedLanguageId = languageId
-
-            if (prefix != null) {
-                baseInjection.prefix = prefix
-            }
-
-            if (suffix != null) {
-                baseInjection.suffix = suffix
-            }
-
-            return baseInjection
-        }
-    }
-
     private fun processAnnotationInjectionInner(annotations: Array<PsiAnnotation>): InjectionInfo? {
         val id = AnnotationUtilEx.calcAnnotationValue(annotations, "value")
         val prefix = AnnotationUtilEx.calcAnnotationValue(annotations, "prefix")

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

@@ -416,6 +416,24 @@ class KotlinInjectionTest : AbstractInjectionTest() {
             )
     )
 
+    fun testSuffixPrefixWithAnnotation() = doInjectionPresentTest(
+        """
+            @org.intellij.lang.annotations.Language("TEXT", prefix = "abc", suffix = "ghi")
+            val test = "<caret>def"
+            """,
+        languageId = PlainTextLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+        shreds = listOf(ShredInfo(range(0, 9), hostRange=range(1, 4), prefix = "abc", suffix = "ghi"))
+    )
+
+    fun testSuffixPrefixInComment() = doInjectionPresentTest(
+        """
+            // language="TEXT" prefix="abc" suffix=ghi
+            val test = "<caret>def"
+            """,
+        languageId = PlainTextLanguage.INSTANCE.id, unInjectShouldBePresent = false,
+        shreds = listOf(ShredInfo(range(0, 9), hostRange=range(1, 4), prefix = "abc", suffix = "ghi"))
+    )
+
     fun testJavaAnnotationsPattern() {
         myFixture.addClass("""
                 @interface Matches { String value(); }