Reuse Kotlin indexes in incremental providersupstream-main

Also simplify finding Java sources.

6 files changed, 140 insertions, 216 deletions

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
index c3d9a7ad..bea08c36 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
@@ -50,7 +50,6 @@ import com.intellij.openapi.roots.PackageIndex
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.psi.PsiFileSystemItem
 import com.intellij.psi.PsiJavaFile
 import com.intellij.psi.PsiManager
@@ -125,7 +124,6 @@ import org.jetbrains.kotlin.platform.konan.NativePlatforms
 import org.jetbrains.kotlin.platform.wasm.WasmPlatforms
 import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
-import java.nio.file.Files
 import java.nio.file.Path
 
 class KotlinSymbolProcessing(
@@ -356,11 +354,10 @@ class KotlinSymbolProcessing(
         compilerConfiguration: CompilerConfiguration
     ): List<KSFile> {
         val project = kotlinCoreProjectEnvironment.project
-        val psiManager = PsiManager.getInstance(project)
         val ktFiles = createSourceFilesFromSourceRoots(
             compilerConfiguration, project, compilerConfiguration.kotlinSourceRoots
         ).toSet().toList()
-        val psiFiles = getPsiFilesFromPaths<PsiFileSystemItem>(
+        val allJavaFiles = getPsiFilesFromPaths<PsiJavaFile>(
             project,
             getSourceFilePaths(compilerConfiguration, includeDirectoryRoot = true)
         )
@@ -378,21 +375,47 @@ class KotlinSymbolProcessing(
             ).update(ktFiles)
 
         // Update Java providers for newly generated source files.
-        reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, psiFiles)
-
-        val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
-        val javaFiles = if (kspConfig is KSPJvmConfig) {
-            val javaRoots = kspConfig.javaSourceRoots + kspConfig.javaOutputDir
-            // Get non-symbolic paths first
-            javaRoots.sortedBy { Files.isSymbolicLink(it.toPath()) }
-                .flatMap { root -> root.walk().filter { it.isFile && it.extension == "java" }.toList() }
-                // This time is for .java files
-                .sortedBy { Files.isSymbolicLink(it.toPath()) }
-                .distinctBy { it.canonicalPath }
-                .mapNotNull { localFileSystem.findFileByPath(it.path)?.let { psiManager.findFile(it) } as? PsiJavaFile }
-        } else {
-            emptyList()
-        }
+        reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, allJavaFiles)
+
+        return ktFiles.map { analyze { KSFileImpl.getCached(it.getFileSymbol()) } } +
+            allJavaFiles.map { KSFileJavaImpl.getCached(it) }
+    }
+
+    private fun prepareNewKSFiles(
+        kotlinCoreProjectEnvironment: KotlinCoreProjectEnvironment,
+        modules: List<KtModule>,
+        compilerConfiguration: CompilerConfiguration,
+        newKotlinFiles: List<File>,
+        newJavaFiles: List<File>,
+    ): List<KSFile> {
+        val project = kotlinCoreProjectEnvironment.project
+        val ktFiles = getPsiFilesFromPaths<KtFile>(
+            project,
+            newKotlinFiles.map { it.toPath() }.toSet()
+        )
+        val javaFiles = getPsiFilesFromPaths<PsiJavaFile>(
+            project,
+            newJavaFiles.map { it.toPath() }.toSet()
+        )
+        val allJavaFiles = getPsiFilesFromPaths<PsiJavaFile>(
+            project,
+            getSourceFilePaths(compilerConfiguration, includeDirectoryRoot = true)
+        )
+
+        // Update Kotlin providers for newly generated source files.
+        (
+            project.getService(
+                KotlinDeclarationProviderFactory::class.java
+            ) as IncrementalKotlinDeclarationProviderFactory
+            ).update(ktFiles)
+        (
+            project.getService(
+                KotlinPackageProviderFactory::class.java
+            ) as IncrementalKotlinPackageProviderFactory
+            ).update(ktFiles)
+
+        // Update Java providers for newly generated source files.
+        reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, allJavaFiles)
 
         return ktFiles.map { analyze { KSFileImpl.getCached(it.getFileSymbol()) } } +
             javaFiles.map { KSFileJavaImpl.getCached(it) }
@@ -533,7 +556,7 @@ class KotlinSymbolProcessing(
             }
 
             // Drop caches
-            KotlinGlobalModificationService.getInstance(project).publishGlobalModuleStateModification()
+            KotlinGlobalModificationService.getInstance(project).publishGlobalSourceModuleStateModification()
             KtAnalysisSessionProvider.getInstance(project).clearCaches()
             psiManager.dropResolveCaches()
             psiManager.dropPsiCaches()
