Skip to content

Commit b665181

Browse files
jbarr21jbarr21
committed
Implement jdeps using K2 APIs
1 parent a00963b commit b665181

14 files changed

+599
-7
lines changed

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/BUILD.bazel

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,12 @@ load("//src/main/kotlin:bootstrap.bzl", "kt_bootstrap_library")
1919
# The compiler binary, this is co-located in the kotlin compiler classloader.
2020
kt_bootstrap_library(
2121
name = "jdeps-gen-lib",
22-
srcs = glob(["*.kt"]),
22+
srcs = glob([
23+
"*.kt",
24+
"k2/*.kt",
25+
"k2/checker/declaration/*.kt",
26+
"k2/checker/expression/*.kt",
27+
]),
2328
visibility = ["//src:__subpackages__"],
2429
deps = [
2530
"//kotlin/compiler:kotlin-compiler",

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/BaseJdepsGenExtension.kt

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import java.nio.file.Paths
1010
abstract class BaseJdepsGenExtension(
1111
protected val configuration: CompilerConfiguration,
1212
) {
13-
1413
protected fun onAnalysisCompleted(
1514
explicitClassesCanonicalPaths: Set<String>,
1615
implicitClassesCanonicalPaths: Set<String>,
@@ -107,9 +106,10 @@ abstract class BaseJdepsGenExtension(
107106
directDeps: List<String>,
108107
targetLabel: String,
109108
): Boolean {
110-
val missingStrictDeps = result.keys
111-
.filter { !directDeps.contains(it) }
112-
.map { JarOwner.readJarOwnerFromManifest(Paths.get(it)) }
109+
val missingStrictDeps =
110+
result.keys
111+
.filter { !directDeps.contains(it) }
112+
.map { JarOwner.readJarOwnerFromManifest(Paths.get(it)) }
113113

114114
if (missingStrictDeps.isNotEmpty()) {
115115
val missingStrictLabels = missingStrictDeps.mapNotNull { it.label }

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/JdepsGenComponentRegistrar.kt

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,43 @@
11
package io.bazel.kotlin.plugin.jdeps
22

3+
import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
4+
import io.bazel.kotlin.plugin.jdeps.k2.JdepsFirExtensions
5+
import io.bazel.kotlin.plugin.jdeps.k2.JdepsGenExtension2
6+
import org.jetbrains.kotlin.codegen.extensions.ClassFileFactoryFinalizerExtension
37
import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
48
import org.jetbrains.kotlin.config.CompilerConfiguration
9+
import org.jetbrains.kotlin.config.languageVersionSettings
510
import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
11+
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
612
import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
13+
import java.nio.file.Paths
714

815
@OptIn(org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi::class)
916
class JdepsGenComponentRegistrar : CompilerPluginRegistrar() {
10-
1117
override val supportsK2: Boolean
12-
get() = false
18+
get() = true
1319

1420
override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
21+
when (configuration.languageVersionSettings.languageVersion.usesK2) {
22+
true -> registerForK2(configuration)
23+
false -> registerForK1(configuration)
24+
}
25+
}
26+
27+
private fun ExtensionStorage.registerForK1(configuration: CompilerConfiguration) {
1528
// Capture all types referenced by the compiler for this module and look up the jar from which
1629
// they were loaded from
1730
val extension = JdepsGenExtension(configuration)
1831
AnalysisHandlerExtension.registerExtension(extension)
1932
StorageComponentContainerContributor.registerExtension(extension)
2033
}
34+
35+
private fun ExtensionStorage.registerForK2(configuration: CompilerConfiguration) {
36+
val projectRoot = Paths.get("").toAbsolutePath().toString() + "/"
37+
ClassUsageRecorder.init(projectRoot)
38+
JdepsGenExtension2(configuration).run {
39+
FirExtensionRegistrarAdapter.registerExtension(JdepsFirExtensions())
40+
ClassFileFactoryFinalizerExtension.registerExtension(this)
41+
}
42+
}
2143
}

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/ClassUsageRecorder.kt

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
package io.bazel.kotlin.plugin.jdeps.k2
2+
3+
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
4+
import org.jetbrains.kotlin.fir.declarations.utils.sourceElement
5+
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
6+
import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
7+
import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
8+
import org.jetbrains.kotlin.fir.types.ConeKotlinType
9+
import org.jetbrains.kotlin.fir.types.FirTypeRef
10+
import org.jetbrains.kotlin.fir.types.classId
11+
import org.jetbrains.kotlin.fir.types.coneType
12+
import org.jetbrains.kotlin.fir.types.forEachType
13+
import org.jetbrains.kotlin.name.ClassId
14+
import java.nio.file.Paths
15+
import java.util.SortedSet
16+
17+
private const val JAR_FILE_SEPARATOR = "!/"
18+
private const val ANONYMOUS = "<anonymous>"
19+
20+
internal object ClassUsageRecorder {
21+
private val explicitClassesCanonicalPaths = mutableSetOf<String>()
22+
private val implicitClassesCanonicalPaths = mutableSetOf<String>()
23+
24+
private val results by lazy { sortedMapOf<String, SortedSet<String>>() }
25+
private val seen = mutableSetOf<ClassId>()
26+
private val javaHome: String by lazy { System.getenv()["JAVA_HOME"] ?: "<not set>" }
27+
private var rootPath: String = Paths.get("").toAbsolutePath().toString() + "/"
28+
29+
fun init(rootPath: String) {
30+
this.rootPath = rootPath
31+
results.clear()
32+
explicitClassesCanonicalPaths.clear()
33+
implicitClassesCanonicalPaths.clear()
34+
seen.clear()
35+
}
36+
37+
fun getExplicitClassesCanonicalPaths(): Set<String> = explicitClassesCanonicalPaths
38+
39+
fun getImplicitClassesCanonicalPaths(): Set<String> = implicitClassesCanonicalPaths
40+
41+
internal fun recordTypeRef(
42+
typeRef: FirTypeRef,
43+
context: CheckerContext,
44+
isExplicit: Boolean = true,
45+
collectTypeArguments: Boolean = true,
46+
visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
47+
) {
48+
recordConeType(typeRef.coneType, context, isExplicit, collectTypeArguments, visited)
49+
}
50+
51+
internal fun recordConeType(
52+
coneKotlinType: ConeKotlinType,
53+
context: CheckerContext,
54+
isExplicit: Boolean = true,
55+
collectTypeArguments: Boolean = true,
56+
visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
57+
) {
58+
if (collectTypeArguments) {
59+
coneKotlinType.forEachType { coneType ->
60+
val classId = coneType.classId ?: return@forEachType
61+
if (ANONYMOUS in classId.toString()) return@forEachType
62+
context.session.symbolProvider
63+
.getClassLikeSymbolByClassId(classId)
64+
?.recordClass(context, isExplicit, collectTypeArguments, visited)
65+
}
66+
} else {
67+
coneKotlinType.classId?.let { classId ->
68+
context.session.symbolProvider.getClassLikeSymbolByClassId(classId)
69+
?.recordClass(context, isExplicit, collectTypeArguments, visited)
70+
}
71+
}
72+
}
73+
74+
internal fun recordClass(
75+
firClass: FirClassLikeSymbol<*>,
76+
context: CheckerContext,
77+
isExplicit: Boolean,
78+
collectTypeArguments: Boolean = true,
79+
visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
80+
) {
81+
val classIdAndIsExplicit = firClass.classId to isExplicit
82+
if (classIdAndIsExplicit in visited) {
83+
return
84+
} else {
85+
visited.add(classIdAndIsExplicit)
86+
}
87+
88+
firClass.sourceElement?.binaryClass()?.let { addFile(it, isExplicit) }
89+
90+
if (firClass is FirClassSymbol<*>) {
91+
firClass.resolvedSuperTypeRefs.forEach {
92+
recordTypeRef(it, context, false, collectTypeArguments, visited)
93+
}
94+
if (collectTypeArguments) {
95+
firClass.typeParameterSymbols.flatMap { it.resolvedBounds }
96+
.forEach { recordTypeRef(it, context, isExplicit, collectTypeArguments, visited) }
97+
}
98+
}
99+
}
100+
101+
internal fun recordClass(
102+
binaryClass: String,
103+
isExplicit: Boolean = true,
104+
) {
105+
addFile(binaryClass, isExplicit)
106+
}
107+
108+
private fun addFile(
109+
path: String,
110+
isExplicit: Boolean,
111+
) {
112+
if (isExplicit) {
113+
explicitClassesCanonicalPaths.add(path)
114+
} else {
115+
implicitClassesCanonicalPaths.add(path)
116+
}
117+
118+
if (path.contains(JAR_FILE_SEPARATOR) && !path.contains(javaHome)) {
119+
val (jarPath, classPath) = path.split(JAR_FILE_SEPARATOR)
120+
// Convert jar files in current directory to relative paths. Remaining absolute are outside
121+
// of project and should be ignored
122+
val relativizedJarPath = Paths.get(jarPath.replace(rootPath, ""))
123+
if (!relativizedJarPath.isAbsolute) {
124+
val occurrences =
125+
results.computeIfAbsent(relativizedJarPath.toString()) { sortedSetOf<String>() }
126+
if (!isJvmClass(classPath)) {
127+
occurrences.add(classPath)
128+
}
129+
}
130+
}
131+
}
132+
}

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/JdepsFirExtensions.kt

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package io.bazel.kotlin.plugin.jdeps.k2
2+
3+
import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.BasicDeclarationChecker
4+
import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.CallableChecker
5+
import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.ClassLikeChecker
6+
import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.FileChecker
7+
import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.FunctionChecker
8+
import io.bazel.kotlin.plugin.jdeps.k2.checker.expression.QualifiedAccessChecker
9+
import io.bazel.kotlin.plugin.jdeps.k2.checker.expression.ResolvedQualifierChecker
10+
import org.jetbrains.kotlin.fir.FirSession
11+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
12+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirBasicDeclarationChecker
13+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirCallableDeclarationChecker
14+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirClassLikeChecker
15+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirFileChecker
16+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirFunctionChecker
17+
import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
18+
import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirQualifiedAccessExpressionChecker
19+
import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirResolvedQualifierChecker
20+
import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
21+
import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
22+
23+
internal class JdepsFirExtensions : FirExtensionRegistrar() {
24+
override fun ExtensionRegistrarContext.configurePlugin() {
25+
+::JdepsFirCheckersExtension
26+
}
27+
}
28+
29+
internal class JdepsFirCheckersExtension(session: FirSession) :
30+
FirAdditionalCheckersExtension(session) {
31+
override val declarationCheckers: DeclarationCheckers =
32+
object : DeclarationCheckers() {
33+
override val basicDeclarationCheckers: Set<FirBasicDeclarationChecker> =
34+
setOf(BasicDeclarationChecker)
35+
36+
override val fileCheckers: Set<FirFileChecker> = setOf(FileChecker)
37+
38+
override val classLikeCheckers: Set<FirClassLikeChecker> = setOf(ClassLikeChecker)
39+
40+
override val functionCheckers: Set<FirFunctionChecker> = setOf(FunctionChecker)
41+
42+
override val callableDeclarationCheckers: Set<FirCallableDeclarationChecker> =
43+
setOf(CallableChecker)
44+
}
45+
46+
override val expressionCheckers: ExpressionCheckers =
47+
object : ExpressionCheckers() {
48+
override val qualifiedAccessExpressionCheckers: Set<FirQualifiedAccessExpressionChecker> =
49+
setOf(QualifiedAccessChecker)
50+
51+
override val resolvedQualifierCheckers: Set<FirResolvedQualifierChecker> =
52+
setOf(ResolvedQualifierChecker)
53+
}
54+
}

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/JdepsGenExtension2.kt

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package io.bazel.kotlin.plugin.jdeps.k2
2+
3+
import io.bazel.kotlin.plugin.jdeps.BaseJdepsGenExtension
4+
import org.jetbrains.kotlin.codegen.ClassFileFactory
5+
import org.jetbrains.kotlin.codegen.extensions.ClassFileFactoryFinalizerExtension
6+
import org.jetbrains.kotlin.config.CompilerConfiguration
7+
8+
internal class JdepsGenExtension2(
9+
configuration: CompilerConfiguration,
10+
) : BaseJdepsGenExtension(configuration), ClassFileFactoryFinalizerExtension {
11+
override fun finalizeClassFactory(factory: ClassFileFactory) {
12+
onAnalysisCompleted(
13+
ClassUsageRecorder.getExplicitClassesCanonicalPaths(),
14+
ClassUsageRecorder.getImplicitClassesCanonicalPaths(),
15+
)
16+
}
17+
}

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/JdepsK2Utils.kt

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package io.bazel.kotlin.plugin.jdeps.k2
2+
3+
import org.jetbrains.kotlin.descriptors.SourceElement
4+
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
5+
import org.jetbrains.kotlin.fir.java.JavaBinarySourceElement
6+
import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
7+
import org.jetbrains.kotlin.fir.types.ConeKotlinType
8+
import org.jetbrains.kotlin.fir.types.FirTypeRef
9+
import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
10+
import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
11+
import org.jetbrains.kotlin.name.ClassId
12+
import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
13+
14+
/**
15+
* Returns whether class is coming from JVM runtime env. There is no need to track these classes.
16+
*
17+
* @param className the class name of the class
18+
* @return whether class is provided by JSM runtime or not
19+
*/
20+
internal fun isJvmClass(className: String): Boolean {
21+
return className.startsWith("java") ||
22+
className.startsWith("modules/java.base/java/")
23+
}
24+
25+
internal fun DeserializedContainerSource.classId(): ClassId? {
26+
return when (this) {
27+
is JvmPackagePartSource -> classId
28+
is KotlinJvmBinarySourceElement -> binaryClass.classId
29+
else -> null
30+
}
31+
}
32+
33+
internal fun SourceElement.binaryClass(): String? {
34+
return when (this) {
35+
is KotlinJvmBinarySourceElement -> binaryClass.location
36+
is JvmPackagePartSource -> this.knownJvmBinaryClass?.location
37+
is JavaBinarySourceElement -> this.javaClass.virtualFile.path
38+
else -> null
39+
}
40+
}
41+
42+
internal fun DeserializedContainerSource.binaryClass(): String? {
43+
return when (this) {
44+
is JvmPackagePartSource -> this.knownJvmBinaryClass?.location
45+
is KotlinJvmBinarySourceElement -> binaryClass.location
46+
else -> null
47+
}
48+
}
49+
50+
// Extension functions to clean up checker logic
51+
52+
internal fun FirClassLikeSymbol<*>.recordClass(
53+
context: CheckerContext,
54+
isExplicit: Boolean = true,
55+
collectTypeArguments: Boolean = true,
56+
visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
57+
) {
58+
ClassUsageRecorder.recordClass(this, context, isExplicit, collectTypeArguments, visited)
59+
}
60+
61+
internal fun FirTypeRef.recordTypeRef(
62+
context: CheckerContext,
63+
isExplicit: Boolean = true,
64+
collectTypeArguments: Boolean = true,
65+
visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
66+
) {
67+
ClassUsageRecorder.recordTypeRef(this, context, isExplicit, collectTypeArguments, visited)
68+
}
69+
70+
internal fun ConeKotlinType.recordConeType(
71+
context: CheckerContext,
72+
isExplicit: Boolean = true,
73+
collectTypeArguments: Boolean = true,
74+
) {
75+
ClassUsageRecorder.recordConeType(this, context, isExplicit, collectTypeArguments)
76+
}
77+
78+
internal fun String.recordClassBinaryPath(isExplicit: Boolean = true) {
79+
ClassUsageRecorder.recordClass(this, isExplicit)
80+
}

‎src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/checker/declaration/BasicDeclarationChecker.kt

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
2+
3+
import io.bazel.kotlin.plugin.jdeps.k2.recordClass
4+
import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
5+
import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
6+
import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
7+
import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirBasicDeclarationChecker
8+
import org.jetbrains.kotlin.fir.declarations.FirDeclaration
9+
import org.jetbrains.kotlin.fir.declarations.toAnnotationClassLikeSymbol
10+
11+
internal object BasicDeclarationChecker : FirBasicDeclarationChecker(MppCheckerKind.Common) {
12+
override fun check(
13+
declaration: FirDeclaration,
14+
context: CheckerContext,
15+
reporter: DiagnosticReporter,
16+
) {
17+
declaration.annotations.forEach { annotation ->
18+
annotation.toAnnotationClassLikeSymbol(context.session)?.recordClass(context)
19+
}
20+
}
21+
}

0 commit comments

Comments
 (0)