Enable injecting predefined and custom Java places in Kotlin

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

@@ -16,12 +16,20 @@
 
 package org.jetbrains.kotlin.idea.injection
 
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.InjectedLanguagePlaces
 import com.intellij.psi.LanguageInjector
 import com.intellij.psi.PsiLanguageInjectionHost
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.util.PsiTreeUtil
+import org.intellij.plugins.intelliLang.Configuration
 import org.intellij.plugins.intelliLang.inject.InjectorUtils
+import org.intellij.plugins.intelliLang.inject.java.JavaLanguageInjectionSupport
+import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtValueArgument
+import org.jetbrains.kotlin.psi.KtValueArgumentList
 import java.util.*
 
 class KotlinLanguageInjector : LanguageInjector {
@@ -30,14 +38,51 @@ class KotlinLanguageInjector : LanguageInjector {
     }
 
     override fun getLanguagesToInject(host: PsiLanguageInjectionHost, injectionPlacesRegistrar: InjectedLanguagePlaces) {
-        if (host !is KtElement) return
         if (!host.isValidHost) return
-        val support = kotlinSupport ?: return
+        val ktHost: KtElement = host as? KtElement ?: return
 
-        val languageId = support.findAnnotationInjectionLanguageId(host) ?: return
+        val injectionInfo =
+                injectWithExplicitCodeInstruction(ktHost)
+                ?: injectWithCall(ktHost)
+                ?: return
 
-        val language = InjectorUtils.getLanguageByString(languageId) ?: return
-
-        injectionPlacesRegistrar.addPlace(language, TextRange.from(0, host.getTextLength()), null, null)
+        val language = InjectorUtils.getLanguageByString(injectionInfo.languageId) ?: return
+        injectionPlacesRegistrar.addPlace(language, TextRange.from(0, ktHost.textLength), injectionInfo.prefix, injectionInfo.suffix)
     }
+
+    private fun injectWithExplicitCodeInstruction(host: KtElement): InjectionInfo? {
+        val support = kotlinSupport ?: return null
+        val languageId = support.findAnnotationInjectionLanguageId(host) ?: return null
+        return InjectionInfo(languageId, null, null)
+    }
+
+    private fun injectWithCall(host: KtElement): InjectionInfo? {
+        val ktHost: KtElement = host
+        val argument = ktHost.parent as? KtValueArgument ?: return null
+
+        val callExpression = PsiTreeUtil.getParentOfType(ktHost, KtCallExpression::class.java) ?: return null
+        val callee = callExpression.calleeExpression ?: return null
+
+        for (reference in callee.references) {
+            ProgressManager.checkCanceled()
+
+            val javaMethod = reference.resolve()
+            if (javaMethod !is PsiMethod) continue
+
+            val argumentIndex = (argument.parent as KtValueArgumentList).arguments.indexOf(argument)
+            val psiParameter = javaMethod.parameterList.parameters.getOrNull(argumentIndex) ?: continue
+
+            val injections = Configuration.getInstance().getInjections(JavaLanguageInjectionSupport.JAVA_SUPPORT_ID)
+
+            for (injection in injections) {
+                if (injection.acceptsPsiElement(psiParameter)) {
+                    return InjectionInfo(injection.injectedLanguageId, injection.prefix, injection.suffix)
+                }
+            }
+        }
+
+        return null
+    }
+
+    private class InjectionInfo(val languageId: String?, val prefix: String?, val suffix: String?)
 }

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

@@ -16,22 +16,21 @@
 
 package org.jetbrains.kotlin.psi
 
-import com.intellij.openapi.module.Module
-import com.intellij.openapi.roots.ContentEntry
-import com.intellij.openapi.roots.ModifiableRootModel
-import com.intellij.openapi.roots.OrderRootType
-import com.intellij.openapi.vfs.VfsUtil
+import com.intellij.injected.editor.EditorWindow
+import com.intellij.lang.html.HTMLLanguage
 import com.intellij.psi.impl.source.resolve.reference.impl.providers.FileReference
 import com.intellij.psi.injection.Injectable
 import com.intellij.testFramework.LightProjectDescriptor
-import com.intellij.testFramework.fixtures.DefaultLightProjectDescriptor
 import junit.framework.TestCase
+import org.intellij.lang.regexp.RegExpLanguage
 import org.intellij.plugins.intelliLang.Configuration
 import org.intellij.plugins.intelliLang.inject.InjectLanguageAction
 import org.intellij.plugins.intelliLang.inject.UnInjectLanguageAction
+import org.intellij.plugins.intelliLang.inject.config.BaseInjection
+import org.intellij.plugins.intelliLang.inject.config.InjectionPlace
 import org.intellij.plugins.intelliLang.references.FileReferenceInjector
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
-import java.io.File
+import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
 
 class KotlinInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
     fun testInjectUnInjectOnSimpleString() {
@@ -50,13 +49,45 @@ class KotlinInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
         TestCase.assertNull(myFixture.getReferenceAtCaretPosition())
     }
 
