Rewrite processClassesWithName with processors

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/KotlinShortNamesCache.kt

@@ -24,10 +24,10 @@ import com.intellij.psi.PsiField
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.PsiShortNamesCache
+import com.intellij.psi.stubs.StubIndex
 import com.intellij.util.Processor
 import com.intellij.util.Processors
 import com.intellij.util.containers.ContainerUtil
-import com.intellij.util.containers.HashSet
 import com.intellij.util.indexing.IdFilter
 import org.jetbrains.kotlin.asJava.LightClassUtil
 import org.jetbrains.kotlin.asJava.defaultImplsChild
@@ -40,17 +40,18 @@ import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.load.java.getPropertyNamesCandidatesByAccessorName
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addToStdlib.sequenceOfLazyValues
-import java.util.*
 
 class KotlinShortNamesCache(private val project: Project) : PsiShortNamesCache() {
     /**
      * Return kotlin class names from project sources which should be visible from java.
      */
     override fun getAllClassNames(): Array<String> {
-        val processor = CancelableArrayCollectProcessor()
+        return CancelableArrayCollectProcessor<String>().also { processor ->
-        processAllClassNames(processor)
+            processAllClassNames(processor)
-        return processor.toArray()
+        }.toArray(arrayOf())
     }
 
     override fun processAllClassNames(processor: Processor<String>): Boolean {
@@ -62,18 +63,20 @@ class KotlinShortNamesCache(private val project: Project) : PsiShortNamesCache()
      * Return class names form kotlin sources in given scope which should be visible as Java classes.
      */
     override fun getClassesByName(name: String, scope: GlobalSearchScope): Array<PsiClass> {
+        return CancelableArrayCollectProcessor<PsiClass>().also { processor ->
+            processClassesWithName(name, processor, scope, null)
+        }.toArray(arrayOf())
+    }
+
+    override fun processClassesWithName(
+        name: String,
+        processor: Processor<in PsiClass>,
+        scope: GlobalSearchScope,
+        filter: IdFilter?
+    ): Boolean {
         val effectiveScope = kotlinDeclarationsVisibleFromJavaScope(scope)
-        val allFqNames = HashSet<FqName?>()
+        val fqNameProcessor = Processor<FqName> { fqName: FqName? ->
-
+            if (fqName == null) return@Processor true
-        KotlinClassShortNameIndex.getInstance().get(name, project, effectiveScope)
-            .mapTo(allFqNames) { it.fqName }
-
-        KotlinFileFacadeShortNameIndex.INSTANCE.get(name, project, effectiveScope)
-            .mapTo(allFqNames) { it.javaFileFacadeFqName }
-
-        val result = ArrayList<PsiClass>()
-        for (fqName in allFqNames) {
-            if (fqName == null) continue
 
             val isInterfaceDefaultImpl = name == JvmAbi.DEFAULT_IMPLS_CLASS_NAME && fqName.shortName().asString() != name
             assert(fqName.shortName().asString() == name || isInterfaceDefaultImpl) {
@@ -83,12 +86,35 @@ class KotlinShortNamesCache(private val project: Project) : PsiShortNamesCache()
             val fqNameToSearch = if (isInterfaceDefaultImpl) fqName.defaultImplsChild() else fqName
 
             val psiClass = JavaElementFinder.getInstance(project).findClass(fqNameToSearch.asString(), effectiveScope)
-            if (psiClass != null) {
+                ?: return@Processor true
-                result.add(psiClass)
+
-            }
+            return@Processor processor.process(psiClass)
         }
 
-        return result.toTypedArray()
+        val allKtClassOrObjectsProcessed = StubIndex.getInstance().processElements(
+            KotlinClassShortNameIndex.getInstance().key,
+            name,
+            project,
+            scope,
+            filter,
+            KtClassOrObject::class.java
+        ) { ktClassOrObject ->
+            fqNameProcessor.process(ktClassOrObject.fqName)
+        }
+        if (!allKtClassOrObjectsProcessed) {
+            return false
+        }
+
+        return StubIndex.getInstance().processElements(
+            KotlinFileFacadeShortNameIndex.getInstance().key,
+            name,
+            project,
+            scope,
+            filter,
+            KtFile::class.java
+        ) { ktFile ->
+            fqNameProcessor.process(ktFile.javaFileFacadeFqName)
+        }
     }
 
     private fun kotlinDeclarationsVisibleFromJavaScope(scope: GlobalSearchScope): GlobalSearchScope {
@@ -145,12 +171,16 @@ class KotlinShortNamesCache(private val project: Project) : PsiShortNamesCache()
         ContainerUtil.process(getMethodsByName(name, scope), processor)
 
     override fun getAllMethodNames(): Array<String> {
-        return CancelableArrayCollectProcessor().also { processor ->
+        return CancelableArrayCollectProcessor<String>().also { processor ->
-            processAllMethodNames(processor, null, null)
+            processAllMethodNames(processor)
-        }.toArray()
+        }.toArray(arrayOf())
     }
 
-    override fun processAllMethodNames(processor: Processor<String>, scope: GlobalSearchScope?, filter: IdFilter?): Boolean {
+    override fun processAllMethodNames(processor: Processor<String>, scope: GlobalSearchScope, filter: IdFilter?): Boolean {
+        return processAllMethodNames(processor)
+    }
+
+    private fun processAllMethodNames(processor: Processor<String>): Boolean {
         if (!KotlinFunctionShortNameIndex.getInstance().processAllKeys(project, processor)) {
             return false
         }
@@ -171,23 +201,27 @@ class KotlinShortNamesCache(private val project: Project) : PsiShortNamesCache()
     }
 
     override fun getAllFieldNames(): Array<String> {
-        return CancelableArrayCollectProcessor().also { processor ->
+        return CancelableArrayCollectProcessor<String>().also { processor ->
-            processAllFieldNames(processor, null, null)
+            processAllFieldNames(processor)
-        }.toArray()
+        }.toArray(arrayOf())
     }
 
-    override fun processAllFieldNames(processor: Processor<String>, scope: GlobalSearchScope?, filter: IdFilter?): Boolean {
+    override fun processAllFieldNames(processor: Processor<String>, scope: GlobalSearchScope, filter: IdFilter?): Boolean {
+        return processAllFieldNames(processor)
+    }
+
+    private fun processAllFieldNames(processor: Processor<String>): Boolean {
         return KotlinPropertyShortNameIndex.getInstance().processAllKeys(project, processor)
     }
 
-    private class CancelableArrayCollectProcessor : Processor<String> {
+    private class CancelableArrayCollectProcessor<T> : Processor<T> {
-        val troveSet = ContainerUtil.newTroveSet<String>()
+        val troveSet = ContainerUtil.newTroveSet<T>()
-        private val processor = Processors.cancelableCollectProcessor<String>(troveSet)
+        private val processor = Processors.cancelableCollectProcessor<T>(troveSet)
 
-        override fun process(value: String): Boolean {
+        override fun process(value: T): Boolean {
             return processor.process(value)
         }
 
-        fun toArray(): Array<String> = troveSet.toArray(arrayOf())
+        fun toArray(a: Array<T>): Array<T> = troveSet.toArray(a)
     }
 }