K2 UAST: simplify PSI declaration provider

In addition to class lookup (done at cec299ac), we can use
JavaFileManager to search classes in a certain package too.

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

@@ -1,155 +0,0 @@
-/*
- * Copyright 2010-2022 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.ide.highlighter.JavaClassFileType
-import com.intellij.openapi.vfs.StandardFileSystems
-import com.intellij.openapi.vfs.VfsUtilCore
-import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.VirtualFileSystem
-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,
-    ): Set<VirtualFile> {
-        val fqNameString = fqName.asString()
-        val fs = StandardFileSystems.local()
-        return binaryModule.getBinaryRoots().flatMap r@{ rootPath ->
-            val root = findRoot(rootPath, fs) ?: return@r emptySet()
-            val files = mutableSetOf<VirtualFile>()
-            VfsUtilCore.iterateChildrenRecursively(
-                root,
-                /*filter=*/filter@{
-                    // Return `false` will skip the children.
-                    if (it == root) return@filter true
-                    // If it is a directory, then check if its path starts with fq name of interest
-                    val relativeFqName = relativeFqName(root, it)
-                    if (it.isDirectory && fqNameString.startsWith(relativeFqName)) {
-                        return@filter true
-                    }
-                    // Otherwise, i.e., if it is a file, we are already in that matched directory (or directory in the middle).
-                    // But, for files at the top-level, double-check if its parent (dir) and fq name of interest match.
-                    if (isPackageName)
-                        relativeFqName(root, it.parent).endsWith(fqNameString)
-                    else // exact class fq name
-                        relativeFqName == fqNameString
-                },
-                /*iterator=*/{
-                    // We reach here after filtering above.
-                    // Directories in the middle, e.g., com/android, can reach too.
-                    if (!it.isDirectory &&
-                        isCompiledFile(it) &&
-                        it in scope
-                    ) {
-                        files.add(it)
-                    }
-                    true
-                }
-            )
-            files
-        }.toSet()
-    }
-
-    private fun findRoot(
-        rootPath: Path,
-        fs: VirtualFileSystem,
-    ): VirtualFile? {
-        return if (Files.isRegularFile(rootPath) && ".jar" in rootPath.toString()) {
-            jarFileSystem.refreshAndFindFileByPath(rootPath.toAbsolutePath().toString() + URLUtil.JAR_SEPARATOR)
-        } else {
-            fs.findFileByPath(rootPath.toAbsolutePath().toString())
-        }
-    }
-
-    private fun relativeFqName(
-        root: VirtualFile,
-        virtualFile: VirtualFile,
-    ): String {
-        return if (root.isDirectory) {
-            val fragments = buildList {
-                var cur = virtualFile
-                while (cur != root) {
-                    add(cur.nameWithoutExtension)
-                    cur = cur.parent
-                }
-            }
-            fragments.reversed().joinToString(".")
-        } else {
-            virtualFile.path.split(URLUtil.JAR_SEPARATOR).lastOrNull()?.replace("/", ".")
-                ?: URLUtil.JAR_SEPARATOR // random string that will bother membership test.
-        }
-    }
-
-    private fun isCompiledFile(
-        virtualFile: VirtualFile,
-    ): Boolean {
-        return virtualFile.extension?.endsWith(JavaClassFileType.INSTANCE.defaultExtension) == true
-    }
-}

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

@@ -13,7 +13,6 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
 import com.intellij.psi.PsiFile
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.FirStandaloneServiceRegistrar
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.KtStaticProjectStructureProvider
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.LLFirStandaloneLibrarySymbolProviderFactory
@@ -23,7 +22,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLi
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.builder.KtModuleProviderBuilder
 import org.jetbrains.kotlin.analysis.project.structure.builder.buildProjectStructureProvider
-import org.jetbrains.kotlin.analysis.project.structure.impl.KtStandaloneProjectStructureProvider
 import org.jetbrains.kotlin.analysis.project.structure.impl.KtSourceModuleImpl
 import org.jetbrains.kotlin.analysis.project.structure.impl.buildKtModuleProviderByCompilerConfiguration
 import org.jetbrains.kotlin.analysis.project.structure.impl.getPsiFilesFromPaths
@@ -108,7 +106,6 @@ public class StandaloneAnalysisAPISessionBuilder(
         extensionDescriptor.registerExtensionPoint(project)
     }
 