@@ -542,10 +565,26 @@ class KotlinSymbolProcessing(
 
             val newFilePaths = codeGenerator.generatedFile.filter { it.extension == "kt" || it.extension == "java" }
                 .map { it.canonicalPath }.toSet()
-            allDirtyKSFiles = prepareAllKSFiles(kotlinCoreProjectEnvironment, modules, compilerConfiguration).filter {
-                it.filePath !in allCleanFilePaths
-            }
-            newKSFiles = allDirtyKSFiles.filter { it.filePath in newFilePaths }
+            newKSFiles = prepareNewKSFiles(
+                kotlinCoreProjectEnvironment,
+                modules,
+                compilerConfiguration,
+                newFilePaths.filter { it.endsWith(".kt") }.map { File(it) }.toList(),
+                newFilePaths.filter { it.endsWith(".java") }.map { File(it) }.toList(),
+            )
+            // Now that caches are dropped, KtSymbols and KS* are invalid. They need to be re-created from PSI.
+            allDirtyKSFiles = allDirtyKSFiles.map {
+                when (it) {
+                    is KSFileImpl -> {
+                        val ktFile = it.ktFileSymbol.psi!! as KtFile
+                        analyze { KSFileImpl.getCached(ktFile.getFileSymbol()) }
+                    }
+                    is KSFileJavaImpl -> {
+                        KSFileJavaImpl.getCached(it.psi)
+                    }
+                    else -> throw IllegalArgumentException("Unknown KSFile implementation: $it")
+                }
+            } + newKSFiles
             incrementalContext.registerGeneratedFiles(newKSFiles)
             codeGenerator.closeFiles()
         }
