Skip to content

Commit 27f57e8

Browse files
ting-yuanKSP Auto Pick
ting-yuan
authored and
KSP Auto Pick
committed
KSP2: Reuse Java indexes
(cherry picked from commit 760fe28)
1 parent 6bba35b commit 27f57e8

File tree

2 files changed

+159
-99
lines changed

2 files changed

+159
-99
lines changed

‎kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt

Lines changed: 10 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import com.google.devtools.ksp.impl.symbol.kotlin.KSFileJavaImpl
3131
import com.google.devtools.ksp.impl.symbol.kotlin.Restorable
3232
import com.google.devtools.ksp.impl.symbol.kotlin.analyze
3333
import com.google.devtools.ksp.processing.*
34+
import com.google.devtools.ksp.standalone.IncrementalJavaFileManager
3435
import com.google.devtools.ksp.standalone.IncrementalKotlinDeclarationProviderFactory
3536
import com.google.devtools.ksp.standalone.IncrementalKotlinPackageProviderFactory
3637
import com.google.devtools.ksp.standalone.KspStandaloneDirectInheritorsProvider
@@ -40,13 +41,10 @@ import com.google.devtools.ksp.symbol.KSFile
4041
import com.google.devtools.ksp.symbol.KSNode
4142
import com.google.devtools.ksp.symbol.Origin
4243
import com.intellij.core.CoreApplicationEnvironment
43-
import com.intellij.core.CorePackageIndex
44-
import com.intellij.ide.highlighter.JavaFileType
4544
import com.intellij.mock.MockProject
4645
import com.intellij.openapi.Disposable
4746
import com.intellij.openapi.application.Application
4847
import com.intellij.openapi.project.Project
49-
import com.intellij.openapi.roots.PackageIndex
5048
import com.intellij.openapi.util.Disposer
5149
import com.intellij.openapi.vfs.StandardFileSystems
5250
import com.intellij.openapi.vfs.VirtualFile
@@ -55,10 +53,8 @@ import com.intellij.psi.PsiJavaFile
5553
import com.intellij.psi.PsiManager
5654
import com.intellij.psi.PsiTreeChangeAdapter
5755
import com.intellij.psi.PsiTreeChangeListener
58-
import com.intellij.psi.impl.file.impl.JavaFileManager
5956
import com.intellij.psi.search.DelegatingGlobalSearchScope
6057
import com.intellij.psi.search.GlobalSearchScope
61-
import com.intellij.psi.search.ProjectScope
6258
import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
6359
import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
6460
import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider
@@ -83,24 +79,15 @@ import org.jetbrains.kotlin.analysis.providers.impl.*
8379
import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoot
8480
import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoots
8581
import org.jetbrains.kotlin.cli.common.config.kotlinSourceRoots
86-
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCliJavaFileManagerImpl
8782
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreApplicationEnvironmentMode
8883
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
89-
import org.jetbrains.kotlin.cli.jvm.compiler.computeDefaultRootModules
9084
import org.jetbrains.kotlin.cli.jvm.compiler.createSourceFilesFromSourceRoots
91-
import org.jetbrains.kotlin.cli.jvm.compiler.getJavaModuleRoots
9285
import org.jetbrains.kotlin.cli.jvm.compiler.setupIdeaStandaloneExecution
9386
import org.jetbrains.kotlin.cli.jvm.config.addJavaSourceRoot
9487
import org.jetbrains.kotlin.cli.jvm.config.addJavaSourceRoots
9588
import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
9689
import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots
9790
import org.jetbrains.kotlin.cli.jvm.config.jvmModularRoots
98-
import org.jetbrains.kotlin.cli.jvm.index.JavaRoot
99-
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesDynamicCompoundIndex
100-
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesIndexImpl
101-
import org.jetbrains.kotlin.cli.jvm.index.SingleJavaFileRootsIndex
102-
import org.jetbrains.kotlin.cli.jvm.modules.CliJavaModuleFinder
103-
import org.jetbrains.kotlin.cli.jvm.modules.JavaModuleGraph
10491
import org.jetbrains.kotlin.config.ApiVersion
10592
import org.jetbrains.kotlin.config.CommonConfigurationKeys
10693
import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -351,7 +338,8 @@ class KotlinSymbolProcessing(
351338
private fun prepareAllKSFiles(
352339
kotlinCoreProjectEnvironment: KotlinCoreProjectEnvironment,
353340
modules: List<KtModule>,
354-
compilerConfiguration: CompilerConfiguration
341+
compilerConfiguration: CompilerConfiguration,
342+
javaFileManager: IncrementalJavaFileManager,
355343
): List<KSFile> {
356344
val project = kotlinCoreProjectEnvironment.project
357345
val ktFiles = createSourceFilesFromSourceRoots(
@@ -375,16 +363,15 @@ class KotlinSymbolProcessing(
375363
).update(ktFiles)
376364

377365
// Update Java providers for newly generated source files.
378-
reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, allJavaFiles)
366+
javaFileManager.initialize(modules, allJavaFiles)
379367

380368
return ktFiles.map { analyze { KSFileImpl.getCached(it.getFileSymbol()) } } +
381369
allJavaFiles.map { KSFileJavaImpl.getCached(it) }
382370
}
383371

384372
private fun prepareNewKSFiles(
385373
kotlinCoreProjectEnvironment: KotlinCoreProjectEnvironment,
386-
modules: List<KtModule>,
387-
compilerConfiguration: CompilerConfiguration,
374+
javaFileManager: IncrementalJavaFileManager,
388375
newKotlinFiles: List<File>,
389376
newJavaFiles: List<File>,
390377
): List<KSFile> {
@@ -397,10 +384,6 @@ class KotlinSymbolProcessing(
397384
project,
398385
newJavaFiles.map { it.toPath() }.toSet()
399386
)
400-
val allJavaFiles = getPsiFilesFromPaths<PsiJavaFile>(
401-
project,
402-
getSourceFilePaths(compilerConfiguration, includeDirectoryRoot = true)
403-
)
404387

405388
// Update Kotlin providers for newly generated source files.
406389
(
@@ -415,7 +398,7 @@ class KotlinSymbolProcessing(
415398
).update(ktFiles)
416399

417400
// Update Java providers for newly generated source files.
418-
reinitJavaFileManager(kotlinCoreProjectEnvironment, modules, allJavaFiles)
401+
javaFileManager.add(javaFiles)
419402

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

472455
// Initializing environments
473-
val allKSFiles = prepareAllKSFiles(kotlinCoreProjectEnvironment, modules, compilerConfiguration)
456+
val javaFileManager = IncrementalJavaFileManager(kotlinCoreProjectEnvironment)
457+
val allKSFiles =
458+
prepareAllKSFiles(kotlinCoreProjectEnvironment, modules, compilerConfiguration, javaFileManager)
474459
val anyChangesWildcard = AnyChanges(kspConfig.projectBaseDir)
475460
val codeGenerator = CodeGeneratorImpl(
476461
kspConfig.classOutputDir,
@@ -499,7 +484,6 @@ class KotlinSymbolProcessing(
499484
var allDirtyKSFiles = incrementalContext.calcDirtyFiles(allKSFiles).toList()
500485
var newKSFiles = allDirtyKSFiles
501486
val initialDirtySet = allDirtyKSFiles.toSet()
502-
val allCleanFilePaths = allKSFiles.filterNot { it in initialDirtySet }.map { it.filePath }.toSet()
503487

504488
val targetPlatform = ResolverAAImpl.ktModule.platform
505489
val symbolProcessorEnvironment = SymbolProcessorEnvironment(
@@ -567,8 +551,7 @@ class KotlinSymbolProcessing(
567551
.map { it.canonicalPath }.toSet()
568552
newKSFiles = prepareNewKSFiles(
569553
kotlinCoreProjectEnvironment,
570-
modules,
571-
compilerConfiguration,
554+
javaFileManager,
572555
newFilePaths.filter { it.endsWith(".kt") }.map { File(it) }.toList(),
573556
newFilePaths.filter { it.endsWith(".java") }.map { File(it) }.toList(),
574557
)
@@ -648,78 +631,6 @@ class DirectoriesScope(
648631
override fun toString() = "All files under: $directories"
649632
}
650633

651-
private fun reinitJavaFileManager(
652-
environment: KotlinCoreProjectEnvironment,
653-
modules: List<KtModule>,
654-
sourceFiles: List<PsiJavaFile>,
655-
) {
656-
val project = environment.project
657-
val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
658-
val javaModuleFinder = CliJavaModuleFinder(null, null, javaFileManager, project, null)
659-
val javaModuleGraph = JavaModuleGraph(javaModuleFinder)
660-
val allSourceFileRoots = sourceFiles.map { JavaRoot(it.virtualFile, JavaRoot.RootType.SOURCE) }
661-
val jdkRoots = getDefaultJdkModuleRoots(javaModuleFinder, javaModuleGraph)
662-
val libraryRoots = StandaloneProjectFactory.getAllBinaryRoots(modules, environment)
663-
664-
val rootsWithSingleJavaFileRoots = buildList {
665-
addAll(libraryRoots)
666-
addAll(allSourceFileRoots)
667-
addAll(jdkRoots)
668-
}
669-
670-
val (roots, singleJavaFileRoots) = rootsWithSingleJavaFileRoots.partition { (file) ->
671-
file.isDirectory || file.extension != JavaFileType.DEFAULT_EXTENSION
672-
}
673-
674-
val corePackageIndex = project.getService(PackageIndex::class.java) as CorePackageIndex
675-
val rootsIndex = JvmDependenciesDynamicCompoundIndex().apply {
676-
addIndex(JvmDependenciesIndexImpl(roots))
677-
indexedRoots.forEach { javaRoot ->
678-
if (javaRoot.file.isDirectory) {
679-
if (javaRoot.type == JavaRoot.RootType.SOURCE) {
680-
// NB: [JavaCoreProjectEnvironment#addSourcesToClasspath] calls:
681-
// 1) [CoreJavaFileManager#addToClasspath], which is used to look up Java roots;
682-
// 2) [CorePackageIndex#addToClasspath], which populates [PackageIndex]; and
683-
// 3) [FileIndexFacade#addLibraryRoot], which conflicts with this SOURCE root when generating a library scope.
684-
// Thus, here we manually call first two, which are used to:
685-
// 1) create [PsiPackage] as a package resolution result; and
686-
// 2) find directories by package name.
687-
// With both supports, annotations defined in package-info.java can be properly propagated.
688-
javaFileManager.addToClasspath(javaRoot.file)
689-
corePackageIndex.addToClasspath(javaRoot.file)
690-
} else {
691-
environment.addSourcesToClasspath(javaRoot.file)
692-
}
693-
}
694-
}
695-
}
696-
697-
javaFileManager.initialize(
698-
rootsIndex,
699-
listOf(
700-
StandaloneProjectFactory.createPackagePartsProvider(
701-
libraryRoots + jdkRoots,
702-
LanguageVersionSettingsImpl(LanguageVersion.LATEST_STABLE, ApiVersion.LATEST)
703-
).invoke(ProjectScope.getLibrariesScope(project))
704-
),
705-
SingleJavaFileRootsIndex(singleJavaFileRoots),
706-
true
707-
)
708-
}
709-
710-
private fun getDefaultJdkModuleRoots(
711-
javaModuleFinder: CliJavaModuleFinder,
712-
javaModuleGraph: JavaModuleGraph
713-
): List<JavaRoot> {
714-
// In contrast to `ClasspathRootsResolver.addModularRoots`, we do not need to handle automatic Java modules because JDK modules
715-
// aren't automatic.
716-
return javaModuleGraph.getAllDependencies(javaModuleFinder.computeDefaultRootModules()).flatMap { moduleName ->
717-
val module = javaModuleFinder.findModule(moduleName) ?: return@flatMap emptyList<JavaRoot>()
718-
val result = module.getJavaModuleRoots()
719-
result
720-
}
721-
}
722-
723634
fun String?.toKotlinVersion(): KotlinVersion {
724635
if (this == null)
725636
return KotlinVersion.CURRENT

‎kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/standalone/IncrementalJavaFileManager.kt

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
package com.google.devtools.ksp.standalone
2+
3+
import com.intellij.core.CorePackageIndex
4+
import com.intellij.ide.highlighter.JavaFileType
5+
import com.intellij.openapi.roots.PackageIndex
6+
import com.intellij.psi.PsiJavaFile
7+
import com.intellij.psi.impl.file.impl.JavaFileManager
8+
import com.intellij.psi.search.ProjectScope
9+
import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.StandaloneProjectFactory
10+
import org.jetbrains.kotlin.analysis.project.structure.KtModule
11+
import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
12+
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCliJavaFileManagerImpl
13+
import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
14+
import org.jetbrains.kotlin.cli.jvm.compiler.computeDefaultRootModules
15+
import org.jetbrains.kotlin.cli.jvm.compiler.getJavaModuleRoots
16+
import org.jetbrains.kotlin.cli.jvm.index.JavaRoot
17+
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesDynamicCompoundIndex
18+
import org.jetbrains.kotlin.cli.jvm.index.JvmDependenciesIndexImpl
19+
import org.jetbrains.kotlin.cli.jvm.index.SingleJavaFileRootsIndex
20+
import org.jetbrains.kotlin.cli.jvm.modules.CliJavaModuleFinder
21+
import org.jetbrains.kotlin.cli.jvm.modules.JavaModuleGraph
22+
import org.jetbrains.kotlin.config.ApiVersion
23+
import org.jetbrains.kotlin.config.LanguageVersion
24+
import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
25+
26+
class IncrementalJavaFileManager(val environment: KotlinCoreProjectEnvironment) {
27+
lateinit var rootsIndex: JvmDependenciesDynamicCompoundIndex
28+
lateinit var packagePartProviders: List<JvmPackagePartProvider>
29+
val singleJavaFileRoots = mutableListOf<JavaRoot>()
30+
31+
fun initialize(
32+
modules: List<KtModule>,
33+
sourceFiles: List<PsiJavaFile>,
34+
) {
35+
val project = environment.project
36+
val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
37+
val javaModuleFinder = CliJavaModuleFinder(null, null, javaFileManager, project, null)
38+
val javaModuleGraph = JavaModuleGraph(javaModuleFinder)
39+
val allSourceFileRoots = sourceFiles.map { JavaRoot(it.virtualFile, JavaRoot.RootType.SOURCE) }
40+
val jdkRoots = getDefaultJdkModuleRoots(javaModuleFinder, javaModuleGraph)
41+
val libraryRoots = StandaloneProjectFactory.getAllBinaryRoots(modules, environment)
42+
43+
val rootsWithSingleJavaFileRoots = buildList {
44+
addAll(libraryRoots)
45+
addAll(allSourceFileRoots)
46+
addAll(jdkRoots)
47+
}
48+
49+
val (roots, newSingleJavaFileRoots) = rootsWithSingleJavaFileRoots.partition { (file) ->
50+
file.isDirectory || file.extension != JavaFileType.DEFAULT_EXTENSION
51+
}
52+
53+
singleJavaFileRoots.addAll(newSingleJavaFileRoots)
54+
55+
rootsIndex = JvmDependenciesDynamicCompoundIndex().apply {
56+
addIndex(JvmDependenciesIndexImpl(roots))
57+
}
58+
59+
val corePackageIndex = project.getService(PackageIndex::class.java) as CorePackageIndex
60+
roots.forEach { javaRoot ->
61+
if (javaRoot.file.isDirectory) {
62+
if (javaRoot.type == JavaRoot.RootType.SOURCE) {
63+
// NB: [JavaCoreProjectEnvironment#addSourcesToClasspath] calls:
64+
// 1) [CoreJavaFileManager#addToClasspath], which is used to look up Java roots;
65+
// 2) [CorePackageIndex#addToClasspath], which populates [PackageIndex]; and
66+
// 3) [FileIndexFacade#addLibraryRoot], which conflicts with this SOURCE root when generating a library scope.
67+
// Thus, here we manually call first two, which are used to:
68+
// 1) create [PsiPackage] as a package resolution result; and
69+
// 2) find directories by package name.
70+
// With both supports, annotations defined in package-info.java can be properly propagated.
71+
javaFileManager.addToClasspath(javaRoot.file)
72+
corePackageIndex.addToClasspath(javaRoot.file)
73+
} else {
74+
environment.addSourcesToClasspath(javaRoot.file)
75+
}
76+
}
77+
}
78+
79+
packagePartProviders = listOf(
80+
StandaloneProjectFactory.createPackagePartsProvider(
81+
libraryRoots + jdkRoots,
82+
LanguageVersionSettingsImpl(LanguageVersion.LATEST_STABLE, ApiVersion.LATEST)
83+
).invoke(ProjectScope.getLibrariesScope(project))
84+
)
85+
86+
javaFileManager.initialize(
87+
rootsIndex,
88+
packagePartProviders,
89+
SingleJavaFileRootsIndex(singleJavaFileRoots),
90+
true
91+
)
92+
}
93+
94+
fun add(sourceFiles: List<PsiJavaFile>) {
95+
val project = environment.project
96+
val javaFileManager = project.getService(JavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
97+
val allSourceFileRoots = sourceFiles.map { JavaRoot(it.virtualFile, JavaRoot.RootType.SOURCE) }
98+
99+
val (roots, newSingleJavaFileRoots) = allSourceFileRoots.partition { (file) ->
100+
file.isDirectory || file.extension != JavaFileType.DEFAULT_EXTENSION
101+
}
102+
103+
singleJavaFileRoots.addAll(newSingleJavaFileRoots)
104+
105+
rootsIndex.apply {
106+
addIndex(JvmDependenciesIndexImpl(roots))
107+
}
108+
109+
val corePackageIndex = project.getService(PackageIndex::class.java) as CorePackageIndex
110+
roots.forEach { javaRoot ->
111+
if (javaRoot.file.isDirectory) {
112+
if (javaRoot.type == JavaRoot.RootType.SOURCE) {
113+
// NB: [JavaCoreProjectEnvironment#addSourcesToClasspath] calls:
114+
// 1) [CoreJavaFileManager#addToClasspath], which is used to look up Java roots;
115+
// 2) [CorePackageIndex#addToClasspath], which populates [PackageIndex]; and
116+
// 3) [FileIndexFacade#addLibraryRoot], which conflicts with this SOURCE root when generating a library scope.
117+
// Thus, here we manually call first two, which are used to:
118+
// 1) create [PsiPackage] as a package resolution result; and
119+
// 2) find directories by package name.
120+
// With both supports, annotations defined in package-info.java can be properly propagated.
121+
javaFileManager.addToClasspath(javaRoot.file)
122+
corePackageIndex.addToClasspath(javaRoot.file)
123+
} else {
124+
environment.addSourcesToClasspath(javaRoot.file)
125+
}
126+
}
127+
}
128+
129+
javaFileManager.initialize(
130+
rootsIndex,
131+
packagePartProviders,
132+
SingleJavaFileRootsIndex(singleJavaFileRoots),
133+
true
134+
)
135+
}
136+
}
137+
138+
private fun getDefaultJdkModuleRoots(
139+
javaModuleFinder: CliJavaModuleFinder,
140+
javaModuleGraph: JavaModuleGraph
141+
): List<JavaRoot> {
142+
// In contrast to `ClasspathRootsResolver.addModularRoots`, we do not need to handle automatic Java modules because JDK modules
143+
// aren't automatic.
144+
return javaModuleGraph.getAllDependencies(javaModuleFinder.computeDefaultRootModules()).flatMap { moduleName ->
145+
val module = javaModuleFinder.findModule(moduleName) ?: return@flatMap emptyList<JavaRoot>()
146+
val result = module.getJavaModuleRoots()
147+
result
148+
}
149+
}

0 commit comments

Comments
 (0)