Add cache for top-level classes' virtual files

The reason is that searching for a class in a VirtualFile-based
directory (i.e. in a package) happens for O(|amount-of-files-in-dir|),
and it may be rather crucial for packages with a huge amount
of classes

Note that it was checked that the cache should not take a lot of heap
since we call findClass only for really existing classes:
see org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaPackageScope.kt:48

For example, there are no more than 10k classes for the whole
intelliJ project (but the cache is flushed for each module)

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCliJavaFileManagerImpl.kt

@@ -23,6 +23,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.*
 import com.intellij.psi.impl.file.PsiPackageImpl
 import com.intellij.psi.search.GlobalSearchScope
+import gnu.trove.THashMap
 import org.jetbrains.kotlin.cli.jvm.index.JavaRoot
 import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesIndex
 import org.jetbrains.kotlin.name.ClassId
@@ -38,17 +39,23 @@ import kotlin.properties.Delegates
 class KotlinCliJavaFileManagerImpl(private val myPsiManager: PsiManager) : CoreJavaFileManager(myPsiManager), KotlinCliJavaFileManager {
     private val perfCounter = PerformanceCounter.create("Find Java class")
     private var index: JvmDependenciesIndex by Delegates.notNull()
-    private val allScope = GlobalSearchScope.allScope(myPsiManager.project)
+    private val topLevelClassesCache: MutableMap<FqName, VirtualFile?> = THashMap()
 
     fun initIndex(packagesCache: JvmDependenciesIndex) {
         this.index = packagesCache
     }
 
     override fun findClass(classId: ClassId, searchScope: GlobalSearchScope): PsiClass? = perfCounter.time {
+        findVirtualFileForTopLevelClass(classId, searchScope)?.findPsiClassInVirtualFile(classId.relativeClassName.asString())
+    }
+
+    private fun findVirtualFileForTopLevelClass(classId: ClassId, searchScope: GlobalSearchScope): VirtualFile? {
         val relativeClassName = classId.relativeClassName.asString()
-        index.findClass(classId) { dir, type ->
+        return topLevelClassesCache.getOrPut(classId.packageFqName.child(classId.relativeClassName.pathSegments().first())) {
-            findClassGivenPackage(allScope, dir, relativeClassName, type)
+            index.findClass(classId) { dir, type ->
-        }?.takeIf { it.containingFile.virtualFile in searchScope }
+                findVirtualFileGivenPackage(dir, relativeClassName, type)
+            }
+        }?.takeIf { it in searchScope }
     }
 
     // this method is called from IDEA to resolve dependencies in Java code
@@ -89,7 +96,10 @@ class KotlinCliJavaFileManagerImpl(private val myPsiManager: PsiManager) : CoreJ
         forEachClassId(qName) { classId ->
             val relativeClassName = classId.relativeClassName.asString()
             index.traverseDirectoriesInPackage(classId.packageFqName) { dir, rootType ->
-                val psiClass = findClassGivenPackage(scope, dir, relativeClassName, rootType)
+                val psiClass =
+                        findVirtualFileGivenPackage(dir, relativeClassName, rootType)
+                            ?.takeIf { it in scope }
+                            ?.findPsiClassInVirtualFile(relativeClassName)
                 if (psiClass != null) {
                     result.add(psiClass)
                 }
@@ -118,10 +128,11 @@ class KotlinCliJavaFileManagerImpl(private val myPsiManager: PsiManager) : CoreJ
         return null
     }
 
-    private fun findClassGivenPackage(
+    private fun findVirtualFileGivenPackage(
-            scope: GlobalSearchScope, packageDir: VirtualFile,
+            packageDir: VirtualFile,
-            classNameWithInnerClasses: String, rootType: JavaRoot.RootType
+            classNameWithInnerClasses: String,
-    ): PsiClass? {
+            rootType: JavaRoot.RootType
+    ): VirtualFile? {
         val topLevelClassName = classNameWithInnerClasses.substringBefore('.')
 
         val vFile = when (rootType) {
@@ -133,11 +144,14 @@ class KotlinCliJavaFileManagerImpl(private val myPsiManager: PsiManager) : CoreJ
             LOG.error("Invalid child of valid parent: ${vFile.path}; ${packageDir.isValid} path=${packageDir.path}")
             return null
         }
-        if (vFile !in scope) {
-            return null
-        }
 
-        val file = myPsiManager.findFile(vFile) as? PsiClassOwner ?: return null
+        return vFile
+    }
+
+    private fun VirtualFile.findPsiClassInVirtualFile(
+            classNameWithInnerClasses: String
+    ): PsiClass? {
+        val file = myPsiManager.findFile(this) as? PsiClassOwner ?: return null
         return findClassInPsiFile(classNameWithInnerClasses, file)
     }