Implement jdeps using K2 APIs

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

@@ -19,7 +19,12 @@ load("//src/main/kotlin:bootstrap.bzl", "kt_bootstrap_library")
 # The compiler binary, this is co-located in the kotlin compiler classloader.
 kt_bootstrap_library(
     name = "jdeps-gen-lib",
-    srcs = glob(["*.kt"]),
+    srcs = glob([
+        "*.kt",
+        "k2/*.kt",
+        "k2/checker/declaration/*.kt",
+        "k2/checker/expression/*.kt",
+    ]),
     visibility = ["//src:__subpackages__"],
     deps = [
         "//kotlin/compiler:kotlin-compiler",

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

@@ -10,7 +10,6 @@ import java.nio.file.Paths
 abstract class BaseJdepsGenExtension(
   protected val configuration: CompilerConfiguration,
 ) {
-
   protected fun onAnalysisCompleted(
     explicitClassesCanonicalPaths: Set<String>,
     implicitClassesCanonicalPaths: Set<String>,
@@ -107,9 +106,10 @@ abstract class BaseJdepsGenExtension(
     directDeps: List<String>,
     targetLabel: String,
   ): Boolean {
-    val missingStrictDeps = result.keys
-      .filter { !directDeps.contains(it) }
-      .map { JarOwner.readJarOwnerFromManifest(Paths.get(it)) }
+    val missingStrictDeps =
+      result.keys
+        .filter { !directDeps.contains(it) }
+        .map { JarOwner.readJarOwnerFromManifest(Paths.get(it)) }
 
     if (missingStrictDeps.isNotEmpty()) {
       val missingStrictLabels = missingStrictDeps.mapNotNull { it.label }

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

@@ -1,21 +1,43 @@
 package io.bazel.kotlin.plugin.jdeps
 
+import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
+import io.bazel.kotlin.plugin.jdeps.k2.JdepsFirExtensions
+import io.bazel.kotlin.plugin.jdeps.k2.JdepsGenExtension2
+import org.jetbrains.kotlin.codegen.extensions.ClassFileFactoryFinalizerExtension
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
+import java.nio.file.Paths
 
 @OptIn(org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi::class)
 class JdepsGenComponentRegistrar : CompilerPluginRegistrar() {
-
   override val supportsK2: Boolean
-    get() = false
+    get() = true
 
   override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
+    when (configuration.languageVersionSettings.languageVersion.usesK2) {
+      true -> registerForK2(configuration)
+      false -> registerForK1(configuration)
+    }
+  }
+
+  private fun ExtensionStorage.registerForK1(configuration: CompilerConfiguration) {
     // Capture all types referenced by the compiler for this module and look up the jar from which
     // they were loaded from
     val extension = JdepsGenExtension(configuration)
     AnalysisHandlerExtension.registerExtension(extension)
     StorageComponentContainerContributor.registerExtension(extension)
   }
+
+  private fun ExtensionStorage.registerForK2(configuration: CompilerConfiguration) {
+    val projectRoot = Paths.get("").toAbsolutePath().toString() + "/"
+    ClassUsageRecorder.init(projectRoot)
+    JdepsGenExtension2(configuration).run {
+      FirExtensionRegistrarAdapter.registerExtension(JdepsFirExtensions())
+      ClassFileFactoryFinalizerExtension.registerExtension(this)
+    }
+  }
 }

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

@@ -0,0 +1,132 @@
+package io.bazel.kotlin.plugin.jdeps.k2
+
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.declarations.utils.sourceElement
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.forEachType
+import org.jetbrains.kotlin.name.ClassId
+import java.nio.file.Paths
+import java.util.SortedSet
+
+private const val JAR_FILE_SEPARATOR = "!/"
+private const val ANONYMOUS = "<anonymous>"
+
+internal object ClassUsageRecorder {
+  private val explicitClassesCanonicalPaths = mutableSetOf<String>()
+  private val implicitClassesCanonicalPaths = mutableSetOf<String>()
+
+  private val results by lazy { sortedMapOf<String, SortedSet<String>>() }
+  private val seen = mutableSetOf<ClassId>()
+  private val javaHome: String by lazy { System.getenv()["JAVA_HOME"] ?: "<not set>" }
+  private var rootPath: String = Paths.get("").toAbsolutePath().toString() + "/"
+
+  fun init(rootPath: String) {
+    this.rootPath = rootPath
+    results.clear()
+    explicitClassesCanonicalPaths.clear()
+    implicitClassesCanonicalPaths.clear()
+    seen.clear()
+  }
+
+  fun getExplicitClassesCanonicalPaths(): Set<String> = explicitClassesCanonicalPaths
+
+  fun getImplicitClassesCanonicalPaths(): Set<String> = implicitClassesCanonicalPaths
+
+  internal fun recordTypeRef(
+    typeRef: FirTypeRef,
+    context: CheckerContext,
+    isExplicit: Boolean = true,
+    collectTypeArguments: Boolean = true,
+    visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
+  ) {
+    recordConeType(typeRef.coneType, context, isExplicit, collectTypeArguments, visited)
+  }
+
+  internal fun recordConeType(
+    coneKotlinType: ConeKotlinType,
+    context: CheckerContext,
+    isExplicit: Boolean = true,
+    collectTypeArguments: Boolean = true,
+    visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
+  ) {
+    if (collectTypeArguments) {
+      coneKotlinType.forEachType { coneType ->
+        val classId = coneType.classId ?: return@forEachType
+        if (ANONYMOUS in classId.toString()) return@forEachType
+        context.session.symbolProvider
+          .getClassLikeSymbolByClassId(classId)
+          ?.recordClass(context, isExplicit, collectTypeArguments, visited)
+      }
+    } else {
+      coneKotlinType.classId?.let { classId ->
+        context.session.symbolProvider.getClassLikeSymbolByClassId(classId)
+          ?.recordClass(context, isExplicit, collectTypeArguments, visited)
+      }
+    }
+  }
+
+  internal fun recordClass(
+    firClass: FirClassLikeSymbol<*>,
+    context: CheckerContext,
+    isExplicit: Boolean,
+    collectTypeArguments: Boolean = true,
+    visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
+  ) {
+    val classIdAndIsExplicit = firClass.classId to isExplicit
+    if (classIdAndIsExplicit in visited) {
+      return
+    } else {
+      visited.add(classIdAndIsExplicit)
+    }
+
+    firClass.sourceElement?.binaryClass()?.let { addFile(it, isExplicit) }
+
+    if (firClass is FirClassSymbol<*>) {
+      firClass.resolvedSuperTypeRefs.forEach {
+        recordTypeRef(it, context, false, collectTypeArguments, visited)
+      }
+      if (collectTypeArguments) {
+        firClass.typeParameterSymbols.flatMap { it.resolvedBounds }
+          .forEach { recordTypeRef(it, context, isExplicit, collectTypeArguments, visited) }
+      }
+    }
+  }
+
+  internal fun recordClass(
+    binaryClass: String,
+    isExplicit: Boolean = true,
+  ) {
+    addFile(binaryClass, isExplicit)
+  }
+
+  private fun addFile(
+    path: String,
+    isExplicit: Boolean,
+  ) {
+    if (isExplicit) {
+      explicitClassesCanonicalPaths.add(path)
+    } else {
+      implicitClassesCanonicalPaths.add(path)
+    }
+
+    if (path.contains(JAR_FILE_SEPARATOR) && !path.contains(javaHome)) {
+      val (jarPath, classPath) = path.split(JAR_FILE_SEPARATOR)
+      // Convert jar files in current directory to relative paths. Remaining absolute are outside
+      // of project and should be ignored
+      val relativizedJarPath = Paths.get(jarPath.replace(rootPath, ""))
+      if (!relativizedJarPath.isAbsolute) {
+        val occurrences =
+          results.computeIfAbsent(relativizedJarPath.toString()) { sortedSetOf<String>() }
+        if (!isJvmClass(classPath)) {
+          occurrences.add(classPath)
+        }
+      }
+    }
+  }
+}

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

@@ -0,0 +1,54 @@
+package io.bazel.kotlin.plugin.jdeps.k2
+
+import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.BasicDeclarationChecker
+import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.CallableChecker
+import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.ClassLikeChecker
+import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.FileChecker
+import io.bazel.kotlin.plugin.jdeps.k2.checker.declaration.FunctionChecker
+import io.bazel.kotlin.plugin.jdeps.k2.checker.expression.QualifiedAccessChecker
+import io.bazel.kotlin.plugin.jdeps.k2.checker.expression.ResolvedQualifierChecker
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirBasicDeclarationChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirCallableDeclarationChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirClassLikeChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirFileChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirFunctionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirQualifiedAccessExpressionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirResolvedQualifierChecker
+import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+
+internal class JdepsFirExtensions : FirExtensionRegistrar() {
+  override fun ExtensionRegistrarContext.configurePlugin() {
+    +::JdepsFirCheckersExtension
+  }
+}
+
+internal class JdepsFirCheckersExtension(session: FirSession) :
+  FirAdditionalCheckersExtension(session) {
+  override val declarationCheckers: DeclarationCheckers =
+    object : DeclarationCheckers() {
+      override val basicDeclarationCheckers: Set<FirBasicDeclarationChecker> =
+        setOf(BasicDeclarationChecker)
+
+      override val fileCheckers: Set<FirFileChecker> = setOf(FileChecker)
+
+      override val classLikeCheckers: Set<FirClassLikeChecker> = setOf(ClassLikeChecker)
+
+      override val functionCheckers: Set<FirFunctionChecker> = setOf(FunctionChecker)
+
+      override val callableDeclarationCheckers: Set<FirCallableDeclarationChecker> =
+        setOf(CallableChecker)
+    }
+
+  override val expressionCheckers: ExpressionCheckers =
+    object : ExpressionCheckers() {
+      override val qualifiedAccessExpressionCheckers: Set<FirQualifiedAccessExpressionChecker> =
+        setOf(QualifiedAccessChecker)
+
+      override val resolvedQualifierCheckers: Set<FirResolvedQualifierChecker> =
+        setOf(ResolvedQualifierChecker)
+    }
+}

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

@@ -0,0 +1,17 @@
+package io.bazel.kotlin.plugin.jdeps.k2
+
+import io.bazel.kotlin.plugin.jdeps.BaseJdepsGenExtension
+import org.jetbrains.kotlin.codegen.ClassFileFactory
+import org.jetbrains.kotlin.codegen.extensions.ClassFileFactoryFinalizerExtension
+import org.jetbrains.kotlin.config.CompilerConfiguration
+
+internal class JdepsGenExtension2(
+  configuration: CompilerConfiguration,
+) : BaseJdepsGenExtension(configuration), ClassFileFactoryFinalizerExtension {
+  override fun finalizeClassFactory(factory: ClassFileFactory) {
+    onAnalysisCompleted(
+      ClassUsageRecorder.getExplicitClassesCanonicalPaths(),
+      ClassUsageRecorder.getImplicitClassesCanonicalPaths(),
+    )
+  }
+}

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

@@ -0,0 +1,80 @@
+package io.bazel.kotlin.plugin.jdeps.k2
+
+import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.java.JavaBinarySourceElement
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
+import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
+
+/**
+ * Returns whether class is coming from JVM runtime env. There is no need to track these classes.
+ *
+ * @param className the class name of the class
+ * @return whether class is provided by JSM runtime or not
+ */
+internal fun isJvmClass(className: String): Boolean {
+  return className.startsWith("java") ||
+    className.startsWith("modules/java.base/java/")
+}
+
+internal fun DeserializedContainerSource.classId(): ClassId? {
+  return when (this) {
+    is JvmPackagePartSource -> classId
+    is KotlinJvmBinarySourceElement -> binaryClass.classId
+    else -> null
+  }
+}
+
+internal fun SourceElement.binaryClass(): String? {
+  return when (this) {
+    is KotlinJvmBinarySourceElement -> binaryClass.location
+    is JvmPackagePartSource -> this.knownJvmBinaryClass?.location
+    is JavaBinarySourceElement -> this.javaClass.virtualFile.path
+    else -> null
+  }
+}
+
+internal fun DeserializedContainerSource.binaryClass(): String? {
+  return when (this) {
+    is JvmPackagePartSource -> this.knownJvmBinaryClass?.location
+    is KotlinJvmBinarySourceElement -> binaryClass.location
+    else -> null
+  }
+}
+
+// Extension functions to clean up checker logic
+
+internal fun FirClassLikeSymbol<*>.recordClass(
+  context: CheckerContext,
+  isExplicit: Boolean = true,
+  collectTypeArguments: Boolean = true,
+  visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
+) {
+  ClassUsageRecorder.recordClass(this, context, isExplicit, collectTypeArguments, visited)
+}
+
+internal fun FirTypeRef.recordTypeRef(
+  context: CheckerContext,
+  isExplicit: Boolean = true,
+  collectTypeArguments: Boolean = true,
+  visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
+) {
+  ClassUsageRecorder.recordTypeRef(this, context, isExplicit, collectTypeArguments, visited)
+}
+
+internal fun ConeKotlinType.recordConeType(
+  context: CheckerContext,
+  isExplicit: Boolean = true,
+  collectTypeArguments: Boolean = true,
+) {
+  ClassUsageRecorder.recordConeType(this, context, isExplicit, collectTypeArguments)
+}
+
+internal fun String.recordClassBinaryPath(isExplicit: Boolean = true) {
+  ClassUsageRecorder.recordClass(this, isExplicit)
+}

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

@@ -0,0 +1,21 @@
+package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
+
+import io.bazel.kotlin.plugin.jdeps.k2.recordClass
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirBasicDeclarationChecker
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.toAnnotationClassLikeSymbol
+
+internal object BasicDeclarationChecker : FirBasicDeclarationChecker(MppCheckerKind.Common) {
+  override fun check(
+    declaration: FirDeclaration,
+    context: CheckerContext,
+    reporter: DiagnosticReporter,
+  ) {
+    declaration.annotations.forEach { annotation ->
+      annotation.toAnnotationClassLikeSymbol(context.session)?.recordClass(context)
+    }
+  }
+}