searcher for more efficient search of Kotlin classes which have methods annotated with @Test

#KT-8557 Fixed

idea/src/META-INF/plugin.xml

@@ -515,6 +515,7 @@
     <overridingMethodsSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinOverridingMethodsWithGenericsSearcher"/>
     <definitionsScopedSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinDefinitionsSearcher"/>
     <annotatedElementsSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinAnnotatedElementsSearcher"/>
+    <classesWithAnnotatedMembersSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinClassesWithAnnotatedMembersSearcher"/>
     <methodReferencesSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinPropertyAccessorsReferenceSearcher"/>
     <methodReferencesSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinConstructorDelegationCallReferenceSearcher"/>
     <methodReferencesSearch implementation="org.jetbrains.kotlin.idea.search.ideaExtensions.KotlinOverridingMethodReferenceSearcher"/>

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

@@ -34,7 +34,6 @@ 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.collectDescendantsOfType
-import org.jetbrains.kotlin.psi.psiUtil.getChildrenOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
@@ -43,47 +42,54 @@ import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 public class KotlinAnnotatedElementsSearcher : QueryExecutor<PsiModifierListOwner, AnnotatedElementsSearch.Parameters> {
 
     override fun execute(p: AnnotatedElementsSearch.Parameters, consumer: Processor<PsiModifierListOwner>): Boolean {
-        val annClass = p.getAnnotationClass()
-        assert(annClass.isAnnotationType(), "Annotation type should be passed to annotated members search")
-
-        val annotationFQN = annClass.getQualifiedName()
-        assert(annotationFQN != null)
-
-        val useScope = p.getScope()
-
-        for (elt in getJetAnnotationCandidates(annClass, useScope)) {
-            if (notJetAnnotationEntry(elt)) continue
-
-            val result = runReadAction(fun(): Boolean {
-                val parentOfType = elt.getStrictParentOfType<JetDeclaration>() ?: return true
-
-                val annotationEntry = elt as JetAnnotationEntry
-
-                val context = annotationEntry.analyze(BodyResolveMode.PARTIAL)
-                val annotationDescriptor = context.get(BindingContext.ANNOTATION, annotationEntry) ?: return true
-
-                val descriptor = annotationDescriptor.getType().getConstructor().getDeclarationDescriptor() ?: return true
-                if (!(DescriptorUtils.getFqName(descriptor).asString() == annotationFQN)) return true
-
-                if (parentOfType is JetClass) {
-                    val lightClass = LightClassUtil.getPsiClass(parentOfType as JetClass?)
-                    if (!consumer.process(lightClass)) return false
+        return processAnnotatedMembers(p.annotationClass, p.scope) { declaration ->
+            when (declaration) {
+                is JetClass -> {
+                    val lightClass = LightClassUtil.getPsiClass(declaration)
+                    consumer.process(lightClass)
                 }
-                else if (parentOfType is JetNamedFunction || parentOfType is JetSecondaryConstructor) {
-                    val wrappedMethod = LightClassUtil.getLightClassMethod(parentOfType as JetFunction)
-                    if (!consumer.process(wrappedMethod)) return false
+                is JetNamedFunction, is JetSecondaryConstructor -> {
+                    val wrappedMethod = LightClassUtil.getLightClassMethod(declaration as JetFunction)
+                    consumer.process(wrappedMethod)
                 }
-                return true
-            })
-            if (!result) return false
+                else -> true
+            }
         }
-
-        return true
     }
 
     companion object {
         private val LOG = Logger.getInstance("#com.intellij.psi.impl.search.AnnotatedMembersSearcher")
 
+        public fun processAnnotatedMembers(annClass: PsiClass, useScope: SearchScope, consumer: (JetDeclaration) -> Boolean): Boolean {
+            assert(annClass.isAnnotationType(), "Annotation type should be passed to annotated members search")
+
+            val annotationFQN = annClass.getQualifiedName()
+            assert(annotationFQN != null)
+
+            for (elt in getJetAnnotationCandidates(annClass, useScope)) {
+                if (notJetAnnotationEntry(elt)) continue
+
+                val result = runReadAction(fun(): Boolean {
+                    val declaration = elt.getStrictParentOfType<JetDeclaration>() ?: return true
+
+                    val annotationEntry = elt as JetAnnotationEntry
+
+                    val context = annotationEntry.analyze(BodyResolveMode.PARTIAL)
+                    val annotationDescriptor = context.get(BindingContext.ANNOTATION, annotationEntry) ?: return true
+
+                    val descriptor = annotationDescriptor.getType().getConstructor().getDeclarationDescriptor() ?: return true
+                    if (!(DescriptorUtils.getFqName(descriptor).asString() == annotationFQN)) return true
+
+                    if (!consumer(declaration)) return false
+
+                    return true
+                })
+                if (!result) return false
+            }
+
+            return true
+        }
+
         /* Return all elements annotated with given annotation name. Aliases don't work now. */
         private fun getJetAnnotationCandidates(annClass: PsiClass, useScope: SearchScope): Collection<PsiElement> {
             return runReadAction(fun(): Collection<PsiElement> {

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.search.ideaExtensions
+
+import com.intellij.psi.PsiClass
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.searches.ClassesWithAnnotatedMembersSearch
+import com.intellij.psi.search.searches.ScopedQueryExecutor
+import com.intellij.util.Processor
+import org.jetbrains.kotlin.asJava.LightClassUtil
+import org.jetbrains.kotlin.idea.JetFileType
+import org.jetbrains.kotlin.idea.search.allScope
+import org.jetbrains.kotlin.psi.JetClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+
+class KotlinClassesWithAnnotatedMembersSearcher : ScopedQueryExecutor<PsiClass, ClassesWithAnnotatedMembersSearch.Parameters> {
+    override fun getScope(param: ClassesWithAnnotatedMembersSearch.Parameters): GlobalSearchScope {
+        return GlobalSearchScope.getScopeRestrictedByFileTypes(param.annotationClass.project.allScope(), JetFileType.INSTANCE)
+    }
+
+    override fun execute(queryParameters: ClassesWithAnnotatedMembersSearch.Parameters, consumer: Processor<PsiClass>): Boolean {
+        val processed = hashSetOf<JetClassOrObject>()
+        return KotlinAnnotatedElementsSearcher.processAnnotatedMembers(queryParameters.annotationClass, queryParameters.scope) { declaration ->
+            val jetClass = declaration.getNonStrictParentOfType<JetClassOrObject>()
+            if (jetClass != null && processed.add(jetClass)) {
+                val lightClass = LightClassUtil.getPsiClass(jetClass)
+                if (lightClass != null) consumer.process(lightClass) else true
+            }
+            else
+                true
+        }
+    }
+}