[Analysis API] optimize KotlinPackageProvider.getSubPackageFqNames

Previously, we queried heavy kotlin package provider two times which affected performance
Now it's being queries only a single time

Also, the commit introduces separation for KotlinPackageProvider between kotlin and platform-specific packages

^KTIJ-24640

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtSymbolByFirBuilder.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.api.fir
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.analysis.api.KtStarTypeProjection
@@ -107,9 +106,7 @@ internal class KtSymbolByFirBuilder constructor(
     private val packageProvider = project.createPackageProvider(GlobalSearchScope.allScope(project))//todo scope
 
     fun createPackageSymbolIfOneExists(packageFqName: FqName): KtFirPackageSymbol? {
-        val exists =
-            packageProvider.doKotlinPackageExists(packageFqName)
-                    || JavaPsiFacade.getInstance(project).findPackage(packageFqName.asString()) != null
+        val exists = packageProvider.doesPackageExist(packageFqName, analysisSession.targetPlatform)
         if (!exists) {
             return null
         }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/scopes/KtFirPackageScope.kt

@@ -26,7 +26,6 @@ import org.jetbrains.kotlin.fir.scopes.impl.FirPackageMemberScope
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
-import org.jetbrains.kotlin.platform.jvm.isJvm
 
 internal class KtFirPackageScope(
     private val fqName: FqName,
@@ -111,21 +110,8 @@ internal class KtFirPackageScope(
 
     override fun getPackageSymbols(nameFilter: KtScopeNameFilter): Sequence<KtPackageSymbol> = withValidityAssertion {
         sequence {
-            if (targetPlatform.isJvm()) {
-                val javaPackage = JavaPsiFacade.getInstance(project).findPackage(fqName.asString())
-                if (javaPackage != null) {
-                    for (psiPackage in javaPackage.getSubPackages(searchScope)) {
-                        val fqName = FqName(psiPackage.qualifiedName)
-                        if (nameFilter(fqName.shortName())) {
-                            yield(builder.createPackageSymbol(fqName))
-                        }
-                    }
-                }
-            }
-            packageProvider.getKotlinSubPackageFqNames(fqName).forEach {
-                if (nameFilter(it)) {
-                    yield(builder.createPackageSymbol(fqName.child(it)))
-                }
+            packageProvider.getSubPackageFqNames(fqName, targetPlatform, nameFilter).forEach {
+                yield(builder.createPackageSymbol(fqName.child(it)))
             }
         }
     }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirPackageSymbol.kt

@@ -65,5 +65,5 @@ class KtPackage(
 ) : PsiPackageImpl(manager, fqName.asString().replace('/', '.')) {
     override fun copy() = KtPackage(manager, fqName, scope)
 
-    override fun isValid(): Boolean = project.createPackageProvider(scope).doKotlinPackageExists(fqName)
+    override fun isValid(): Boolean = project.createPackageProvider(scope).doesKotlinOnlyPackageExist(fqName)
 }

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/configurators/AnalysisApiBaseTestServiceRegistrar.kt

@@ -39,7 +39,7 @@ object AnalysisApiBaseTestServiceRegistrar: AnalysisApiTestServiceRegistrar()  {
             registerService(KtModuleScopeProvider::class.java, KtModuleScopeProviderImpl())
             registerService(KotlinAnnotationsResolverFactory::class.java, KotlinStaticAnnotationsResolverFactory(allKtFiles))
             registerService(KotlinDeclarationProviderFactory::class.java, KotlinStaticDeclarationProviderFactory(project, allKtFiles))
-            registerService(KotlinPackageProviderFactory::class.java, KotlinStaticPackageProviderFactory(allKtFiles))
+            registerService(KotlinPackageProviderFactory::class.java, KotlinStaticPackageProviderFactory(project, allKtFiles))
             registerService(KotlinReferenceProvidersService::class.java, HLApiReferenceProviderService::class.java)
             registerService(KotlinResolutionScopeProvider::class.java, KotlinByModulesResolutionScopeProvider::class.java)
         }

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/KotlinPackageProvider.kt

@@ -10,19 +10,67 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.TargetPlatform
 
 /**
- * Provides information about Kotlin packages in the given scope. Can be constructed via [KotlinPackageProviderFactory]
- * [doKotlinPackageExists] is called very often by a FIR compiler. And implementations should consider cache results.
+ * Provides information about packages that are visible to Kotlin in the given scope. Can be constructed via [KotlinPackageProviderFactory].
+ * The FIR compiler calls [doesKotlinOnlyPackageExist]  very often, so the implementations should consider caching the results.
  */
 public abstract class KotlinPackageProvider {
     /**
-     * Checks if a package with given [FqName] exists in current [GlobalSearchScope].
-     * Note, that for Kotlin it is not mandatory for a package to correspond to a directory structure like in Java.
-     * So, a package [FqName] is determined by Kotlin files package directive.
+     * Checks if a package with given [FqName] exists in current [GlobalSearchScope] with a view from a given [platform].
+     *
+     * This includes Kotlin packages as well as platform-specific (i.e., JVM packages) that match the [platform].
+     * Generally, the result is equal to [doesKotlinOnlyPackageExist] || [doesPlatformSpecificPackageExist].
      */
-    public abstract fun doKotlinPackageExists(packageFqName: FqName): Boolean
-    public abstract fun getKotlinSubPackageFqNames(packageFqName: FqName): Set<Name>
+    public abstract fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean
+
+    /**
+     * Checks if a package with a given [FqName] exists in the current [GlobalSearchScope].
+     *
+     * The package should contain Kotlin declarations inside.
+     *
+     * Note that for Kotlin, a package doesn't need to correspond to a directory structure like in Java.
+     * So, a package [FqName] is determined by a Kotlin file package directive.
+     */
+    public abstract fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean
+
+    /**
+     * Checks if a platform-specific (e.g., Java packages for Kotlin/JVM) package with [FqName] exists in the current [GlobalSearchScope].
+     */
+    public abstract fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean
+
+
+    /**
+     * Returns the list of subpackages for a given package, which satisfies [nameFilter].
+     *
+     * The returned sub-package list contains sub-packages visible to Kotlin. (e.g., for Kotlin/JVM, it should include Java packages)
+     *
+     * Generally, the result is equal to [getKotlinOnlySubPackagesFqNames] union with [getPlatformSpecificSubPackagesFqNames]
+     */
+    public abstract fun getSubPackageFqNames(
+        packageFqName: FqName,
+        platform: TargetPlatform,
+        nameFilter: (Name) -> Boolean
+    ): Set<Name>
+
+    /**
+     * Returns the list of subpackages for a given package, which satisfies [nameFilter].
+     *
+     * The returned sub-package list contains all packages with some Kotlin declarations inside.
+     */
+    public abstract fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name>
+
+    /**
+     * Returns the platform-specific (e.g., Java packages for Kotlin/JVM) list of subpackages for a given package, which satisfies [nameFilter].
+     *
+     * The returned sub-package list contains sub-packages visible to Kotlin. (e.g., for Kotlin/JVM, it should include Java packages)
+     */
+    public abstract fun getPlatformSpecificSubPackagesFqNames(
+        packageFqName: FqName,
+        platform: TargetPlatform,
+        nameFilter: (Name) -> Boolean
+    ): Set<Name>
 }
 
 public abstract class KotlinPackageProviderFactory {

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

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.providers.impl
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElementFinder
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.jvm.isJvm
+
+
+public abstract class KotlinPackageProviderBase(
+    protected val
+    project: Project,
+    protected val searchScope: GlobalSearchScope
+) : KotlinPackageProvider() {
+
+    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
+        return doesPlatformSpecificPackageExist(packageFqName, platform) || doesKotlinOnlyPackageExist(packageFqName)
+    }
+
+    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
+        when {
+            platform.isJvm() -> {
+                val fqNameString = packageFqName.asString()
+                forEachNonKotlinPsiElementFinder { finder ->
+                    val psiPackage = finder.findPackage(fqNameString)
+                    if (psiPackage != null) {
+                        // we cannot easily check if some PsiPackage is in GlobalSearchScope or not
+                        return true
+                    }
+                }
+                return false
+            }
+            else -> {
+                // non-JVM platforms are not supported yet
+                return false
+            }
+        }
+    }
+
+    override fun getSubPackageFqNames(packageFqName: FqName, platform: TargetPlatform, nameFilter: (Name) -> Boolean): Set<Name> =
+        buildSet {
+            addAll(getKotlinOnlySubPackagesFqNames(packageFqName, nameFilter))
+            addAll(getPlatformSpecificSubPackagesFqNames(packageFqName, platform, nameFilter))
+        }
+
+
+    override fun getPlatformSpecificSubPackagesFqNames(
+        packageFqName: FqName,
+        platform: TargetPlatform,
+        nameFilter: (Name) -> Boolean
+    ): Set<Name> = when {
+        platform.isJvm() -> {
+            val fqNameString = packageFqName.asString()
+            buildSet {
+                forEachNonKotlinPsiElementFinder { finder ->
+                    val psiPackage = finder.findPackage(fqNameString) ?: return@forEachNonKotlinPsiElementFinder
+                    for (subPackage in finder.getSubPackages(psiPackage, searchScope)) {
+                        val name = subPackage.name?.let(Name::identifierIfValid) ?: continue
+                        if (!nameFilter(name)) continue
+                        add(name)
+                    }
+                }
+            }
+        }
+        else -> {
+            // non-JVM platforms are not supported yet
+            emptySet()
+        }
+    }
+
+    private inline fun forEachNonKotlinPsiElementFinder(action: (PsiElementFinder) -> Unit) {
+        for (finder in PsiElementFinder.EP.getPoint(project).extensionList) {
+            if (finder::class.java.name == KOTLIN_JAVA_ELEMENT_FINDER_CLASS_NAME) {
+                // ignore kotlin package finder
+                // kotlin packages will be handled by KotlinPackageProvider.doesKotlinOnlyPackageExist/getKotlinOnlySubPackagesFqNames
+                continue
+            }
+            action(finder)
+        }
+    }
+
+    private companion object {
+        private const val KOTLIN_JAVA_ELEMENT_FINDER_CLASS_NAME = "org.jetbrains.kotlin.asJava.finder.JavaElementFinder"
+    }
+}
+

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.analysis.providers.impl
 
+import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProviderFactory
@@ -13,10 +14,11 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFile
 
 public class KotlinStaticPackageProvider(
+    project: Project,
     scope: GlobalSearchScope,
     files: Collection<KtFile>
-) : KotlinPackageProvider() {
-    private val packageToSubPackageNames: Map<FqName, Set<Name>> = run {
+) : KotlinPackageProviderBase(project, scope) {
+    private val kotlinPackageToSubPackages: Map<FqName, Set<Name>> = run {
         val filesInScope = files.filter { scope.contains(it.virtualFile) }
         val packages = mutableMapOf<FqName, MutableSet<Name>>()
         filesInScope.forEach { file ->
@@ -30,17 +32,20 @@ public class KotlinStaticPackageProvider(
         packages
     }
 
-    override fun doKotlinPackageExists(packageFqName: FqName): Boolean {
-        return packageFqName in packageToSubPackageNames
+    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean {
+        return packageFqName in kotlinPackageToSubPackages
     }
 
-    override fun getKotlinSubPackageFqNames(packageFqName: FqName): Set<Name> {
-        return packageToSubPackageNames[packageFqName] ?: emptySet()
+    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> {
+        return kotlinPackageToSubPackages[packageFqName]?.filterTo(mutableSetOf()) { nameFilter(it) } ?: emptySet()
     }
 }
 
-public class KotlinStaticPackageProviderFactory(private val files: Collection<KtFile>) : KotlinPackageProviderFactory() {
+public class KotlinStaticPackageProviderFactory(
+    private val project: Project,
+    private val files: Collection<KtFile>
+) : KotlinPackageProviderFactory() {
     override fun createPackageProvider(searchScope: GlobalSearchScope): KotlinPackageProvider {
-        return KotlinStaticPackageProvider(searchScope, files)
+        return KotlinStaticPackageProvider(project, searchScope, files)
     }
 }

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt

@@ -130,7 +130,7 @@ public class StandaloneAnalysisAPISessionBuilder(
                     kotlinCoreProjectEnvironment.environment.jarFileSystem as CoreJarFileSystem
                 )
             )
-            registerService(KotlinPackageProviderFactory::class.java, KotlinStaticPackageProviderFactory(ktFiles))
+            registerService(KotlinPackageProviderFactory::class.java, KotlinStaticPackageProviderFactory(project, ktFiles))
 
             registerService(KtAnalysisSessionProvider::class.java, KtFirAnalysisSessionProvider(this))
             registerService(

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneUtils.kt

@@ -161,7 +161,7 @@ internal fun configureProjectEnvironment(
     )
     project.picoContainer.registerComponentInstance(
         KotlinPackageProviderFactory::class.qualifiedName,
-        KotlinStaticPackageProviderFactory(ktFiles)
+        KotlinStaticPackageProviderFactory(project, ktFiles)
     )
     project.picoContainer.registerComponentInstance(
         PackagePartProviderFactory::class.qualifiedName,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProviderHelper.kt

@@ -92,7 +92,7 @@ internal class LLFirProviderHelper(
 
     fun getPackage(fqName: FqName): FqName? {
         if (!allowKotlinPackage && fqName.isKotlinPackage()) return null
-        return fqName.takeIf(packageProvider::doKotlinPackageExists)
+        return fqName.takeIf(packageProvider::doesKotlinOnlyPackageExist)
     }
 }
 

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/SymbolKotlinAsJavaSupport.kt

@@ -79,11 +79,11 @@ class SymbolKotlinAsJavaSupport(project: Project) : KotlinAsJavaSupportBase<KtMo
             .toSet()
 
     override fun packageExists(fqName: FqName, scope: GlobalSearchScope): Boolean =
-        project.createPackageProvider(scope).doKotlinPackageExists(fqName)
+        project.createPackageProvider(scope).doesKotlinOnlyPackageExist(fqName)
 
     override fun getSubPackages(fqn: FqName, scope: GlobalSearchScope): Collection<FqName> =
         project.createPackageProvider(scope)
-            .getKotlinSubPackageFqNames(fqn)
+            .getKotlinOnlySubPackagesFqNames(fqn, nameFilter = { true })
             .map { fqn.child(it) }
 
     override fun createInstanceOfLightScript(script: KtScript): KtLightClass? = error("Should not be called")

core/compiler.common/src/org/jetbrains/kotlin/name/Name.java

@@ -74,6 +74,11 @@ public final class Name implements Comparable<Name> {
         return true;
     }
 
+    public static @Nullable Name identifierIfValid(@NotNull String name) {
+        if (!isValidIdentifier(name)) return null;
+        return identifier(name);
+    }
+
     @NotNull
     public static Name special(@NotNull String name) {
         if (!name.startsWith("<")) {