KotlinAnnotatedElementsSearcher: cleanup after J2K; correctly handle consumer returning false

idea/src/org/jetbrains/kotlin/idea/search/ideaExtensions/KotlinAnnotatedElementsSearcher.kt

@@ -16,9 +16,7 @@
 
 package org.jetbrains.kotlin.idea.search.ideaExtensions
 
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.Logger
-import com.intellij.openapi.util.Computable
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiModifierListOwner
@@ -26,19 +24,18 @@ import com.intellij.psi.impl.search.AnnotatedElementsSearcher
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.SearchScope
 import com.intellij.psi.search.searches.AnnotatedElementsSearch
-import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.PsiUtilCore
 import com.intellij.util.Processor
 import com.intellij.util.indexing.FileBasedIndex
 import org.jetbrains.kotlin.asJava.LightClassUtil
-import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.stubindex.JetAnnotationsIndex
+import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
-import java.util.ArrayList
 
 public class KotlinAnnotatedElementsSearcher : AnnotatedElementsSearcher() {
 
@@ -54,30 +51,28 @@ public class KotlinAnnotatedElementsSearcher : AnnotatedElementsSearcher() {
         for (elt in getJetAnnotationCandidates(annClass, useScope)) {
             if (notJetAnnotationEntry(elt)) continue
 
-            ApplicationManager.getApplication().runReadAction(object : Runnable {
+            val result = runReadAction(fun(): Boolean {
-                override fun run() {
+                val parentOfType = elt.getStrictParentOfType<JetDeclaration>() ?: return true
-                    val parentOfType = PsiTreeUtil.getParentOfType<JetDeclaration>(elt, javaClass<JetDeclaration>()) ?: return
 
-                    val annotationEntry = elt as JetAnnotationEntry
+                val annotationEntry = elt as JetAnnotationEntry
 
-                    val context = annotationEntry.analyze(BodyResolveMode.PARTIAL)
+                val context = annotationEntry.analyze(BodyResolveMode.PARTIAL)
-                    val annotationDescriptor = context.get<JetAnnotationEntry, AnnotationDescriptor>(BindingContext.ANNOTATION, annotationEntry)
+                val annotationDescriptor = context.get(BindingContext.ANNOTATION, annotationEntry) ?: return true
-                    if (annotationDescriptor == null) return
 
-                    val descriptor = annotationDescriptor.getType().getConstructor().getDeclarationDescriptor()
+                val descriptor = annotationDescriptor.getType().getConstructor().getDeclarationDescriptor() ?: return true
-                    if (descriptor == null) return
+                if (!(DescriptorUtils.getFqName(descriptor).asString() == annotationFQN)) return true
-                    if (!(DescriptorUtils.getFqName(descriptor).asString() == annotationFQN)) return
 
-                    if (parentOfType is JetClass) {
+                if (parentOfType is JetClass) {
-                        val lightClass = LightClassUtil.getPsiClass(parentOfType as JetClass?)
+                    val lightClass = LightClassUtil.getPsiClass(parentOfType as JetClass?)
-                        consumer.process(lightClass)
+                    if (!consumer.process(lightClass)) return false
-                    }
-                    else if (parentOfType is JetNamedFunction || parentOfType is JetSecondaryConstructor) {
-                        val wrappedMethod = LightClassUtil.getLightClassMethod(parentOfType as JetFunction)
-                        consumer.process(wrappedMethod)
-                    }
                 }
+                else if (parentOfType is JetNamedFunction || parentOfType is JetSecondaryConstructor) {
+                    val wrappedMethod = LightClassUtil.getLightClassMethod(parentOfType as JetFunction)
+                    if (!consumer.process(wrappedMethod)) return false
+                }
+                return true
             })
+            if (!result) return false
         }
 
         return true
@@ -88,23 +83,21 @@ public class KotlinAnnotatedElementsSearcher : AnnotatedElementsSearcher() {
 
         /* Return all elements annotated with given annotation name. Aliases don't work now. */
         private fun getJetAnnotationCandidates(annClass: PsiClass, useScope: SearchScope): Collection<PsiElement> {
-            return ApplicationManager.getApplication().runReadAction(object : Computable<Collection<PsiElement>> {
+            return runReadAction(fun(): Collection<PsiElement> {
-                override fun compute(): Collection<PsiElement> {
+                if (useScope is GlobalSearchScope) {
-                    if (useScope is GlobalSearchScope) {
+                    val name = annClass.getName() ?: return emptyList()
-                        val name = annClass.getName() ?: return emptyList()
+                    val annotationEntries = JetAnnotationsIndex.getInstance().get(name, annClass.getProject(), useScope)
-                        val annotationEntries = JetAnnotationsIndex.getInstance().get(name, annClass.getProject(), useScope)
 
-                        // Add annotations 'test' as often used alias when we search Test annotation
+                    // Add annotations 'test' as often used alias when we search Test annotation
-                        if (name == "Test") {
+                    if (name == "Test") {
-                            annotationEntries.addAll(JetAnnotationsIndex.getInstance().get(name.toLowerCase(), annClass.getProject(), useScope))
+                        annotationEntries.addAll(JetAnnotationsIndex.getInstance().get(name.toLowerCase(), annClass.getProject(), useScope))
-                        }
-
-                        return annotationEntries
                     }
 
-                    // TODO getJetAnnotationCandidates works only with global search scope
+                    return annotationEntries
-                    return ArrayList()
                 }
+
+                // TODO getJetAnnotationCandidates works only with global search scope
+                return emptyList()
             })
         }