From fa1d2a891bcde5af0354e0e06176406a44e894c4 Mon Sep 17 00:00:00 2001
From: Nicolay Mitropolsky <Nicolay.Mitropolsky@jetbrains.com>
Date: Wed, 28 Feb 2018 14:01:30 +0300
Subject: [PATCH] KotlinLanguageInjector: `getNameReference` extracted to work
 with annotations

---
 .../idea/injection/KotlinLanguageInjector.kt  | 25 +++++++------------
 1 file changed, 9 insertions(+), 16 deletions(-)

diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
index 769214676f6..dfa68796bf1 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
@@ -248,10 +248,7 @@ class KotlinLanguageInjector(
         val argument = ktHost.parent as? KtValueArgument ?: return null
 
         val callExpression = PsiTreeUtil.getParentOfType(ktHost, KtCallElement::class.java) ?: return null
-        var callee = callExpression.calleeExpression ?: return null
-
-        if (callee is KtConstructorCalleeExpression)
-            callee = callee.constructorReferenceExpression ?: return null
+        val callee = getNameReference(callExpression.calleeExpression) ?: return null
 
         if (isAnalyzeOff()) return null
 
@@ -276,18 +273,19 @@ class KotlinLanguageInjector(
         return null
     }
 
+    private fun getNameReference(callee: KtExpression?): KtNameReferenceExpression? {
+        if (callee is KtConstructorCalleeExpression)
+            return callee.constructorReferenceExpression as? KtNameReferenceExpression
+        return callee as? KtNameReferenceExpression
+    }
+
     private fun injectInAnnotationCall(host: KtElement): InjectionInfo? {
         if (!annotationInjectionsEnabled) return null
         val argument = host.parent as? KtValueArgument ?: return null
         val annotationEntry = argument.parent.parent as? KtCallElement ?: return null
         if (!fastCheckInjectionsExists(annotationEntry)) return null
         val calleeExpression = annotationEntry.calleeExpression ?: return null
-        val calleeReference = when (calleeExpression) {
-            is KtConstructorCalleeExpression -> calleeExpression.constructorReferenceExpression?.mainReference // for top annotations
-            is KtNameReferenceExpression -> calleeExpression.mainReference // for nested annotations
-            else -> return null
-        }
-        val callee = calleeReference?.resolve()
+        val callee = getNameReference(calleeExpression)?.mainReference?.resolve()
         when (callee) {
             is PsiClass -> {
                 val psiClass = callee as? PsiClass ?: return null
@@ -386,12 +384,7 @@ class KotlinLanguageInjector(
                                }, false)
 
     private fun fastCheckInjectionsExists(annotationEntry: KtCallElement): Boolean {
-        val referencedName = when (annotationEntry) {
-            is KtAnnotationEntry -> // for top annotations
-                (annotationEntry.typeReference?.typeElement as? KtUserType)?.referencedName ?: return false
-            else -> // for nested annotations
-                (annotationEntry.calleeExpression as? KtNameReferenceExpression)?.getReferencedName() ?: return false
-        }
+        val referencedName = getNameReference(annotationEntry.calleeExpression)?.getReferencedName() ?: return false
         val annotationShortName = annotationEntry.containingKtFile.aliasImportMap()[referencedName].singleOrNull() ?: referencedName
         return annotationShortName in injectableTargetClassShortNames.value
     }