Convert ClassUsageRecorder & checker objects to classes

Author: jbarr21

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

@@ -34,9 +34,9 @@ class JdepsGenComponentRegistrar : CompilerPluginRegistrar() {
 
   private fun ExtensionStorage.registerForK2(configuration: CompilerConfiguration) {
     val projectRoot = Paths.get("").toAbsolutePath().toString() + "/"
-    ClassUsageRecorder.init(projectRoot)
-    JdepsGenExtension2(configuration).run {
-      FirExtensionRegistrarAdapter.registerExtension(JdepsFirExtensions())
+    val classUsageRecorder = ClassUsageRecorder(rootPath = projectRoot)
+    JdepsGenExtension2(classUsageRecorder, configuration).run {
+      FirExtensionRegistrarAdapter.registerExtension(JdepsFirExtensions(classUsageRecorder))
       ClassFileFactoryFinalizerExtension.registerExtension(this)
     }
   }

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

@@ -17,26 +17,14 @@ 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>()
+class ClassUsageRecorder(
+  internal val explicitClassesCanonicalPaths: MutableSet<String> = mutableSetOf(),
+  internal val implicitClassesCanonicalPaths: MutableSet<String> = mutableSetOf(),
+  private val seen: MutableSet<ClassId> = mutableSetOf(),
+  private val results: MutableMap<String, SortedSet<String>> = sortedMapOf(),
+  private val rootPath: String = Paths.get("").toAbsolutePath().toString() + "/",
+) {
   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,
@@ -61,20 +49,22 @@ internal object ClassUsageRecorder {
         if (ANONYMOUS in classId.toString()) return@forEachType
         context.session.symbolProvider
           .getClassLikeSymbolByClassId(classId)
-          ?.recordClass(context, isExplicit, collectTypeArguments, visited)
+          ?.let { recordClass(it, context, isExplicit, collectTypeArguments, visited) }
       }
     } else {
       coneKotlinType.classId?.let { classId ->
-        context.session.symbolProvider.getClassLikeSymbolByClassId(classId)
-          ?.recordClass(context, isExplicit, collectTypeArguments, visited)
+        if (!classId.isLocal) {
+          context.session.symbolProvider.getClassLikeSymbolByClassId(classId)
+            ?.let { recordClass(it, context, isExplicit, collectTypeArguments, visited) }
+        }
       }
     }
   }
 
   internal fun recordClass(
     firClass: FirClassLikeSymbol<*>,
     context: CheckerContext,
-    isExplicit: Boolean,
+    isExplicit: Boolean = true,
     collectTypeArguments: Boolean = true,
     visited: MutableSet<Pair<ClassId, Boolean>> = mutableSetOf(),
   ) {

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

@@ -20,35 +20,40 @@ import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirResolvedQualifie
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 
-internal class JdepsFirExtensions : FirExtensionRegistrar() {
+internal class JdepsFirExtensions(
+  private val classUsageRecorder: ClassUsageRecorder,
+) : 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)
+  internal inner class JdepsFirCheckersExtension(session: FirSession) :
+    FirAdditionalCheckersExtension(session) {
+    override val declarationCheckers: DeclarationCheckers =
+      object : DeclarationCheckers() {
+        override val basicDeclarationCheckers: Set<FirBasicDeclarationChecker> =
+          setOf(BasicDeclarationChecker(classUsageRecorder))
 
-      override val fileCheckers: Set<FirFileChecker> = setOf(FileChecker)
+        override val fileCheckers: Set<FirFileChecker> =
+          setOf(FileChecker(classUsageRecorder))
 
-      override val classLikeCheckers: Set<FirClassLikeChecker> = setOf(ClassLikeChecker)
+        override val classLikeCheckers: Set<FirClassLikeChecker> =
+          setOf(ClassLikeChecker(classUsageRecorder))
 
-      override val functionCheckers: Set<FirFunctionChecker> = setOf(FunctionChecker)
+        override val functionCheckers: Set<FirFunctionChecker> =
+          setOf(FunctionChecker(classUsageRecorder))
 
-      override val callableDeclarationCheckers: Set<FirCallableDeclarationChecker> =
-        setOf(CallableChecker)
-    }
+        override val callableDeclarationCheckers: Set<FirCallableDeclarationChecker> =
+          setOf(CallableChecker(classUsageRecorder))
+      }
 
-  override val expressionCheckers: ExpressionCheckers =
-    object : ExpressionCheckers() {
-      override val qualifiedAccessExpressionCheckers: Set<FirQualifiedAccessExpressionChecker> =
-        setOf(QualifiedAccessChecker)
+    override val expressionCheckers: ExpressionCheckers =
+      object : ExpressionCheckers() {
+        override val qualifiedAccessExpressionCheckers: Set<FirQualifiedAccessExpressionChecker> =
+          setOf(QualifiedAccessChecker(classUsageRecorder))
 
-      override val resolvedQualifierCheckers: Set<FirResolvedQualifierChecker> =
-        setOf(ResolvedQualifierChecker)
-    }
+        override val resolvedQualifierCheckers: Set<FirResolvedQualifierChecker> =
+          setOf(ResolvedQualifierChecker(classUsageRecorder))
+      }
+  }
 }

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

