Restore per-class and android resource tracking that got lost with K2

Author: oliviernotteghem

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

@@ -1,28 +1,47 @@
 package io.bazel.kotlin.plugin.jdeps
 
+import com.google.common.io.ByteStreams
 import com.google.devtools.build.lib.view.proto.Deps
+import com.google.protobuf.ByteString
 import io.bazel.kotlin.builder.utils.jars.JarOwner
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import java.io.BufferedOutputStream
 import java.io.File
 import java.nio.file.Paths
+import java.security.MessageDigest
+import java.util.jar.JarFile
+import java.util.stream.Collectors
 
 abstract class BaseJdepsGenExtension(
   protected val configuration: CompilerConfiguration,
 ) {
   protected fun onAnalysisCompleted(
     explicitClassesCanonicalPaths: Set<String>,
     implicitClassesCanonicalPaths: Set<String>,
+    usedResources: Set<String>,
   ) {
     val directDeps = configuration.getList(JdepsGenConfigurationKeys.DIRECT_DEPENDENCIES)
     val targetLabel = configuration.getNotNull(JdepsGenConfigurationKeys.TARGET_LABEL)
     val explicitDeps = createDepsMap(explicitClassesCanonicalPaths)
 
-    doWriteJdeps(directDeps, targetLabel, explicitDeps, implicitClassesCanonicalPaths)
+    doWriteJdeps(directDeps, targetLabel, explicitDeps, implicitClassesCanonicalPaths, usedResources)
 
     doStrictDeps(configuration, targetLabel, directDeps, explicitDeps)
   }
 
+  /**
+   * Compute hash of internal jar class ABI definition.
+   */
+  protected fun getHashFromJarEntry(
+    jarPath: String,
+    internalPath: String,
+  ): ByteArray {
+    val jarFile = JarFile(jarPath)
+    val entry = jarFile.getEntry(internalPath)
+    val bytes = ByteStreams.toByteArray(jarFile.getInputStream(entry))
+    return MessageDigest.getInstance("SHA-256").digest(bytes)
+  }
+
   /**
    * Returns a map of jars to classes loaded from those jars.
    */
@@ -43,7 +62,10 @@ abstract class BaseJdepsGenExtension(
     targetLabel: String,
     explicitDeps: Map<String, List<String>>,
     implicitClassesCanonicalPaths: Set<String>,
+    usedResources: Set<String>,
   ) {
+    val trackClassUsage = configuration.getNotNull(JdepsGenConfigurationKeys.TRACK_CLASS_USAGE).equals("on")
+    val trackResourceUsage = configuration.getNotNull(JdepsGenConfigurationKeys.TRACK_RESOURCE_USAGE).equals("on")
     val implicitDeps = createDepsMap(implicitClassesCanonicalPaths)
 
     // Build and write out deps.proto
@@ -53,10 +75,26 @@ abstract class BaseJdepsGenExtension(
     rootBuilder.success = true
     rootBuilder.ruleLabel = targetLabel
 
-    explicitDeps.forEach { (jarPath, _) ->
+    explicitDeps.toSortedMap().forEach { (jarPath, usedClasses) ->
       val dependency = Deps.Dependency.newBuilder()
       dependency.kind = Deps.Dependency.Kind.EXPLICIT
       dependency.path = jarPath
+
+      if (trackClassUsage) {
+        // Add tracked classes and their (compile time) hash into final output, as needed for
+        // compilation avoidance.
+        usedClasses.stream().sorted().collect(Collectors.toList()).forEach { it ->
+          val name = it.replace(".class", "").replace("/", ".")
+          val hash = ByteString.copyFrom(getHashFromJarEntry(jarPath, it))
+          val usedClass: Deps.UsedClass = Deps.UsedClass.newBuilder()
+            .setFullyQualifiedName(name)
+            .setInternalPath(it)
+            .setHash(hash)
+            .build()
+          dependency.addUsedClasses(usedClass)
+        }
+      }
+
       rootBuilder.addDependency(dependency)
     }
 
@@ -67,6 +105,12 @@ abstract class BaseJdepsGenExtension(
       rootBuilder.addDependency(dependency)
     }
 
+    if (trackResourceUsage) {
+      usedResources.sorted().forEach { resource ->
+        rootBuilder.addUsedResources(resource)
+      }
+    }
+
     BufferedOutputStream(File(jdepsOutput).outputStream()).use {
       it.write(rootBuilder.buildSorted().toByteArray())
     }

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

@@ -377,7 +377,7 @@ class JdepsGenExtension(
     bindingTrace: BindingTrace,
     files: Collection<KtFile>,
   ): AnalysisResult? {
-    onAnalysisCompleted(explicitClassesCanonicalPaths, implicitClassesCanonicalPaths)
+    onAnalysisCompleted(explicitClassesCanonicalPaths, implicitClassesCanonicalPaths, usedResources)
 
     return super.analysisCompleted(project, module, bindingTrace, files)
   }
@@ -454,19 +454,6 @@ class JdepsGenExtension(
     }
   }
 
-  /**
-   * Compute hash of internal jar class ABI definition.
-   */
-  private fun getHashFromJarEntry(
-    jarPath: String,
-    internalPath: String,
-  ): ByteArray {
-    val jarFile = JarFile(jarPath)
-    val entry = jarFile.getEntry(internalPath)
-    val bytes = ByteStreams.toByteArray(jarFile.getInputStream(entry))
-    return MessageDigest.getInstance("SHA-256").digest(bytes)
-  }
-
   private fun doStrictDeps(
     compilerConfiguration: CompilerConfiguration,
     targetLabel: String,

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

@@ -20,6 +20,7 @@ private const val ANONYMOUS = "<anonymous>"
 class ClassUsageRecorder(
   internal val explicitClassesCanonicalPaths: MutableSet<String> = mutableSetOf(),
   internal val implicitClassesCanonicalPaths: MutableSet<String> = mutableSetOf(),
+  internal val usedResources: 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() + "/",
@@ -99,6 +100,10 @@ class ClassUsageRecorder(
     addFile(binaryClass, isExplicit)
   }
 
+  internal fun recordResource(resourceId: String) {
+    usedResources.add(resourceId)
+  }
+
   private fun addFile(
     path: String,
     isExplicit: Boolean,

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

@@ -14,6 +14,7 @@ internal class JdepsGenExtension2(
     onAnalysisCompleted(
       classUsageRecorder.explicitClassesCanonicalPaths,
       classUsageRecorder.implicitClassesCanonicalPaths,
+      classUsageRecorder.usedResources,
     )
   }
 }

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

@@ -1,6 +1,8 @@
 package io.bazel.kotlin.plugin.jdeps.k2
 
 import org.jetbrains.kotlin.descriptors.SourceElement
+import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.java.JavaBinarySourceElement
 import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement

src/main/kotlin/io/bazel/kotlin/plugin/jdeps/k2/checker/expression/QualifiedAccessChecker.kt

@@ -8,13 +8,15 @@ import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirQualifiedAccessExpressionChecker
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.arguments
 import org.jetbrains.kotlin.fir.expressions.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.references.toResolvedFunctionSymbol
 import org.jetbrains.kotlin.fir.types.isExtensionFunctionType
 import org.jetbrains.kotlin.fir.types.isUnit
 import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.KtSourceElement
 
 internal class QualifiedAccessChecker(
   private val classUsageRecorder: ClassUsageRecorder,
@@ -66,5 +68,17 @@ internal class QualifiedAccessChecker(
         classUsageRecorder.recordConeType(it, context, isExplicit = false)
       }
     }
+
+    // track android resources
+    getResourceName(expression)?.let {
+      classUsageRecorder.recordResource(it)
+    }
+  }
+
+  private fun getResourceName(expression: FirQualifiedAccessExpression): String? {
+    return (expression as? FirPropertyAccessExpression)
+      ?.source
+      ?.getElementTextInContextForDebug()
+      ?.takeIf { ".R." in it || it.startsWith("R.") }
   }
 }