K2 UAST: Improve performance of PSI declaration provider

When it comes to finding regular class (of `PsiClass`) from the given
`ClassId`, we've reinvented the wheel in a worse way. `JavaFileManager`
is there to find `PsiClass`, and we already build/populate index for
that while initializing AA standalone mode's own project environment.
We can simply load that manager and send a query, that's it.

However, our efforts so far on binary PSI declaration providers are not
entirely useless. `JavaFileManager` can't handle top-level declarations
and search for all callables/classes in a given package in general. So,
we can still use existing logic for non-ClassId queries. We just need to
cache the intermediate results: fq name to `VirtualFile`s.

To fully utilize the aforementioned cache, we also need to cache an
instance of `KotlinPsiDeclarationProvider`.

Combing these three ideas, now running K2 UAST on a very large module is
"on par", e.g.,
K1:
//tools/adt/idea/android:intellij.android.core_lint_test PASSED in 112.8s
  Stats over 7 runs: max = 112.8s, min = 104.3s, avg = 107.9s, dev = 3.6s
K2 (before):
//tools/adt/idea/android:intellij.android.core_lint_test PASSED in 303.0s
  Stats over 7 runs: max = 303.0s, min = 286.8s, avg = 294.4s, dev = 5.0s
K2 (after):
//tools/adt/idea/android:intellij.android.core_lint_test PASSED in 112.8s
  Stats over 7 runs: max = 112.8s, min = 105.7s, avg = 108.9s, dev = 2.4s

analysis/symbol-light-classes/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticPsiDeclarationFromBinaryModuleProvider.kt

@@ -11,7 +11,9 @@ import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
 import com.intellij.psi.*
 import com.intellij.psi.impl.compiled.ClsClassImpl
 import com.intellij.psi.impl.compiled.ClsFileImpl
+import com.intellij.psi.impl.file.impl.JavaFileManager
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.analysis.decompiled.light.classes.ClsJavaStubByVirtualFileCache
 import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
 import org.jetbrains.kotlin.analysis.providers.KotlinPsiDeclarationProvider
@@ -24,6 +26,8 @@ import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.decapitalizeSmart
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.ConcurrentMap
 
 private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     private val project: Project,
@@ -34,17 +38,23 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
 ) : KotlinPsiDeclarationProvider(), AbstractDeclarationFromBinaryModuleProvider {
     private val psiManager by lazyPub { PsiManager.getInstance(project) }
 
-    private fun clsClassImplsByFqName(
+    private val javaFileManager by lazyPub { project.getService(JavaFileManager::class.java) }
+
+    private val virtualFileCache = ContainerUtil.createConcurrentSoftMap<KtBinaryModule, ConcurrentMap<FqName, Set<VirtualFile>>>()
+
+    private fun clsClassImplsInPackage(
         fqName: FqName,
-        isPackageName: Boolean = true,
     ): Collection<ClsClassImpl> {
         return binaryModules
-            .flatMap {
-                val virtualFilesFromKotlinModule = if (isPackageName) virtualFilesFromKotlinModule(it, fqName) else emptySet()
-                // NB: this assumes Kotlin module has a valid `kotlin_module` info,
-                // i.e., package part info for the given `fqName` points to exact class paths we're looking for,
-                // and thus it's redundant to walk through the folders in an exhaustive way.
-                virtualFilesFromKotlinModule.ifEmpty { virtualFilesFromModule(it, fqName, isPackageName) }
+            .flatMap { binaryModule ->
+                val mapPerModule = virtualFileCache.getOrPut(binaryModule) { ConcurrentHashMap() }
+                mapPerModule.getOrPut(fqName) {
+                    val virtualFilesFromKotlinModule = virtualFilesFromKotlinModule(binaryModule, fqName)
+                    // NB: this assumes Kotlin module has a valid `kotlin_module` info,
+                    // i.e., package part info for the given `fqName` points to exact class paths we're looking for,
+                    // and thus it's redundant to walk through the folders in an exhaustive way.
+                    virtualFilesFromKotlinModule.ifEmpty { virtualFilesFromModule(binaryModule, fqName, isPackageName = true) }
+                }
             }
             .mapNotNull {
                 createClsJavaClassFromVirtualFile(it)
@@ -72,14 +82,14 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
                 parentClsClass.innerClasses.find { it.name == innerClassName }
             }
         }
-        return clsClassImplsByFqName(classId.asSingleFqName(), isPackageName = false)
+        return listOfNotNull(javaFileManager.findClass(classId.asFqNameString(), scope))
     }
 
     // TODO(dimonchik0036): support 'is' accessor
     override fun getProperties(callableId: CallableId): Collection<PsiMember> {
         val classes = callableId.classId?.let { classId ->
             getClassesByClassId(classId)
-        } ?: clsClassImplsByFqName(callableId.packageName)
+        } ?: clsClassImplsInPackage(callableId.packageName)
         return classes.flatMap { psiClass ->
             psiClass.children
                 .filterIsInstance<PsiMember>()
@@ -100,7 +110,7 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     override fun getFunctions(callableId: CallableId): Collection<PsiMethod> {
         val classes = callableId.classId?.let { classId ->
             getClassesByClassId(classId)
-        } ?: clsClassImplsByFqName(callableId.packageName)
+        } ?: clsClassImplsInPackage(callableId.packageName)
         return classes.flatMap { psiClass ->
             psiClass.methods.filter { psiMethod ->
                 psiMethod.name == callableId.callableName.identifier
@@ -116,8 +126,11 @@ class KotlinStaticPsiDeclarationProviderFactory(
     private val binaryModules: Collection<KtBinaryModule>,
     private val jarFileSystem: CoreJarFileSystem,
 ) : KotlinPsiDeclarationProviderFactory() {
-    override fun createPsiDeclarationProvider(searchScope: GlobalSearchScope): KotlinPsiDeclarationProvider {
-        return KotlinStaticPsiDeclarationFromBinaryModuleProvider(
+    // TODO: For now, [createPsiDeclarationProvider] is always called with the project scope, hence singleton.
+    //  If we come up with a better / optimal search scope, we may need a different way to cache scope-to-provider mapping.
+    private val provider: KotlinStaticPsiDeclarationFromBinaryModuleProvider by lazy {
+        val searchScope = GlobalSearchScope.allScope(project)
+        KotlinStaticPsiDeclarationFromBinaryModuleProvider(
             project,
             searchScope,
             project.createPackagePartProvider(searchScope),
@@ -125,4 +138,18 @@ class KotlinStaticPsiDeclarationProviderFactory(
             jarFileSystem,
         )
     }
+
+    override fun createPsiDeclarationProvider(searchScope: GlobalSearchScope): KotlinPsiDeclarationProvider {
+        return if (searchScope == provider.scope) {
+            provider
+        } else {
+            KotlinStaticPsiDeclarationFromBinaryModuleProvider(
+                project,
+                searchScope,
+                project.createPackagePartProvider(searchScope),
+                binaryModules,
+                jarFileSystem,
+            )
+        }
+    }
 }