@@ -6,12 +6,13 @@ import org.jetbrains.kotlin.codegen.extensions.ClassFileFactoryFinalizerExtensio
 import org.jetbrains.kotlin.config.CompilerConfiguration
 
 internal class JdepsGenExtension2(
+  private val classUsageRecorder: ClassUsageRecorder,
   configuration: CompilerConfiguration,
 ) : BaseJdepsGenExtension(configuration), ClassFileFactoryFinalizerExtension {
   override fun finalizeClassFactory(factory: ClassFileFactory) {
     onAnalysisCompleted(
-      ClassUsageRecorder.getExplicitClassesCanonicalPaths(),
-      ClassUsageRecorder.getImplicitClassesCanonicalPaths(),
+      classUsageRecorder.explicitClassesCanonicalPaths,
+      classUsageRecorder.implicitClassesCanonicalPaths,
     )
   }
 }

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

@@ -1,11 +1,7 @@
 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
@@ -46,35 +42,3 @@ internal fun DeserializedContainerSource.binaryClass(): String? {
     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

@@ -1,21 +1,25 @@
 package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
 
-import io.bazel.kotlin.plugin.jdeps.k2.recordClass
+import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
 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) {
+internal class BasicDeclarationChecker(
+  private val classUsageRecorder: ClassUsageRecorder,
+) : FirBasicDeclarationChecker(MppCheckerKind.Common) {
   override fun check(
     declaration: FirDeclaration,
     context: CheckerContext,
     reporter: DiagnosticReporter,
   ) {
     declaration.annotations.forEach { annotation ->
-      annotation.toAnnotationClassLikeSymbol(context.session)?.recordClass(context)
+      annotation.toAnnotationClassLikeSymbol(context.session)?.let {
+        classUsageRecorder.recordClass(it, context)
+      }
     }
   }
 }

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

@@ -1,6 +1,6 @@
 package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
 
-import io.bazel.kotlin.plugin.jdeps.k2.recordTypeRef
+import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
@@ -9,7 +9,9 @@ import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.utils.isExtension
 
-internal object CallableChecker : FirCallableDeclarationChecker(MppCheckerKind.Common) {
+internal class CallableChecker(
+  private val classUsageRecorder: ClassUsageRecorder,
+) : FirCallableDeclarationChecker(MppCheckerKind.Common) {
   /**
    * Tracks the return type & type parameters of a callable declaration. Function parameters are
    * tracked in [FunctionChecker].
@@ -20,19 +22,20 @@ internal object CallableChecker : FirCallableDeclarationChecker(MppCheckerKind.C
     reporter: DiagnosticReporter,
   ) {
     // return type
-    declaration.returnTypeRef.recordTypeRef(context)
+    declaration.returnTypeRef.let { classUsageRecorder.recordTypeRef(it, context) }
 
     // type params
     declaration.typeParameters.forEach { typeParam ->
       typeParam.symbol.resolvedBounds.forEach { typeParamBound ->
-        typeParamBound.recordTypeRef(context)
+        typeParamBound.let { classUsageRecorder.recordTypeRef(it, context) }
       }
     }
 
     // receiver param for extensions
     if (declaration !is FirAnonymousFunction) {
-      declaration.receiverParameter?.typeRef
-        ?.recordTypeRef(context, isExplicit = declaration.isExtension)
+      declaration.receiverParameter?.typeRef?.let {
+        classUsageRecorder.recordTypeRef(it, context, isExplicit = declaration.isExtension)
+      }
     }
   }
 }

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

@@ -1,24 +1,25 @@
 package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
 
-import io.bazel.kotlin.plugin.jdeps.k2.recordClass
-import io.bazel.kotlin.plugin.jdeps.k2.recordConeType
+import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
 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.FirClassLikeChecker
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.resolve.getSuperTypes
 
-internal object ClassLikeChecker : FirClassLikeChecker(MppCheckerKind.Common) {
+internal class ClassLikeChecker(
+  private val classUsageRecorder: ClassUsageRecorder,
+) : FirClassLikeChecker(MppCheckerKind.Common) {
   override fun check(
     declaration: FirClassLikeDeclaration,
     context: CheckerContext,
     reporter: DiagnosticReporter,
   ) {
-    declaration.symbol.recordClass(context)
+    declaration.symbol.let { classUsageRecorder.recordClass(it, context) }
     // [recordClass] also handles supertypes, but this marks direct supertypes as explicit
     declaration.symbol.getSuperTypes(context.session, recursive = false).forEach {
-      it.recordConeType(context)
+      classUsageRecorder.recordConeType(it, context)
     }
   }
 }

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

@@ -2,7 +2,6 @@ package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
 
 import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
 import io.bazel.kotlin.plugin.jdeps.k2.binaryClass
-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
@@ -18,7 +17,9 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeAliasSymbol
 import org.jetbrains.kotlin.name.ClassId
 
-internal object FileChecker : FirFileChecker(MppCheckerKind.Common) {
+internal class FileChecker(
+  private val classUsageRecorder: ClassUsageRecorder,
+) : FirFileChecker(MppCheckerKind.Common) {
   override fun check(
     declaration: FirFile,
     context: CheckerContext,
@@ -27,15 +28,19 @@ internal object FileChecker : FirFileChecker(MppCheckerKind.Common) {
     declaration.imports.filterIsInstance<FirResolvedImport>().forEach { import ->
       // check for classlike import (class, interface, object, enum, annotation, etc)
       if (import.resolvesToClass(context)) {
-        import.classId()?.resolveToClass(context)?.recordClass(context)
+        import.classId()?.resolveToClass(context)?.let {
+          classUsageRecorder.recordClass(it, context)
+        }
       } else {
         // check for function import
         val callableBinaryClass = import.resolveToFun(context)?.containerSource?.binaryClass()
         if (callableBinaryClass != null) {
-          ClassUsageRecorder.recordClass(callableBinaryClass)
+          classUsageRecorder.recordClass(callableBinaryClass)
         } else {
           // for other symbols, track the parent class
-          import.resolvedParentClassId?.resolveToClass(context)?.recordClass(context)
+          import.resolvedParentClassId?.resolveToClass(context)?.let {
+            classUsageRecorder.recordClass(it, context)
+          }
         }
       }
     }

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

@@ -1,13 +1,15 @@
 package io.bazel.kotlin.plugin.jdeps.k2.checker.declaration
 
-import io.bazel.kotlin.plugin.jdeps.k2.recordTypeRef
+import io.bazel.kotlin.plugin.jdeps.k2.ClassUsageRecorder
 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.FirFunctionChecker
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 
-internal object FunctionChecker : FirFunctionChecker(MppCheckerKind.Common) {
+internal class FunctionChecker(
+  private val classUsageRecorder: ClassUsageRecorder,
+) : FirFunctionChecker(MppCheckerKind.Common) {
   /**
    * Tracks the value parameters of a function declaration. Return type & type parameters are
    * tracked in [CallableChecker].
@@ -19,7 +21,7 @@ internal object FunctionChecker : FirFunctionChecker(MppCheckerKind.Common) {
   ) {
     // function parameters
     declaration.valueParameters.forEach { valueParam ->
-      valueParam.returnTypeRef.recordTypeRef(context)
+      valueParam.returnTypeRef.let { classUsageRecorder.recordTypeRef(it, context) }
     }
   }
 }