-    @OptIn(KtAnalysisApiInternals::class)
     private fun registerProjectServices(
         sourceKtFiles: List<KtFile>,
         packagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
@@ -161,15 +158,10 @@ public class StandaloneAnalysisAPISessionBuilder(
     }
 
     private fun registerPsiDeclarationFromBinaryModuleProvider() {
-        val standaloneProjectStructureProvider = projectStructureProvider as KtStandaloneProjectStructureProvider
         kotlinCoreProjectEnvironment.project.apply {
             registerService(
                 KotlinPsiDeclarationProviderFactory::class.java,
-                KotlinStaticPsiDeclarationProviderFactory(
+                KotlinStaticPsiDeclarationProviderFactory::class.java
-                    this,
-                    standaloneProjectStructureProvider.binaryModules,
-                    kotlinCoreProjectEnvironment.environment.jarFileSystem as CoreJarFileSystem
-                )
             )
         }
     }

analysis/analysis-api-standalone/tests/org/jetbrains/kotlin/analysis/api/standalone/fir/test/configurators/StandaloneModeTestServiceRegistrar.kt

@@ -6,18 +6,13 @@
 package org.jetbrains.kotlin.analysis.api.standalone.fir.test.configurators
 
 import com.intellij.mock.MockProject
-import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.standalone.KtAlwaysAccessibleLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.LLFirStandaloneLibrarySymbolProviderFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLibrarySymbolProviderFactory
-import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
-import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPsiDeclarationProviderFactory
 import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticPsiDeclarationProviderFactory
-import org.jetbrains.kotlin.analysis.test.framework.services.KtTestProjectStructureProvider
-import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestServiceRegistrar
 import org.jetbrains.kotlin.test.services.TestServices
 
@@ -31,19 +26,10 @@ public object StandaloneModeTestServiceRegistrar : AnalysisApiTestServiceRegistr
     }
 
     override fun registerProjectModelServices(project: MockProject, testServices: TestServices) {
-        val projectStructureProvider = ProjectStructureProvider.getInstance(project)
-        val binaryModules =
-            (projectStructureProvider as? KtTestProjectStructureProvider)?.allKtModules?.filterIsInstance<KtBinaryModule>()
-                ?: emptyList()
-        val projectEnvironment = testServices.environmentManager.getProjectEnvironment()
         project.apply {
             registerService(
                 KotlinPsiDeclarationProviderFactory::class.java,
-                KotlinStaticPsiDeclarationProviderFactory(
+                KotlinStaticPsiDeclarationProviderFactory::class.java
-                    this,
-                    binaryModules,
-                    projectEnvironment.environment.jarFileSystem as CoreJarFileSystem
-                )
             )
         }
     }

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

@@ -6,75 +6,49 @@
 package org.jetbrains.kotlin.analysis.providers.impl
 
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
 import com.intellij.psi.*
-import com.intellij.psi.impl.compiled.ClsClassImpl
-import com.intellij.psi.impl.compiled.ClsFileImpl
 import com.intellij.psi.impl.file.impl.JavaFileManager
 import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.util.containers.ContainerUtil
-import org.jetbrains.kotlin.analysis.decompiled.light.classes.ClsJavaStubByVirtualFileCache
-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.analysis.providers.createPackagePartProvider
-import org.jetbrains.kotlin.asJava.builder.ClsWrapperStubPsiFactory
 import org.jetbrains.kotlin.asJava.classes.lazyPub
 import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
 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
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.jvm.KotlinCliJavaFileManager
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.decapitalizeSmart
 import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.ConcurrentMap
 
 private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     private val project: Project,
