Optimize top level class searching in LazyJavaPackageScope

Do not try to search something, that is known not to exist

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

@@ -20,9 +20,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.util.containers.IntArrayList
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
-import java.util.ArrayList
-import java.util.EnumSet
-import java.util.HashMap
+import java.util.*
 
 data class JavaRoot(val file: VirtualFile, val type: JavaRoot.RootType, val prefixFqName: FqName? = null) {
     enum class RootType {
@@ -99,6 +97,24 @@ class JvmDependenciesIndex(_roots: List<JavaRoot>) {
         }
     }
 
+    fun collectKnownClassNamesInPackage(
+            packageFqName: FqName
+    ): Set<String> {
+        var result = hashSetOf<String>()
+        traverseDirectoriesInPackage(packageFqName, continueSearch = {
+            dir, rootType ->
+
+            for (child in dir.children) {
+                if (child.extension != "class" && child.extension != "java") continue
+                result.add(child.nameWithoutExtension)
+            }
+
+            true
+        })
+
+        return result
+    }
+
     private data class HandleResult<T : Any>(val result: T?, val continueSearch: Boolean)
 
     private fun <T : Any> search(

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

@@ -120,6 +120,8 @@ class KotlinCliJavaFileManagerImpl(private val myPsiManager: PsiManager)
         return findClassInPsiFile(classNameWithInnerClasses, file)
     }
 
+    override fun knownClassNamesInPackage(packageFqName: FqName) = index.collectKnownClassNamesInPackage(packageFqName)
+
     companion object {
         private val LOG = Logger.getInstance(KotlinCliJavaFileManagerImpl::class.java)
 

compiler/frontend.java/src/org/jetbrains/kotlin/load/java/JavaClassFinderImpl.java

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade;
 
 import javax.annotation.PostConstruct;
 import javax.inject.Inject;
+import java.util.Set;
 
 public class JavaClassFinderImpl implements JavaClassFinder {
     private Project project;
@@ -108,4 +109,10 @@ public class JavaClassFinderImpl implements JavaClassFinder {
         PsiPackage psiPackage = javaFacade.findPackage(fqName.asString(), javaSearchScope);
         return psiPackage == null ? null : new JavaPackageImpl(psiPackage, javaSearchScope);
     }
+
+    @Nullable
+    @Override
+    public Set<String> knownClassNamesInPackage(@NotNull FqName packageFqName) {
+        return javaFacade.knownClassNamesInPackage(packageFqName);
+    }
 }

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/KotlinCliJavaFileManager.kt

@@ -20,7 +20,10 @@ import com.intellij.psi.PsiClass
 import com.intellij.psi.impl.file.impl.JavaFileManager
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 
 interface KotlinCliJavaFileManager : JavaFileManager {
     fun findClass(classId: ClassId, searchScope: GlobalSearchScope): PsiClass?
+
+    fun knownClassNamesInPackage(packageFqName: FqName): Set<String>?
 }

compiler/frontend.java/src/org/jetbrains/kotlin/resolve/jvm/KotlinJavaPsiFacade.java

@@ -40,12 +40,15 @@ import kotlin.collections.ArraysKt;
 import kotlin.collections.CollectionsKt;
 import kotlin.jvm.functions.Function1;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus;
 import org.jetbrains.kotlin.name.ClassId;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 
 public class KotlinJavaPsiFacade {
@@ -112,6 +115,17 @@ public class KotlinJavaPsiFacade {
         return null;
     }
 
+    @Nullable
+    public Set<String> knownClassNamesInPackage(@NotNull FqName packageFqName) {
+        KotlinPsiElementFinderWrapper[] finders = finders();
+
+        if (finders.length == 1) {
+            return ((KotlinPsiElementFinderImpl) finders[0]).knownClassNamesInPackage(packageFqName);
+        }
+
+        return null;
+    }
+
     @NotNull
     private PsiClass[] findClassesInDumbMode(@NotNull String qualifiedName, @NotNull GlobalSearchScope scope) {
         String packageName = StringUtil.getPackageName(qualifiedName);
@@ -328,6 +342,15 @@ public class KotlinJavaPsiFacade {
             return findClass(classId.asSingleFqName().asString(), scope);
         }
 
+        @Nullable
+        public Set<String> knownClassNamesInPackage(@NotNull FqName packageFqName) {
+            if (isCliFileManager) {
+                return ((KotlinCliJavaFileManager) javaFileManager).knownClassNamesInPackage(packageFqName);
+            }
+
+            return null;
+        }
+
         @Override
         public PsiPackage findPackage(@NotNull String qualifiedName, @NotNull GlobalSearchScope scope) {
             if (isCliFileManager) {

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/JavaClassFinder.java

@@ -18,15 +18,22 @@ package org.jetbrains.kotlin.load.java;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.ReadOnly;
 import org.jetbrains.kotlin.load.java.structure.JavaClass;
 import org.jetbrains.kotlin.load.java.structure.JavaPackage;
 import org.jetbrains.kotlin.name.ClassId;
 import org.jetbrains.kotlin.name.FqName;
 
+import java.util.Set;
+
 public interface JavaClassFinder {
     @Nullable
     JavaClass findClass(@NotNull ClassId classId);
 
     @Nullable
     JavaPackage findPackage(@NotNull FqName fqName);
+
+    @ReadOnly
+    @Nullable
+    Set<String> knownClassNamesInPackage(@NotNull FqName packageFqName);
 }

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/lazy/descriptors/LazyJavaPackageScope.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
+import org.jetbrains.kotlin.storage.NullableLazyValue
 import org.jetbrains.kotlin.storage.getValue
 
 class LazyJavaPackageScope(
@@ -38,6 +39,12 @@ class LazyJavaPackageScope(
         override val ownerDescriptor: LazyJavaPackageFragment
 ) : LazyJavaStaticScope(c) {
 
+    // Null means that it's impossible to determine list of class names in package, i.e. in IDE where special finders exist
+    // But for compiler though we can determine full list of class names by getting all class-file names in classpath and sources
+    private val knownClassNamesInPackage: NullableLazyValue<Set<String>> = c.storageManager.createNullableLazyValue {
+        c.components.finder.knownClassNamesInPackage(ownerDescriptor.fqName)
+    }
+
     private val partToFacade = c.storageManager.createLazyValue {
         val result = hashMapOf<String, String>()
         kotlinClasses@for (kotlinClass in ownerDescriptor.kotlinBinaryClasses) {
@@ -88,6 +95,12 @@ class LazyJavaPackageScope(
         if (!SpecialNames.isSafeIdentifier(name)) return null
 
         recordLookup(name, location)
+
+        val knownClassNamesInPackage = knownClassNamesInPackage()
+        if (knownClassNamesInPackage != null && name.asString() !in knownClassNamesInPackage) {
+            return null
+        }
+
         return classes(name)
     }
 

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/kotlin/KotlinClassFinder.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.load.kotlin
 
 import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 
 interface KotlinClassFinder {
     fun findKotlinClass(classId: ClassId): KotlinJvmBinaryClass?

core/descriptors.runtime/src/org/jetbrains/kotlin/load/java/reflect/ReflectJavaClassFinder.kt

@@ -40,6 +40,8 @@ class ReflectJavaClassFinder(private val classLoader: ClassLoader) : JavaClassFi
         // We don't know which packages our class loader has and has not, so we behave as if it contains any package in the world
         return ReflectJavaPackage(fqName)
     }
+
+    override fun knownClassNamesInPackage(packageFqName: FqName): Set<String>? = null
 }
 
 fun ClassLoader.tryLoadClass(fqName: String) =