Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

KSP2: Reuse Java indexes #1895

Merged
merged 1 commit into from
May 17, 2024
Merged

KSP2: Reuse Java indexes #1895

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 10 additions & 99 deletions kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import com.google.devtools.ksp.impl.symbol.kotlin.KSFileJavaImpl
import com.google.devtools.ksp.impl.symbol.kotlin.Restorable
import com.google.devtools.ksp.impl.symbol.kotlin.analyze
import com.google.devtools.ksp.processing.*
import com.google.devtools.ksp.standalone.IncrementalJavaFileManager
import com.google.devtools.ksp.standalone.IncrementalKotlinDeclarationProviderFactory
import com.google.devtools.ksp.standalone.IncrementalKotlinPackageProviderFactory
import com.google.devtools.ksp.standalone.KspStandaloneDirectInheritorsProvider
Expand All @@ -40,13 +41,10 @@ import com.google.devtools.ksp.symbol.KSFile
import com.google.devtools.ksp.symbol.KSNode
import com.google.devtools.ksp.symbol.Origin
import com.intellij.core.CoreApplicationEnvironment
import com.intellij.core.CorePackageIndex
import com.intellij.ide.highlighter.JavaFileType
import com.intellij.mock.MockProject
import com.intellij.openapi.Disposable
import com.intellij.openapi.application.Application
import com.intellij.openapi.project.Project
import com.intellij.openapi.roots.PackageIndex
import com.intellij.openapi.util.Disposer
import com.intellij.openapi.vfs.StandardFileSystems
import com.intellij.openapi.vfs.VirtualFile
Expand All @@ -55,10 +53,8 @@ import com.intellij.psi.PsiJavaFile
import com.intellij.psi.PsiManager
import com.intellij.psi.PsiTreeChangeAdapter
import com.intellij.psi.PsiTreeChangeListener
import com.intellij.psi.impl.file.impl.JavaFileManager
import com.intellij.psi.search.DelegatingGlobalSearchScope
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.search.ProjectScope
import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider
Expand All @@ -83,24 +79,15 @@ import org.jetbrains.kotlin.analysis.providers.impl.*
import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoot
import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoots
import org.jetbrains.kotlin.cli.common.config.kotlinSourceRoots
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCliJavaFileManagerImpl
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreApplicationEnvironmentMode
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
import org.jetbrains.kotlin.cli.jvm.compiler.computeDefaultRootModules
import org.jetbrains.kotlin.cli.jvm.compiler.createSourceFilesFromSourceRoots
import org.jetbrains.kotlin.cli.jvm.compiler.getJavaModuleRoots
import org.jetbrains.kotlin.cli.jvm.compiler.setupIdeaStandaloneExecution
import org.jetbrains.kotlin.cli.jvm.config.addJavaSourceRoot
import org.jetbrains.kotlin.cli.jvm.config.addJavaSourceRoots
import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots
import org.jetbrains.kotlin.cli.jvm.config.jvmModularRoots
import org.jetbrains.kotlin.cli.jvm.index.JavaRoot
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesDynamicCompoundIndex
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesIndexImpl
import org.jetbrains.kotlin.cli.jvm.index.SingleJavaFileRootsIndex
import org.jetbrains.kotlin.cli.jvm.modules.CliJavaModuleFinder
import org.jetbrains.kotlin.cli.jvm.modules.JavaModuleGraph
import org.jetbrains.kotlin.config.ApiVersion
import org.jetbrains.kotlin.config.CommonConfigurationKeys
import org.jetbrains.kotlin.config.CompilerConfiguration
Expand Down Expand Up @@ -351,7 +338,8 @@ class KotlinSymbolProcessing(
private fun prepareAllKSFiles(
kotlinCoreProjectEnvironment: KotlinCoreProjectEnvironment,
modules: List<KtModule>,
compilerConfiguration: CompilerConfiguration
compilerConfiguration: CompilerConfiguration,
javaFileManager: IncrementalJavaFileManager,
): List<KSFile> {
val project = kotlinCoreProjectEnvironment.project
val ktFiles = createSourceFilesFromSourceRoots(
Expand All @@ -375,16 +363,15 @@ class KotlinSymbolProcessing(
).update(ktFiles)

// Update Java providers for newly generated source files.
reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, allJavaFiles)
javaFileManager.initialize(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,
javaFileManager: IncrementalJavaFileManager,
newKotlinFiles: List<File>,
newJavaFiles: List<File>,
): List<KSFile> {
Expand All @@ -397,10 +384,6 @@ class KotlinSymbolProcessing(
project,
newJavaFiles.map { it.toPath() }.toSet()
)
val allJavaFiles = getPsiFilesFromPaths<PsiJavaFile>(
project,
getSourceFilePaths(compilerConfiguration, includeDirectoryRoot = true)
)

// Update Kotlin providers for newly generated source files.
(
Expand All @@ -415,7 +398,7 @@ class KotlinSymbolProcessing(
).update(ktFiles)

// Update Java providers for newly generated source files.
reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, allJavaFiles)
javaFileManager.add(javaFiles)

return ktFiles.map { analyze { KSFileImpl.getCached(it.getFileSymbol()) } } +
javaFiles.map { KSFileJavaImpl.getCached(it) }
Expand Down Expand Up @@ -470,7 +453,9 @@ class KotlinSymbolProcessing(
ResolverAAImpl.ktModule = modules.single() as KtSourceModule

// Initializing environments
val allKSFiles = prepareAllKSFiles(kotlinCoreProjectEnvironment, modules, compilerConfiguration)
val javaFileManager = IncrementalJavaFileManager(kotlinCoreProjectEnvironment)
val allKSFiles =
prepareAllKSFiles(kotlinCoreProjectEnvironment, modules, compilerConfiguration, javaFileManager)
val anyChangesWildcard = AnyChanges(kspConfig.projectBaseDir)
val codeGenerator = CodeGeneratorImpl(
kspConfig.classOutputDir,
Expand Down Expand Up @@ -499,7 +484,6 @@ class KotlinSymbolProcessing(
var allDirtyKSFiles = incrementalContext.calcDirtyFiles(allKSFiles).toList()
var newKSFiles = allDirtyKSFiles
val initialDirtySet = allDirtyKSFiles.toSet()
val allCleanFilePaths = allKSFiles.filterNot { it in initialDirtySet }.map { it.filePath }.toSet()

val targetPlatform = ResolverAAImpl.ktModule.platform
val symbolProcessorEnvironment = SymbolProcessorEnvironment(
Expand Down Expand Up @@ -567,8 +551,7 @@ class KotlinSymbolProcessing(
.map { it.canonicalPath }.toSet()
newKSFiles = prepareNewKSFiles(
kotlinCoreProjectEnvironment,
modules,
compilerConfiguration,
javaFileManager,
newFilePaths.filter { it.endsWith(".kt") }.map { File(it) }.toList(),
newFilePaths.filter { it.endsWith(".java") }.map { File(it) }.toList(),
)
Expand Down Expand Up @@ -648,78 +631,6 @@ class DirectoriesScope(
override fun toString() = "All files under: $directories"
}

private fun reinitJavaFileManager(
environment: KotlinCoreProjectEnvironment,
modules: List<KtModule>,
sourceFiles: List<PsiJavaFile>,
) {
val project = environment.project
val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
val javaModuleFinder = CliJavaModuleFinder(null, null, javaFileManager, project, null)
val javaModuleGraph = JavaModuleGraph(javaModuleFinder)
val allSourceFileRoots = sourceFiles.map { JavaRoot(it.virtualFile, JavaRoot.RootType.SOURCE) }
val jdkRoots = getDefaultJdkModuleRoots(javaModuleFinder, javaModuleGraph)
val libraryRoots = StandaloneProjectFactory.getAllBinaryRoots(modules, environment)

val rootsWithSingleJavaFileRoots = buildList {
addAll(libraryRoots)
addAll(allSourceFileRoots)
addAll(jdkRoots)
}

val (roots, singleJavaFileRoots) = rootsWithSingleJavaFileRoots.partition { (file) ->
file.isDirectory || file.extension != JavaFileType.DEFAULT_EXTENSION
}

val corePackageIndex = project.getService(PackageIndex::class.java) as CorePackageIndex
val rootsIndex = JvmDependenciesDynamicCompoundIndex().apply {
addIndex(JvmDependenciesIndexImpl(roots))
indexedRoots.forEach { javaRoot ->
if (javaRoot.file.isDirectory) {
if (javaRoot.type == JavaRoot.RootType.SOURCE) {
// NB: [JavaCoreProjectEnvironment#addSourcesToClasspath] calls:
// 1) [CoreJavaFileManager#addToClasspath], which is used to look up Java roots;
// 2) [CorePackageIndex#addToClasspath], which populates [PackageIndex]; and
// 3) [FileIndexFacade#addLibraryRoot], which conflicts with this SOURCE root when generating a library scope.
// Thus, here we manually call first two, which are used to:
// 1) create [PsiPackage] as a package resolution result; and
// 2) find directories by package name.
// With both supports, annotations defined in package-info.java can be properly propagated.
javaFileManager.addToClasspath(javaRoot.file)
corePackageIndex.addToClasspath(javaRoot.file)
} else {
environment.addSourcesToClasspath(javaRoot.file)
}
}
}
}

javaFileManager.initialize(
rootsIndex,
listOf(
StandaloneProjectFactory.createPackagePartsProvider(
libraryRoots + jdkRoots,
LanguageVersionSettingsImpl(LanguageVersion.LATEST_STABLE, ApiVersion.LATEST)
).invoke(ProjectScope.getLibrariesScope(project))
),
SingleJavaFileRootsIndex(singleJavaFileRoots),
true
)
}

private fun getDefaultJdkModuleRoots(
javaModuleFinder: CliJavaModuleFinder,
javaModuleGraph: JavaModuleGraph
): List<JavaRoot> {
// In contrast to `ClasspathRootsResolver.addModularRoots`, we do not need to handle automatic Java modules because JDK modules
// aren't automatic.
return javaModuleGraph.getAllDependencies(javaModuleFinder.computeDefaultRootModules()).flatMap { moduleName ->
val module = javaModuleFinder.findModule(moduleName) ?: return@flatMap emptyList<JavaRoot>()
val result = module.getJavaModuleRoots()
result
}
}

fun String?.toKotlinVersion(): KotlinVersion {
if (this == null)
return KotlinVersion.CURRENT
Expand Down
149 changes: 149 additions & 0 deletions ...ysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalJavaFileManager.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
package com.google.devtools.ksp.standalone

import com.intellij.core.CorePackageIndex
import com.intellij.ide.highlighter.JavaFileType
import com.intellij.openapi.roots.PackageIndex
import com.intellij.psi.PsiJavaFile
import com.intellij.psi.impl.file.impl.JavaFileManager
import com.intellij.psi.search.ProjectScope
import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.StandaloneProjectFactory
import org.jetbrains.kotlin.analysis.project.structure.KtModule
import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCliJavaFileManagerImpl
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
import org.jetbrains.kotlin.cli.jvm.compiler.computeDefaultRootModules
import org.jetbrains.kotlin.cli.jvm.compiler.getJavaModuleRoots
import org.jetbrains.kotlin.cli.jvm.index.JavaRoot
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesDynamicCompoundIndex
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesIndexImpl
import org.jetbrains.kotlin.cli.jvm.index.SingleJavaFileRootsIndex
import org.jetbrains.kotlin.cli.jvm.modules.CliJavaModuleFinder
import org.jetbrains.kotlin.cli.jvm.modules.JavaModuleGraph
import org.jetbrains.kotlin.config.ApiVersion
import org.jetbrains.kotlin.config.LanguageVersion
import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl

class IncrementalJavaFileManager(val environment: KotlinCoreProjectEnvironment) {
lateinit var rootsIndex: JvmDependenciesDynamicCompoundIndex
lateinit var packagePartProviders: List<JvmPackagePartProvider>
val singleJavaFileRoots = mutableListOf<JavaRoot>()

fun initialize(
modules: List<KtModule>,
sourceFiles: List<PsiJavaFile>,
) {
val project = environment.project
val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
val javaModuleFinder = CliJavaModuleFinder(null, null, javaFileManager, project, null)
val javaModuleGraph = JavaModuleGraph(javaModuleFinder)
val allSourceFileRoots = sourceFiles.map { JavaRoot(it.virtualFile, JavaRoot.RootType.SOURCE) }
val jdkRoots = getDefaultJdkModuleRoots(javaModuleFinder, javaModuleGraph)
val libraryRoots = StandaloneProjectFactory.getAllBinaryRoots(modules, environment)

val rootsWithSingleJavaFileRoots = buildList {
addAll(libraryRoots)
addAll(allSourceFileRoots)
addAll(jdkRoots)
}

val (roots, newSingleJavaFileRoots) = rootsWithSingleJavaFileRoots.partition { (file) ->
file.isDirectory || file.extension != JavaFileType.DEFAULT_EXTENSION
}

singleJavaFileRoots.addAll(newSingleJavaFileRoots)

rootsIndex = JvmDependenciesDynamicCompoundIndex().apply {
addIndex(JvmDependenciesIndexImpl(roots))
}

val corePackageIndex = project.getService(PackageIndex::class.java) as CorePackageIndex
roots.forEach { javaRoot ->
if (javaRoot.file.isDirectory) {
if (javaRoot.type == JavaRoot.RootType.SOURCE) {
// NB: [JavaCoreProjectEnvironment#addSourcesToClasspath] calls:
// 1) [CoreJavaFileManager#addToClasspath], which is used to look up Java roots;
// 2) [CorePackageIndex#addToClasspath], which populates [PackageIndex]; and
// 3) [FileIndexFacade#addLibraryRoot], which conflicts with this SOURCE root when generating a library scope.
// Thus, here we manually call first two, which are used to:
// 1) create [PsiPackage] as a package resolution result; and
// 2) find directories by package name.
// With both supports, annotations defined in package-info.java can be properly propagated.
javaFileManager.addToClasspath(javaRoot.file)
corePackageIndex.addToClasspath(javaRoot.file)
} else {
environment.addSourcesToClasspath(javaRoot.file)
}
}
}

packagePartProviders = listOf(
StandaloneProjectFactory.createPackagePartsProvider(
libraryRoots + jdkRoots,
LanguageVersionSettingsImpl(LanguageVersion.LATEST_STABLE, ApiVersion.LATEST)
).invoke(ProjectScope.getLibrariesScope(project))
)

javaFileManager.initialize(
rootsIndex,
packagePartProviders,
SingleJavaFileRootsIndex(singleJavaFileRoots),
true
)
}

fun add(sourceFiles: List<PsiJavaFile>) {
val project = environment.project
val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
val allSourceFileRoots = sourceFiles.map { JavaRoot(it.virtualFile, JavaRoot.RootType.SOURCE) }

val (roots, newSingleJavaFileRoots) = allSourceFileRoots.partition { (file) ->
file.isDirectory || file.extension != JavaFileType.DEFAULT_EXTENSION
}

singleJavaFileRoots.addAll(newSingleJavaFileRoots)

rootsIndex.apply {
addIndex(JvmDependenciesIndexImpl(roots))
}

val corePackageIndex = project.getService(PackageIndex::class.java) as CorePackageIndex
roots.forEach { javaRoot ->
if (javaRoot.file.isDirectory) {
if (javaRoot.type == JavaRoot.RootType.SOURCE) {
// NB: [JavaCoreProjectEnvironment#addSourcesToClasspath] calls:
// 1) [CoreJavaFileManager#addToClasspath], which is used to look up Java roots;
// 2) [CorePackageIndex#addToClasspath], which populates [PackageIndex]; and
// 3) [FileIndexFacade#addLibraryRoot], which conflicts with this SOURCE root when generating a library scope.
// Thus, here we manually call first two, which are used to:
// 1) create [PsiPackage] as a package resolution result; and
// 2) find directories by package name.
// With both supports, annotations defined in package-info.java can be properly propagated.
javaFileManager.addToClasspath(javaRoot.file)
corePackageIndex.addToClasspath(javaRoot.file)
} else {
environment.addSourcesToClasspath(javaRoot.file)
}
}
}

javaFileManager.initialize(
rootsIndex,
packagePartProviders,
SingleJavaFileRootsIndex(singleJavaFileRoots),
true
)
}
}

private fun getDefaultJdkModuleRoots(
javaModuleFinder: CliJavaModuleFinder,
javaModuleGraph: JavaModuleGraph
): List<JavaRoot> {
// In contrast to `ClasspathRootsResolver.addModularRoots`, we do not need to handle automatic Java modules because JDK modules
// aren't automatic.
return javaModuleGraph.getAllDependencies(javaModuleFinder.computeDefaultRootModules()).flatMap { moduleName ->
val module = javaModuleFinder.findModule(moduleName) ?: return@flatMap emptyList<JavaRoot>()
val result = module.getJavaModuleRoots()
result
}
}
Loading