AA: populate index for declarations upfront

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationProvider.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.providers.impl
 
+import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProviderFactory
@@ -14,54 +15,120 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 
 public class KotlinStaticDeclarationProvider(
     scope: GlobalSearchScope,
     ktFiles: Collection<KtFile>
 ) : KotlinDeclarationProvider() {
-    private val filesInScope = ktFiles.filter { scope.contains(it.virtualFile) }
 
-    private fun filesByPackage(packageFqName: FqName) =
+    private val index = Index()
-        filesInScope.asSequence()
+
-            .filter { it.packageFqName == packageFqName }
+    private class Index {
+        val facadeFileMap: MutableMap<FqName, MutableSet<KtFile>> = mutableMapOf()
+        val classMap: MutableMap<FqName, MutableSet<KtClassOrObject>> = mutableMapOf()
+        val typeAliasMap: MutableMap<FqName, MutableSet<KtTypeAlias>> = mutableMapOf()
+        val topLevelFunctionMap: MutableMap<FqName, MutableSet<KtNamedFunction>> = mutableMapOf()
+        val topLevelPropertyMap: MutableMap<FqName, MutableSet<KtProperty>> = mutableMapOf()
+    }
+
+    private class KtDeclarationRecorder(private val index: Index) : KtVisitorVoid() {
+
+        override fun visitElement(element: PsiElement) {
+            element.acceptChildren(this)
+        }
+
+        override fun visitKtFile(file: KtFile) {
+            if (file.hasTopLevelCallables()) {
+                index.facadeFileMap.computeIfAbsent(file.packageFqName) {
+                    mutableSetOf()
+                }.add(file)
+            }
+            super.visitKtFile(file)
+        }
+
+        override fun visitClassOrObject(classOrObject: KtClassOrObject) {
+            classOrObject.getClassId()?.let { classId ->
+                index.classMap.computeIfAbsent(classId.packageFqName) {
+                    mutableSetOf()
+                }.add(classOrObject)
+            }
+            super.visitClassOrObject(classOrObject)
+        }
+
+        override fun visitTypeAlias(typeAlias: KtTypeAlias) {
+            typeAlias.getClassId()?.let { classId ->
+                index.typeAliasMap.computeIfAbsent(classId.packageFqName) {
+                    mutableSetOf()
+                }.add(typeAlias)
+            }
+            super.visitTypeAlias(typeAlias)
+        }
+
+        override fun visitNamedFunction(function: KtNamedFunction) {
+            if (function.isTopLevel) {
+                val packageFqName = (function.parent as KtFile).packageFqName
+                index.topLevelFunctionMap.computeIfAbsent(packageFqName) {
+                    mutableSetOf()
+                }.add(function)
+            }
+            super.visitNamedFunction(function)
+        }
+
+        override fun visitProperty(property: KtProperty) {
+            if (property.isTopLevel) {
+                val packageFqName = (property.parent as KtFile).packageFqName
+                index.topLevelPropertyMap.computeIfAbsent(packageFqName) {
+                    mutableSetOf()
+                }.add(property)
+            }
+            super.visitProperty(property)
+        }
+    }
+
+    init {
+        val recorder = KtDeclarationRecorder(index)
+        ktFiles
+            .filter {
+                scope.contains(it.virtualFile)
+            }
+            .forEach {
+                it.accept(recorder)
+            }
+    }
 
     override fun getClassesByClassId(classId: ClassId): Collection<KtClassOrObject> =
-        filesByPackage(classId.packageFqName).flatMap { file ->
+        index.classMap[classId.packageFqName]
-            file.collectDescendantsOfType<KtClassOrObject> { ktClass ->
+            ?.filter { it.getClassId() == classId }
-                ktClass.getClassId() == classId
+            ?: emptyList()
-            }
-        }.toList()
 
     override fun getTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias> =
-        filesByPackage(classId.packageFqName).flatMap { file ->
+        index.typeAliasMap[classId.packageFqName]
-            file.collectDescendantsOfType<KtTypeAlias> { typeAlias ->
+            ?.filter { it.getClassId() == classId }
-                typeAlias.getClassId() == classId
+            ?: emptyList()
-            }
+
-        }.toList()
+    override fun getClassNamesInPackage(packageFqName: FqName): Set<Name> =
+        index.classMap[packageFqName]
+            ?.mapNotNullTo(mutableSetOf()) { it.nameAsName }
+            ?: emptySet()
 
     override fun getTypeAliasNamesInPackage(packageFqName: FqName): Set<Name> =
-        filesByPackage(packageFqName)
+        index.typeAliasMap[packageFqName]
-            .flatMap { it.declarations }
+            ?.mapNotNullTo(mutableSetOf()) { it.nameAsName }
-            .filterIsInstance<KtTypeAlias>()
+            ?: emptySet()
-            .mapNotNullTo(mutableSetOf()) { it.nameAsName }
 
     override fun getPropertyNamesInPackage(packageFqName: FqName): Set<Name> =
-        filesByPackage(packageFqName)
+        index.topLevelPropertyMap[packageFqName]
-            .flatMap { it.declarations }
+            ?.mapNotNullTo(mutableSetOf()) { it.nameAsName }
-            .filterIsInstance<KtProperty>()
+            ?: emptySet()
-            .mapNotNullTo(mutableSetOf()) { it.nameAsName }
 
     override fun getFunctionsNamesInPackage(packageFqName: FqName): Set<Name> =
-        filesByPackage(packageFqName)
+        index.topLevelFunctionMap[packageFqName]
-            .flatMap { it.declarations }
+            ?.mapNotNullTo(mutableSetOf()) { it.nameAsName }
-            .filterIsInstance<KtNamedFunction>()
+            ?: emptySet()
-            .mapNotNullTo(mutableSetOf()) { it.nameAsName }
 
     override fun getFacadeFilesInPackage(packageFqName: FqName): Collection<KtFile> =
-        filesByPackage(packageFqName)
+        index.facadeFileMap[packageFqName]
-            .filter { file -> file.hasTopLevelCallables() }
+            ?: emptyList()
-            .toSet()
 
     override fun findFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
         if (facadeFqName.shortNameOrSpecial().isSpecial) return emptyList()
@@ -70,29 +137,19 @@ public class KotlinStaticDeclarationProvider(
     }
 
     override fun getTopLevelProperties(callableId: CallableId): Collection<KtProperty> =
-        filesByPackage(callableId.packageName)
+        index.topLevelPropertyMap[callableId.packageName]
-            .flatMap { it.declarations }
+            ?.filter { it.nameAsName == callableId.callableName }
-            .filterIsInstance<KtProperty>()
+            ?: emptyList()
-            .filter { it.nameAsName == callableId.callableName }
-            .toList()
 
     override fun getTopLevelFunctions(callableId: CallableId): Collection<KtNamedFunction> =
-        filesByPackage(callableId.packageName)
+        index.topLevelFunctionMap[callableId.packageName]
-            .flatMap { it.declarations }
+            ?.filter { it.nameAsName == callableId.callableName }
-            .filterIsInstance<KtNamedFunction>()
+            ?: emptyList()
-            .filter { it.nameAsName == callableId.callableName }
-            .toList()
 
-
-    override fun getClassNamesInPackage(packageFqName: FqName): Set<Name> =
-        filesByPackage(packageFqName)
-            .flatMap { it.declarations }
-            .filterIsInstance<KtClassOrObject>()
-            .mapNotNullTo(mutableSetOf()) { it.nameAsName }
 }
 
 public class KotlinStaticDeclarationProviderFactory(private val files: Collection<KtFile>) : KotlinDeclarationProviderFactory() {
     override fun createDeclarationProvider(searchScope: GlobalSearchScope): KotlinDeclarationProvider {
         return KotlinStaticDeclarationProvider(searchScope, files)
     }
-}
+}