[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-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)
     }
 }