-    fun testInjectionWithCommentOnProperty() = doTestInjection(
+    fun testInjectionOnJavaPredefinedMethodWithAnnotation() = testInjectionPresent(
+            """
+            |val test1 = java.util.regex.Pattern.compile("<caret>pattern")
+            """,
+            RegExpLanguage.INSTANCE.id,
+            unInjectShouldBePresent = false
+    )
+
+    fun testInjectionOnJavaCustomInjectionWithAnnotation() {
+        val customInjection = BaseInjection("java")
+        customInjection.injectedLanguageId = HTMLLanguage.INSTANCE.id
+        val elementPattern = customInjection.compiler.createElementPattern(
+                """psiParameter().ofMethod(2, psiMethod().withName("replace").withParameters("int", "int", "java.lang.String").definedInClass("java.lang.StringBuilder"))""",
+                "HTML temp rule")
+        customInjection.setInjectionPlaces(InjectionPlace(elementPattern, true))
+
+        try {
+            Configuration.getInstance().replaceInjections(listOf(customInjection), listOf(), true)
+
+            testInjectionPresent(
+                    """
+                    |val stringBuilder = StringBuilder().replace(0, 0, "<caret><html></html>")
+                    """,
+                    HTMLLanguage.INSTANCE.id,
+                    unInjectShouldBePresent = false
+            )
+        }
+        finally {
+            Configuration.getInstance().replaceInjections(listOf(), listOf(customInjection), true)
+        }
+    }
+
+    fun testInjectionWithCommentOnProperty() = testInjectionPresent(
             """
             |//language=file-reference
             |val test = "<caret>simple"
             """)
 
-    fun testInjectionWithMultipleCommentsOnFun() = doTestInjection(
+    fun testInjectionWithMultipleCommentsOnFun() = testInjectionPresent(
             """
             |// Some comment
             |// Other comment
@@ -64,7 +95,7 @@ class KotlinInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
             |fun test() = "<caret>simple"
             """)
 
-    fun testInjectionWithAnnotationOnPropertyWithAnnotation() = doTestInjection(
+    fun testInjectionWithAnnotationOnPropertyWithAnnotation() = testInjectionPresent(
             """
             |@org.intellij.lang.annotations.Language("file-reference")
             |val test = "<caret>simple"
@@ -189,6 +220,18 @@ class KotlinInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
 //            |
 //            |fun template(): String = "<caret><html></html>"
 //            """
+//    )
+
+//    fun testRemoveInjectionOnQualifiedNameWithAnnotation() = doRemoveInjectionTest(
+//            """
+//            |import org.intellij.lang.annotations.Language
+//            |
+//            |@Language("RegExp")
+//            |val s = java.util.regex.Pattern.compile("Hi")
+//            """,
+//            """
+//            |val test1 = java.util.regex.Pattern.compile("Hi")
+//            """
 //    )
 
     fun testRemoveInjectionWithComment() = doRemoveInjectionTest(
@@ -243,17 +286,27 @@ class KotlinInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
 
     override fun getProjectDescriptor(): LightProjectDescriptor {
         if (getTestName(true).endsWith("WithAnnotation")) {
-            return PROJECT_DESCRIPTOR_WITH_LANGUAGE_ANNOTATION_LIB
+            return KotlinLightProjectDescriptor.INSTANCE
         }
 
         return JAVA_LATEST
     }
 
-    private fun doTestInjection(text: String) {
+    private fun testInjectionPresent(text: String, languageId: String? = null, unInjectShouldBePresent: Boolean = true) {
         myFixture.configureByText("${getTestName(true)}.kt", text.trimMargin())
 
-        TestCase.assertFalse(InjectLanguageAction().isAvailable(project, myFixture.editor, myFixture.file))
-        TestCase.assertTrue(UnInjectLanguageAction().isAvailable(project, myFixture.editor, myFixture.file))
+        TestCase.assertFalse("Injection action is available. There's probably no injection at caret place",
+                             InjectLanguageAction().isAvailable(project, myFixture.editor, myFixture.file))
+
+        if (languageId != null) {
+            val injectedFile = (editor as? EditorWindow)?.injectedFile
+            assertEquals("Wrong injection language", languageId, injectedFile?.language?.id)
+        }
+
+        if (unInjectShouldBePresent) {
+            TestCase.assertTrue("UnInjection action is not available. There's no injection at caret place or some other troubles.",
+                                UnInjectLanguageAction().isAvailable(project, myFixture.editor, myFixture.file))
+        }
     }
 
     private fun doRemoveInjectionTest(before: String, after: String) {
@@ -285,16 +338,4 @@ class KotlinInjectionTest : KotlinLightCodeInsightFixtureTestCase() {
             configuration.isSourceModificationAllowed = allowed
         }
     }
-}
-
-private val PROJECT_DESCRIPTOR_WITH_LANGUAGE_ANNOTATION_LIB = object : DefaultLightProjectDescriptor() {
-    override fun configureModule(module: Module, model: ModifiableRootModel, contentEntry: ContentEntry) {
-        super.configureModule(module, model, contentEntry)
-
-        val jarUrl = VfsUtil.getUrlForLibraryRoot(File("ideaSDK/lib/annotations.jar"))
-        val libraryModel = model.moduleLibraryTable.modifiableModel.createLibrary("annotations").modifiableModel
-        libraryModel.addRoot(jarUrl, OrderRootType.CLASSES)
-
-        libraryModel.commit()
-    }
-}
+}