AA: use package part info (if available) when searching for PSI from binary module

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

@@ -14,19 +14,71 @@ import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.io.URLUtil
 import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
+import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.name.FqName
 import java.nio.file.Files
 import java.nio.file.Path
 
 public interface AbstractDeclarationFromBinaryModuleProvider {
     public val scope: GlobalSearchScope
+    public val packagePartProvider: PackagePartProvider
     public val jarFileSystem: CoreJarFileSystem
 
+    /**
+     * Collect [VirtualFile]s that belong to the package with the given [FqName],
+     * from the given [KtBinaryModule], which is supposed to be a Kotlin module (i.e., with `kotlin_module` info),
+     * and properly registered to [PackagePartProvider]. Otherwise, returns an empty set.
+     *
+     * This util is useful to collect files for the package that may have multi-file facades.
+     * E.g., for `kotlin.collection`, regular classes would be under `kotlin/collection` folder.
+     * But, there could be more classes under irregular places, like `.../jdk8/...`,
+     * which would still have `kotlin.collection` as a package, if it is part of multi-file facades.
+     *
+     * To cover such cases with a normal, exhaustive directory lookup used in [virtualFilesFromModule], we will end up
+     * traversing _all_ folders, which is inefficient if package part information is available in `kotlin_module`.
+     */
+    public fun virtualFilesFromKotlinModule(
+        binaryModule: KtBinaryModule,
+        fqName: FqName,
+    ): Set<VirtualFile> {
+        val fqNameString = fqName.asString()
+        val packageParts = packagePartProvider.findPackageParts(fqNameString)
+        return if (packageParts.isNotEmpty()) {
+            binaryModule.getBinaryRoots().flatMap r@{ rootPath ->
+                if (!Files.isRegularFile(rootPath) || ".jar" !in rootPath.toString()) return@r emptySet<VirtualFile>()
+                buildSet {
+                    packageParts.forEach { packagePart ->
+                        add(
+                            jarFileSystem.refreshAndFindFileByPath(
+                                rootPath.toAbsolutePath().toString() + URLUtil.JAR_SEPARATOR + packagePart + ".class"
+                            ) ?: return@r emptySet<VirtualFile>()
+                        )
+                    }
+                }
+            }.toSet()
+        } else
+            emptySet()
+    }
+
+    /**
+     * Collect [VirtualFile]s that belong to the package with the given [FqName],
+     * from the given [KtBinaryModule], which has general `jar` files as roots, e.g., `android.jar` (for a specific API version)
+     *
+     * If the given [FqName] is a specific class name, returns a set with the corresponding [VirtualFile].
+     *
+     * This util assumes that classes will be under the folder where the folder path and package name match.
+     * To avoid exhaustive traversal, this util only visits folders that are parts of the given package name.
+     * E.g., for `android.os`, this will visit `android` and `android/os` directories only,
+     * and will return [VirtualFile]s for all classes under `android/os`.
+     *
+     * For a query with a class name, e.g., `android.os.Bundle`, this will visit `android` and `android/os` directories too,
+     * to search for that specific class.
+     */
     public fun virtualFilesFromModule(
         binaryModule: KtBinaryModule,
         fqName: FqName,
         isPackageName: Boolean,
-    ): Collection<VirtualFile> {
+    ): Set<VirtualFile> {
         val fqNameString = fqName.asString()
         val fs = StandardFileSystems.local()
         return binaryModule.getBinaryRoots().flatMap r@{ rootPath ->
@@ -62,7 +114,7 @@ public interface AbstractDeclarationFromBinaryModuleProvider {
                 }
             )
             files
-        }
+        }.toSet()
     }
 
     private fun findRoot(

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

@@ -9,10 +9,13 @@ import com.intellij.core.CoreApplicationEnvironment
 import com.intellij.mock.MockProject
 import com.intellij.openapi.application.Application
 import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
+import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 
 public class StandaloneAnalysisAPISession internal constructor(
     kotlinCoreProjectEnvironment: KotlinCoreProjectEnvironment,
+    public val createPackagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
 ) {
     // TODO: better to limit exposure? Current usages are: addExtension, jarFileSystem
     public val coreApplicationEnvironment: CoreApplicationEnvironment = kotlinCoreProjectEnvironment.environment

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

@@ -195,16 +195,19 @@ public class StandaloneAnalysisAPISessionBuilder(
         val project = kotlinCoreProjectEnvironment.project
         val ktFiles = allSourceFiles.filterIsInstance<KtFile>()
         val libraryRoots = StandaloneProjectFactory.getAllBinaryRoots(modules, kotlinCoreProjectEnvironment)
-        registerProjectServices(
+        val createPackagePartProvider =
-            ktFiles,
             StandaloneProjectFactory.createPackagePartsProvider(
                 project,
                 libraryRoots,
             )
+        registerProjectServices(
+            ktFiles,
+            createPackagePartProvider,
         )
 
         return StandaloneAnalysisAPISession(
             kotlinCoreProjectEnvironment,
+            createPackagePartProvider,
         )
     }
 }

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
 import org.jetbrains.kotlin.analysis.providers.KotlinPsiDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPsiDeclarationProviderFactory
 import org.jetbrains.kotlin.asJava.builder.ClsWrapperStubPsiFactory
+import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -25,6 +26,7 @@ import org.jetbrains.kotlin.util.capitalizeDecapitalize.decapitalizeSmart
 private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     private val project: Project,
     override val scope: GlobalSearchScope,
+    override val packagePartProvider: PackagePartProvider,
     private val binaryModules: Collection<KtBinaryModule>,
     override val jarFileSystem: CoreJarFileSystem,
 ) : KotlinPsiDeclarationProvider(), AbstractDeclarationFromBinaryModuleProvider {
@@ -36,7 +38,11 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     ): Collection<ClsClassImpl> {
         return binaryModules
             .flatMap {
-                virtualFilesFromModule(it, fqName, isPackageName)
+                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) }
             }
             .mapNotNull {
                 createClsJavaClassFromVirtualFile(it)
@@ -99,10 +105,17 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
 //  We need a session or facade that maintains such information
 class KotlinStaticPsiDeclarationProviderFactory(
     private val project: Project,
+    private val createPackagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
     private val binaryModules: Collection<KtBinaryModule>,
     private val jarFileSystem: CoreJarFileSystem,
 ) : KotlinPsiDeclarationProviderFactory() {
     override fun createPsiDeclarationProvider(searchScope: GlobalSearchScope): KotlinPsiDeclarationProvider {
-        return KotlinStaticPsiDeclarationFromBinaryModuleProvider(project, searchScope, binaryModules, jarFileSystem)
+        return KotlinStaticPsiDeclarationFromBinaryModuleProvider(
+            project,
+            searchScope,
+            createPackagePartProvider.invoke(searchScope),
+            binaryModules,
+            jarFileSystem,
+        )
     }
 }