KotlinLanguageInjector: getNameReference extracted to work with annotations

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
+        val callee = getNameReference(callExpression.calleeExpression) ?: return null
-
-        if (callee is KtConstructorCalleeExpression)
-            callee = callee.constructorReferenceExpression ?: 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) {
+        val callee = getNameReference(calleeExpression)?.mainReference?.resolve()
-            is KtConstructorCalleeExpression -> calleeExpression.constructorReferenceExpression?.mainReference // for top annotations
-            is KtNameReferenceExpression -> calleeExpression.mainReference // for nested annotations
-            else -> return null
-        }
-        val callee = calleeReference?.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) {
+        val referencedName = getNameReference(annotationEntry.calleeExpression)?.getReferencedName() ?: return false
-            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 annotationShortName = annotationEntry.containingKtFile.aliasImportMap()[referencedName].singleOrNull() ?: referencedName
         return annotationShortName in injectableTargetClassShortNames.value
     }