Skip to content

Commit

Permalink
Prepare Jdeps extension for K2 implementation (bazelbuild#1164)
Browse files Browse the repository at this point in the history
* Fix K2 jdeps test args
* Sort Deps in jdeps

* Refactor jdeps extension to prepare for K2 impl
  • Loading branch information
jbarr21 authored and oliviernotteghem committed Aug 2, 2024
1 parent 7044c0d commit 8da989d
Show file tree
Hide file tree
Showing 4 changed files with 389 additions and 219 deletions.
149 changes: 149 additions & 0 deletions src/main/kotlin/io/bazel/kotlin/plugin/jdeps/BaseJdepsGenExtension.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
package io.bazel.kotlin.plugin.jdeps

import com.google.devtools.build.lib.view.proto.Deps
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

abstract class BaseJdepsGenExtension(
protected val configuration: CompilerConfiguration,
) {

protected fun onAnalysisCompleted(
explicitClassesCanonicalPaths: Set<String>,
implicitClassesCanonicalPaths: 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)

doStrictDeps(configuration, targetLabel, directDeps, explicitDeps)
}

/**
* Returns a map of jars to classes loaded from those jars.
*/
private fun createDepsMap(classes: Set<String>): Map<String, List<String>> {
val jarsToClasses = mutableMapOf<String, MutableList<String>>()
classes.forEach {
val parts = it.split("!/")
val jarPath = parts[0]
if (jarPath.endsWith(".jar")) {
jarsToClasses.computeIfAbsent(jarPath) { ArrayList() }.add(parts[1])
}
}
return jarsToClasses
}

private fun doWriteJdeps(
directDeps: MutableList<String>,
targetLabel: String,
explicitDeps: Map<String, List<String>>,
implicitClassesCanonicalPaths: Set<String>,
) {
val implicitDeps = createDepsMap(implicitClassesCanonicalPaths)

// Build and write out deps.proto
val jdepsOutput = configuration.getNotNull(JdepsGenConfigurationKeys.OUTPUT_JDEPS)

val rootBuilder = Deps.Dependencies.newBuilder()
rootBuilder.success = true
rootBuilder.ruleLabel = targetLabel

val unusedDeps = directDeps.subtract(explicitDeps.keys)
unusedDeps.forEach { jarPath ->
val dependency = Deps.Dependency.newBuilder()
dependency.kind = Deps.Dependency.Kind.UNUSED
dependency.path = jarPath
rootBuilder.addDependency(dependency)
}

explicitDeps.forEach { (jarPath, _) ->
val dependency = Deps.Dependency.newBuilder()
dependency.kind = Deps.Dependency.Kind.EXPLICIT
dependency.path = jarPath
rootBuilder.addDependency(dependency)
}

implicitDeps.keys.subtract(explicitDeps.keys).forEach {
val dependency = Deps.Dependency.newBuilder()
dependency.kind = Deps.Dependency.Kind.IMPLICIT
dependency.path = it
rootBuilder.addDependency(dependency)
}

BufferedOutputStream(File(jdepsOutput).outputStream()).use {
it.write(rootBuilder.buildSorted().toByteArray())
}
}

private fun doStrictDeps(
compilerConfiguration: CompilerConfiguration,
targetLabel: String,
directDeps: MutableList<String>,
explicitDeps: Map<String, List<String>>,
) {
when (compilerConfiguration.getNotNull(JdepsGenConfigurationKeys.STRICT_KOTLIN_DEPS)) {
"warn" -> checkStrictDeps(explicitDeps, directDeps, targetLabel)
"error" -> {
if (checkStrictDeps(explicitDeps, directDeps, targetLabel)) {
error(
"Strict Deps Violations - please fix",
)
}
}
}
}

/**
* Prints strict deps warnings and returns true if violations were found.
*/
private fun checkStrictDeps(
result: Map<String, List<String>>,
directDeps: List<String>,
targetLabel: String,
): Boolean {
val missingStrictDeps = result.keys
.filter { !directDeps.contains(it) }
.map { JarOwner.readJarOwnerFromManifest(Paths.get(it)) }

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

val open = "\u001b[35m\u001b[1m"
val close = "\u001b[0m"

var command =
"""
$open ** Please add the following dependencies:$close
${
missingStrictDeps.map { it.label ?: it.jar }.joinToString(" ")
} to $targetLabel
"""

if (missingStrictLabels.isNotEmpty()) {
command += """$open ** You can use the following buildozer command:$close
buildozer 'add deps ${
missingStrictLabels.joinToString(" ")
}' $targetLabel
"""
}

println(command.trimIndent())
return true
}
return false
}
}

private fun Deps.Dependencies.Builder.buildSorted(): Deps.Dependencies {
val sortedDeps = dependencyList.sortedBy { it.path }
sortedDeps.forEachIndexed { index, dep ->
setDependency(index, dep)
}
return build()
}
178 changes: 5 additions & 173 deletions src/main/kotlin/io/bazel/kotlin/plugin/jdeps/JdepsGenExtension.kt
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
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 com.intellij.openapi.project.Project
import com.intellij.psi.PsiElement
import io.bazel.kotlin.builder.utils.jars.JarOwner
import org.jetbrains.kotlin.analyzer.AnalysisResult
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.container.StorageComponentContainer
Expand Down Expand Up @@ -53,14 +49,6 @@ import org.jetbrains.kotlin.types.KotlinType
import org.jetbrains.kotlin.types.TypeConstructor
import org.jetbrains.kotlin.types.getAbbreviation
import org.jetbrains.kotlin.types.typeUtil.supertypes
import java.io.BufferedOutputStream
import java.io.File
import java.nio.file.Paths
import java.security.MessageDigest
import java.util.*
import java.util.jar.JarFile
import java.util.stream.Collectors