-    override val scope: GlobalSearchScope,
+    val scope: GlobalSearchScope,
-    override val packagePartProvider: PackagePartProvider,
+    private val packagePartProvider: PackagePartProvider,
-    private val binaryModules: Collection<KtBinaryModule>,
+) : KotlinPsiDeclarationProvider() {
-    override val jarFileSystem: CoreJarFileSystem,
-) : KotlinPsiDeclarationProvider(), AbstractDeclarationFromBinaryModuleProvider {
-    private val psiManager by lazyPub { PsiManager.getInstance(project) }
 
     private val javaFileManager by lazyPub { project.getService(JavaFileManager::class.java) }
 
-    private val virtualFileCache = ContainerUtil.createConcurrentSoftMap<KtBinaryModule, ConcurrentMap<FqName, Set<VirtualFile>>>()
+    private val classesInPackageCache = ConcurrentHashMap<FqName, Collection<PsiClass>>()
 
-    private fun clsClassImplsInPackage(
+    private fun getClassesInPackage(fqName: FqName): Collection<PsiClass> {
-        fqName: FqName,
+        return classesInPackageCache.getOrPut(fqName) {
-    ): Collection<ClsClassImpl> {
+            // `javaFileManager.findPackage(fqName).classes` triggers reading decompiled text from stub for built-in,
-        return binaryModules
+            // which will fail since such stubs are fake, i.e., no mirror to render decompiled text.
-            .flatMap { binaryModule ->
+            // Instead, we will find/use potential class names in the package, while considering package parts.
-                val mapPerModule = virtualFileCache.getOrPut(binaryModule) { ConcurrentHashMap() }
+            val packageParts =
-                mapPerModule.getOrPut(fqName) {
+                packagePartProvider.findPackageParts(fqName.asString()).map { it.replace("/", ".") }
-                    val virtualFilesFromKotlinModule = virtualFilesFromKotlinModule(binaryModule, fqName)
+            val fqNames = packageParts.ifEmpty {
-                    // NB: this assumes Kotlin module has a valid `kotlin_module` info,
+                (javaFileManager as? KotlinCliJavaFileManager)?.knownClassNamesInPackage(fqName)?.map { name ->
-                    // i.e., package part info for the given `fqName` points to exact class paths we're looking for,
+                    fqName.child(Name.identifier(name)).asString()
-                    // and thus it's redundant to walk through the folders in an exhaustive way.
-                    virtualFilesFromKotlinModule.ifEmpty { virtualFilesFromModule(binaryModule, fqName, isPackageName = true) }
                 }
-            }
+            } ?: return@getOrPut emptyList()
-            .distinct()
+            fqNames.flatMap { fqName ->
-            .mapNotNull {
+                javaFileManager.findClasses(fqName, scope).asIterable()
-                createClsJavaClassFromVirtualFile(it)
+            }.distinct()
-            }
-    }
-
-    private fun createClsJavaClassFromVirtualFile(
-        classFile: VirtualFile,
-    ): ClsClassImpl? {
-        val javaFileStub = ClsJavaStubByVirtualFileCache.getInstance(project).get(classFile) ?: return null
-        javaFileStub.psiFactory = ClsWrapperStubPsiFactory.INSTANCE
-        val fakeFile = object : ClsFileImpl(ClassFileViewProvider(psiManager, classFile)) {
-            override fun getStub() = javaFileStub
-
-            override fun isPhysical() = false
         }
-        javaFileStub.psi = fakeFile
-        return fakeFile.classes.single() as ClsClassImpl
     }
 
     override fun getClassesByClassId(classId: ClassId): Collection<PsiClass> {
@@ -97,7 +71,7 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
             // property in companion object is actually materialized at the containing class.
             val classFromOuterClassID = classId.outerClassId?.let { getClassesByClassId(it) } ?: emptyList()
             classFromCurrentClassId + classFromOuterClassID
-        } ?: clsClassImplsInPackage(callableId.packageName)
+        } ?: getClassesInPackage(callableId.packageName)
         return classes.flatMap { psiClass ->
             psiClass.children
                 .filterIsInstance<PsiMember>()
@@ -127,7 +101,7 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     override fun getFunctions(callableId: CallableId): Collection<PsiMethod> {
         val classes = callableId.classId?.let { classId ->
             getClassesByClassId(classId)
-        } ?: clsClassImplsInPackage(callableId.packageName)
+        } ?: getClassesInPackage(callableId.packageName)
         return classes.flatMap { psiClass ->
             psiClass.methods.filter { psiMethod ->
                 psiMethod.name == callableId.callableName.identifier
@@ -136,23 +110,17 @@ private class KotlinStaticPsiDeclarationFromBinaryModuleProvider(
     }
 }
 
-// TODO: we can't register this in IDE yet due to non-trivial parameters: lib modules and jar file system.
-//  We need a session or facade that maintains such information
 class KotlinStaticPsiDeclarationProviderFactory(
     private val project: Project,
-    private val binaryModules: Collection<KtBinaryModule>,
-    private val jarFileSystem: CoreJarFileSystem,
 ) : KotlinPsiDeclarationProviderFactory() {
     // TODO: For now, [createPsiDeclarationProvider] is always called with the project scope, hence singleton.
     //  If we come up with a better / optimal search scope, we may need a different way to cache scope-to-provider mapping.
-    private val provider: KotlinStaticPsiDeclarationFromBinaryModuleProvider by lazy {
+    private val provider: KotlinStaticPsiDeclarationFromBinaryModuleProvider by lazyPub {
         val searchScope = GlobalSearchScope.allScope(project)
         KotlinStaticPsiDeclarationFromBinaryModuleProvider(
             project,
             searchScope,
             project.createPackagePartProvider(searchScope),
-            binaryModules,
-            jarFileSystem,
         )
     }
 
@@ -164,8 +132,6 @@ class KotlinStaticPsiDeclarationProviderFactory(
                 project,
                 searchScope,
                 project.createPackagePartProvider(searchScope),
-                binaryModules,
-                jarFileSystem,
             )
         }
     }