@@ -612,7 +651,7 @@ class DirectoriesScope(
 private fun reinitJavaFileManager(
     environment: KotlinCoreProjectEnvironment,
     modules: List<KtModule>,
-    sourceFiles: List<PsiFileSystemItem>,
+    sourceFiles: List<PsiJavaFile>,
 ) {
     val project = environment.project
     val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinDeclarationProvider.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinDeclarationProvider.kt
deleted file mode 100644
index 005209fa..00000000
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinDeclarationProvider.kt
+++ /dev/null
@@ -1,107 +0,0 @@
-package com.google.devtools.ksp.standalone
-
-import com.intellij.openapi.project.Project
-import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
-import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
-import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProviderFactory
-import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticDeclarationProviderFactory
-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.psi.KtClassLikeDeclaration
-import org.jetbrains.kotlin.psi.KtClassOrObject
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.KtScript
-import org.jetbrains.kotlin.psi.KtTypeAlias
-
-class IncrementalKotlinDeclarationProvider(var del: KotlinDeclarationProvider) : KotlinDeclarationProvider() {
-    override val hasSpecificCallablePackageNamesComputation: Boolean
-        get() = del.hasSpecificCallablePackageNamesComputation
-    override val hasSpecificClassifierPackageNamesComputation: Boolean
-        get() = del.hasSpecificClassifierPackageNamesComputation
-
-    override fun findFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
-        return del.findFilesForFacade(facadeFqName)
-    }
-
-    override fun findFilesForFacadeByPackage(packageFqName: FqName): Collection<KtFile> {
-        return del.findFilesForFacadeByPackage(packageFqName)
-    }
-
-    override fun findFilesForScript(scriptFqName: FqName): Collection<KtScript> {
-        return del.findFilesForScript(scriptFqName)
-    }
-
-    override fun findInternalFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
-        return del.findInternalFilesForFacade(facadeFqName)
-    }
-
-    override fun getAllClassesByClassId(classId: ClassId): Collection<KtClassOrObject> {
-        return del.getAllClassesByClassId(classId)
-    }
-
-    override fun getAllTypeAliasesByClassId(classId: ClassId): Collection<KtTypeAlias> {
-        return del.getAllTypeAliasesByClassId(classId)
-    }
-
-    override fun getClassLikeDeclarationByClassId(classId: ClassId): KtClassLikeDeclaration? {
-        return del.getClassLikeDeclarationByClassId(classId)
-    }
-
-    override fun getTopLevelCallableFiles(callableId: CallableId): Collection<KtFile> {
-        return del.getTopLevelCallableFiles(callableId)
-    }
-
-    override fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name> {
-        return del.getTopLevelCallableNamesInPackage(packageFqName)
-    }
-
-    override fun getTopLevelFunctions(callableId: CallableId): Collection<KtNamedFunction> {
-        return del.getTopLevelFunctions(callableId)
-    }
-
-    override fun getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName: FqName): Set<Name> {
-        return del.getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName)
-    }
-
-    override fun getTopLevelProperties(callableId: CallableId): Collection<KtProperty> {
-        return del.getTopLevelProperties(callableId)
-    }
-}
-
-class IncrementalKotlinDeclarationProviderFactory(
-    private val project: Project,
-) : KotlinDeclarationProviderFactory() {
-    var provider: IncrementalKotlinDeclarationProvider? = null
-    private lateinit var scope: GlobalSearchScope
-    private var contextualModule: KtModule? = null
-    private var files: Collection<KtFile> = emptyList()
-    lateinit var staticFactory: KotlinDeclarationProviderFactory
-
-    override fun createDeclarationProvider(
-        scope: GlobalSearchScope,
-        contextualModule: KtModule?
-    ): KotlinDeclarationProvider {
-        this.scope = scope
-        this.contextualModule = contextualModule
-        return IncrementalKotlinDeclarationProvider(createDelegateProvider()).also {
-            provider = it
-        }
-    }
-
-    fun update(files: Collection<KtFile>) {
-        this.files = files
-        this.staticFactory = KotlinStaticDeclarationProviderFactory(project, files)
-        provider?.let {
-            it.del = createDelegateProvider()
-        }
-    }
-
-    private fun createDelegateProvider(): KotlinDeclarationProvider {
-        return staticFactory.createDeclarationProvider(scope, contextualModule)
-    }
-}
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinDeclarationProviderFactory.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinDeclarationProviderFactory.kt
new file mode 100644
index 00000000..4cf534b2
--- /dev/null
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinDeclarationProviderFactory.kt
@@ -0,0 +1,42 @@
+package com.google.devtools.ksp.standalone
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProviderFactory
+import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticDeclarationProviderFactory
+import org.jetbrains.kotlin.analysis.providers.impl.declarationProviders.CompositeKotlinDeclarationProvider
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtTypeAlias
+
+class IncrementalKotlinDeclarationProviderFactory(
+    private val project: Project,
+) : KotlinDeclarationProviderFactory() {
+    private val staticFactories: MutableList<KotlinStaticDeclarationProviderFactory> = mutableListOf()
+
+    override fun createDeclarationProvider(
+        scope: GlobalSearchScope,
+        contextualModule: KtModule?
+    ): KotlinDeclarationProvider {
+        val providers = staticFactories.map { it.createDeclarationProvider(scope, contextualModule) }
+        return CompositeKotlinDeclarationProvider.create(providers)
+    }
+
+    fun update(files: Collection<KtFile>) {
+        val staticFactory = KotlinStaticDeclarationProviderFactory(project, files)
+        staticFactories.add(staticFactory)
+    }
+
+    fun getDirectInheritorCandidates(baseClassName: Name): Set<KtClassOrObject> =
+        staticFactories.flatMapTo(mutableSetOf()) {
+            it.getDirectInheritorCandidates(baseClassName)
+        }
+
+    fun getInheritableTypeAliases(aliasedName: Name): Set<KtTypeAlias> =
+        staticFactories.flatMapTo(mutableSetOf()) {
+            it.getInheritableTypeAliases(aliasedName)
+        }
+}
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinPackageProvider.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinPackageProvider.kt
deleted file mode 100644
index 947354d0..00000000
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinPackageProvider.kt
+++ /dev/null
@@ -1,73 +0,0 @@
-package com.google.devtools.ksp.standalone
-
-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
-import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticPackageProviderFactory
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.platform.TargetPlatform
-import org.jetbrains.kotlin.psi.KtFile
-
-class IncrementalKotlinPackageProvider(var del: KotlinPackageProvider) : KotlinPackageProvider() {
-    override fun doesKotlinOnlyPackageExist(packageFqName: FqName): Boolean {
-        return del.doesKotlinOnlyPackageExist(packageFqName)
-    }
-
-    override fun doesPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
-        return del.doesPackageExist(packageFqName, platform)
-    }
-
-    override fun doesPlatformSpecificPackageExist(packageFqName: FqName, platform: TargetPlatform): Boolean {
-        return del.doesPlatformSpecificPackageExist(packageFqName, platform)
-    }
-
-    override fun getKotlinOnlySubPackagesFqNames(packageFqName: FqName, nameFilter: (Name) -> Boolean): Set<Name> {
-        return del.getKotlinOnlySubPackagesFqNames(packageFqName, nameFilter)
-    }
-
-    override fun getPlatformSpecificSubPackagesFqNames(
-        packageFqName: FqName,
-        platform: TargetPlatform,
-        nameFilter: (Name) -> Boolean
-    ): Set<Name> {
-        return del.getPlatformSpecificSubPackagesFqNames(packageFqName, platform, nameFilter)
-    }
-
-    override fun getSubPackageFqNames(
-        packageFqName: FqName,
-        platform: TargetPlatform,
-        nameFilter: (Name) -> Boolean
-    ): Set<Name> {
-        return del.getSubPackageFqNames(packageFqName, platform, nameFilter)
-    }
-}
-
-class IncrementalKotlinPackageProviderFactory(
-    private val project: Project,
-) : KotlinPackageProviderFactory() {
-    private var provider: IncrementalKotlinPackageProvider? = null
-    private lateinit var scope: GlobalSearchScope
-    private var files: Collection<KtFile> = emptyList()
-    private lateinit var staticFactory: KotlinPackageProviderFactory
-
-    override fun createPackageProvider(searchScope: GlobalSearchScope): KotlinPackageProvider {
-        this.scope = searchScope
-        return IncrementalKotlinPackageProvider(createDelegateProvider()).also {
-            provider = it
-        }
-    }
-
-    fun update(files: Collection<KtFile>) {
-        this.files = files
-        this.staticFactory = KotlinStaticPackageProviderFactory(project, files)
-        provider?.let {
-            it.del = createDelegateProvider()
-        }
-    }
-
-    private fun createDelegateProvider(): KotlinPackageProvider {
-        return staticFactory.createPackageProvider(scope)
-    }
-}
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinPackageProviderFactory.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinPackageProviderFactory.kt
new file mode 100644
index 00000000..81b26586
--- /dev/null
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalKotlinPackageProviderFactory.kt
@@ -0,0 +1,25 @@
+package com.google.devtools.ksp.standalone
+
+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
+import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticPackageProviderFactory
+import org.jetbrains.kotlin.analysis.providers.impl.packageProviders.CompositeKotlinPackageProvider
+import org.jetbrains.kotlin.psi.KtFile
+
+class IncrementalKotlinPackageProviderFactory(
+    private val project: Project,
+) : KotlinPackageProviderFactory() {
+    private val staticFactories: MutableList<KotlinStaticPackageProviderFactory> = mutableListOf()
+
+    override fun createPackageProvider(searchScope: GlobalSearchScope): KotlinPackageProvider {
+        val providers = staticFactories.map { it.createPackageProvider(searchScope) }
+        return CompositeKotlinPackageProvider.create(providers)
+    }
+
+    fun update(files: Collection<KtFile>) {
+        val staticFactory = KotlinStaticPackageProviderFactory(project, files)
+        staticFactories.add(staticFactory)
+    }
+}
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/KspStandaloneDirectInheritorsProvider.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/KspStandaloneDirectInheritorsProvider.kt
index 3f114bfa..c8dec4c9 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/KspStandaloneDirectInheritorsProvider.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/KspStandaloneDirectInheritorsProvider.kt
@@ -24,15 +24,13 @@ import org.jetbrains.kotlin.psi.psiUtil.contains
 // TODO: copied from upstream as a workaround, remove after upstream fixes standalone session builder for KSP.
 @OptIn(LLFirInternals::class, SymbolInternals::class)
 class KspStandaloneDirectInheritorsProvider(private val project: Project) : KotlinDirectInheritorsProvider {
-    private val staticDeclarationProviderFactory by lazy {
-        (
-            (KotlinDeclarationProviderFactory.getInstance(project) as? IncrementalKotlinDeclarationProviderFactory)
-                ?.staticFactory as? KotlinStaticDeclarationProviderFactory
-            ) ?: error(
-            "KotlinStandaloneDirectInheritorsProvider" +
-                "` expects the following declaration provider factory to be" +
-                " registered: `${KotlinStaticDeclarationProviderFactory::class.simpleName}`"
-        )
+    private val declarationProviderFactory by lazy {
+        (KotlinDeclarationProviderFactory.getInstance(project) as? IncrementalKotlinDeclarationProviderFactory)
+            ?: error(
+                "KotlinStandaloneDirectInheritorsProvider" +
+                    "` expects the following declaration provider factory to be" +
+                    " registered: `${KotlinStaticDeclarationProviderFactory::class.simpleName}`"
+            )
     }
 
     override fun getDirectKotlinInheritors(
@@ -45,7 +43,7 @@ class KspStandaloneDirectInheritorsProvider(private val project: Project) : Kotl
         val aliases = mutableSetOf(classId.shortClassName)
         calculateAliases(classId.shortClassName, aliases)
 
-        val possibleInheritors = aliases.flatMap { staticDeclarationProviderFactory.getDirectInheritorCandidates(it) }
+        val possibleInheritors = aliases.flatMap { declarationProviderFactory.getDirectInheritorCandidates(it) }
 
         if (possibleInheritors.isEmpty()) {
             return emptyList()
@@ -71,7 +69,7 @@ class KspStandaloneDirectInheritorsProvider(private val project: Project) : Kotl
     }
 
     private fun calculateAliases(aliasedName: Name, aliases: MutableSet<Name>) {
-        staticDeclarationProviderFactory.getInheritableTypeAliases(aliasedName).forEach { alias ->
+        declarationProviderFactory.getInheritableTypeAliases(aliasedName).forEach { alias ->
             val aliasName = alias.nameAsSafeName
             val isNewAliasName = aliases.add(aliasName)
             if (isNewAliasName) {