/**
* Kotlin compiler extension that tracks classes (and corresponding classpath jars) needed to
Expand All @@ -80,9 +68,10 @@ import java.util.stream.Collectors
* @param configuration the current compilation configuration
*/
class JdepsGenExtension(
val configuration: CompilerConfiguration,
) :
AnalysisHandlerExtension, StorageComponentContainerContributor {
configuration: CompilerConfiguration,
) : BaseJdepsGenExtension(configuration),
AnalysisHandlerExtension,
StorageComponentContainerContributor {

companion object {

Expand Down Expand Up @@ -375,165 +364,8 @@ class JdepsGenExtension(
bindingTrace: BindingTrace,
files: Collection<KtFile>,
): AnalysisResult? {
val directDeps = configuration.getList(JdepsGenConfigurationKeys.DIRECT_DEPENDENCIES)
val targetLabel = configuration.getNotNull(JdepsGenConfigurationKeys.TARGET_LABEL)
val explicitDeps = createDepsMap(explicitClassesCanonicalPaths)

doWriteJdeps(directDeps, targetLabel, explicitDeps)

doStrictDeps(configuration, targetLabel, directDeps, explicitDeps)
onAnalysisCompleted(explicitClassesCanonicalPaths, implicitClassesCanonicalPaths)

return super.analysisCompleted(project, module, bindingTrace, files)
}

/**
* Returns a map of jars to classes loaded from those jars.
*/
private fun createDepsMap(classes: Set<String>): Map<String, List<String>> {
val jarsToClasses = mutableMapOf<String, MutableList<String>>()
classes.forEach {
val parts = it.split("!/")
val jarPath = parts[0]
if (jarPath.endsWith(".jar")) {
jarsToClasses.computeIfAbsent(jarPath) { ArrayList() }.add(parts[1])
}
}
return jarsToClasses
}

private fun doWriteJdeps(
directDeps: MutableList<String>,
targetLabel: String,
explicitDeps: Map<String, List<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
val jdepsOutput = configuration.getNotNull(JdepsGenConfigurationKeys.OUTPUT_JDEPS)

val rootBuilder = Deps.Dependencies.newBuilder()
rootBuilder.success = true
rootBuilder.ruleLabel = targetLabel

val unusedDeps = directDeps.subtract(explicitDeps.keys)
unusedDeps.sorted().forEach { jarPath ->
val dependency = Deps.Dependency.newBuilder()
dependency.kind = Deps.Dependency.Kind.UNUSED
dependency.path = jarPath
rootBuilder.addDependency(dependency)
}

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)
}

implicitDeps.keys.subtract(explicitDeps.keys).sorted().forEach {
val dependency = Deps.Dependency.newBuilder()
dependency.kind = Deps.Dependency.Kind.IMPLICIT
dependency.path = it
rootBuilder.addDependency(dependency)
}

if (trackResourceUsage) {
usedResources.sorted().forEach { resource ->
rootBuilder.addUsedResources(resource)
}
}

BufferedOutputStream(File(jdepsOutput).outputStream()).use {
it.write(rootBuilder.build().toByteArray())
}
}

/**
* 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,
directDeps: MutableList<String>,
explicitDeps: Map<String, List<String>>,
) {
when (compilerConfiguration.getNotNull(JdepsGenConfigurationKeys.STRICT_KOTLIN_DEPS)) {
"warn" -> checkStrictDeps(explicitDeps, directDeps, targetLabel)
"error" -> {
if (checkStrictDeps(explicitDeps, directDeps, targetLabel)) {
error(
"Strict Deps Violations - please fix",
)
}
}
}
}

/**
* Prints strict deps warnings and returns true if violations were found.
*/
private fun checkStrictDeps(
result: Map<String, List<String>>,
directDeps: List<String>,
targetLabel: String,
): Boolean {
val missingStrictDeps = result.keys
.filter { !directDeps.contains(it) }
.map { JarOwner.readJarOwnerFromManifest(Paths.get(it)) }

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

val open = "\u001b[35m\u001b[1m"
val close = "\u001b[0m"

var command =
"""
$open ** Please add the following dependencies:$close
${
missingStrictDeps.map { it.label ?: it.jar }.joinToString(" ")
} to $targetLabel
"""

if (missingStrictLabels.isNotEmpty()) {
command += """$open ** You can use the following buildozer command:$close
buildozer 'add deps ${
missingStrictLabels.joinToString(" ")
}' $targetLabel
"""
}

println(command.trimIndent())
return true
}
return false
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import io.bazel.kotlin.builder.toolchain.CompilationTaskContext;
import io.bazel.kotlin.model.CompilationTaskInfo;
import io.bazel.kotlin.model.JvmCompilationTask;
import io.bazel.kotlin.model.KotlinToolchainInfo;

import java.util.EnumSet;
import java.util.HashSet;
Expand Down Expand Up @@ -155,7 +156,7 @@ public TaskBuilder compileJava() {
}

public TaskBuilder compileKotlin() {
taskBuilder.setInfo(CompilationTaskInfo.newBuilder().addDebug("trace").addDebug("timings"));
taskBuilder.getInfoBuilder().addDebug("trace").addDebug("timings");
taskBuilder.setCompileKotlin(true);
return this;
}
Expand Down Expand Up @@ -241,7 +242,10 @@ public TaskBuilder incrementalData() {
}

public TaskBuilder useK2() {
taskBuilder.getInfoBuilder().addPassthroughFlags("-Xuse-k2");
taskBuilder.getInfoBuilder()
.getToolchainInfoBuilder()
.getCommonBuilder()
.setLanguageVersion("2.0");
return this;
}
}
Expand Down
Loading

0 comments on commit 8da989d

Please sign in to comment.