From 12df326a75fbe992b9c2e54ab9ac5bd46cc76f66 Mon Sep 17 00:00:00 2001
From: "Olga.Klisho" <Olga.Klisho@jetbrains.com>
Date: Sat, 12 Dec 2020 12:57:34 +0300
Subject: [PATCH] (IDEA-248355) MethodSource detection in Kotlin cheap check

GitOrigin-RevId: 2ef97172d22dbd85bf6f0ccc18e3cfeef3880eff
---
 .../references/JUnitReferenceContributor.java | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/plugins/junit/src/com/intellij/execution/junit/codeInsight/references/JUnitReferenceContributor.java b/plugins/junit/src/com/intellij/execution/junit/codeInsight/references/JUnitReferenceContributor.java
index 64aa0a4e8087..25e655786a20 100644
--- a/plugins/junit/src/com/intellij/execution/junit/codeInsight/references/JUnitReferenceContributor.java
+++ b/plugins/junit/src/com/intellij/execution/junit/codeInsight/references/JUnitReferenceContributor.java
@@ -74,18 +74,29 @@ public class JUnitReferenceContributor extends PsiReferenceContributor {
 
     @Override
     public boolean isAcceptable(Object element, PsiElement context) {
-      if (context != null && context.getParent() instanceof PsiArrayInitializerMemberValue) {
-        context = context.getParent();
-      }
       UElement type = UastContextKt.toUElement(context, UElement.class);
       if (type == null) return false;
-      if (!(type.getUastParent() instanceof UNamedExpression)) return false;
-      UNamedExpression uPair = (UNamedExpression)type.getUastParent();
-      if (!(uPair.getUastParent() instanceof UAnnotation)) return false;
+      UAnnotation annotation = isCheapEnoughToSearch(type);
+      if (annotation == null || !myAnnotation.equals(annotation.getQualifiedName())) return false;
+      UNamedExpression uPair = UastUtils.getParentOfType(type, UNamedExpression.class);
+      if (uPair == null) return false;
       String name = ObjectUtils.notNull(uPair.getName(), PsiAnnotation.DEFAULT_REFERENCED_METHOD_NAME);
-      if (!myParameterName.equals(name)) return false;
-      UAnnotation uAnnotation = (UAnnotation)uPair.getUastParent();
-      return myAnnotation.equals(uAnnotation.getQualifiedName());
+      return myParameterName.equals(name);
+    }
+
+    private static UAnnotation isCheapEnoughToSearch(UElement element) {
+      for (int i = 0; i < 5; i++) {
+        if (element instanceof UDeclarationsExpression) {
+          return null;
+        }
+        if (element instanceof UAnnotation) {
+          return (UAnnotation)element;
+        }
+        if (element.getUastParent() != null) {
+          element = element.getUastParent();
+        } else return null;
+      }
+      return null;
     }
 
     @Override