Skip to content

Commit

Permalink
Improve Kotlin jdeps logic to find more explicit and implicit depende…
Browse files Browse the repository at this point in the history
…ncies (bazelbuild#17)

# Problem
Kotlin jdeps do not collect many direct and indirect dependencies that are required for compilation. These missing jdep found dependencies were discovered as Airbnb tests a "jdeps compile-time pruning" optimization which provides significant IntelliJ indexing, UI, and compilation performance benefits in our fork of the [Bazel IntelliJ Plugin](https://github.com/bazelbuild/intellij).

# Solution
As a followup to bazelbuild#1045, write test cases to repro missing dependencies found in Airbnb's monorepo, and then fix the underlying issues. In addition to these fixes helping our "jdeps compile-time pruning" optimization, these fixes should also help when the normal Bazel IntelliJ Plugin is used, as it also consults jdeps when deciding what jars to import into IntelliJ.

# Test Plan:
* Wrote Kotlin and Java test cases to repro missing dependencies found in the Airbnb monorepo
* Updated Airbnb's minimal rules_kotlin fork to use the changes in this PR, which has been successfully running CI for over a week.
  • Loading branch information
Steve Cosenza committed Feb 26, 2024
1 parent 06a36ed commit 72b1641
Show file tree
Hide file tree
Showing 25 changed files with 520 additions and 57 deletions.
2 changes: 1 addition & 1 deletion .bazelrc
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@ common --enable_bzlmod=true
build --strategy=KotlinCompile=worker
build --test_output=all
build --verbose_failures

common --jvmopt=-Djava.security.manager=allow
174 changes: 122 additions & 52 deletions src/main/kotlin/io/bazel/kotlin/plugin/jdeps/JdepsGenExtension.kt
Original file line number Diff line number Diff line change
Expand Up @@ -11,36 +11,43 @@ import org.jetbrains.kotlin.container.useInstance
import org.jetbrains.kotlin.descriptors.ClassDescriptor
import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithSource
import org.jetbrains.kotlin.descriptors.ValueDescriptor
import org.jetbrains.kotlin.descriptors.CallableDescriptor
import org.jetbrains.kotlin.descriptors.FunctionDescriptor
import org.jetbrains.kotlin.descriptors.ModuleDescriptor
import org.jetbrains.kotlin.descriptors.ParameterDescriptor
import org.jetbrains.kotlin.descriptors.PropertyDescriptor
import org.jetbrains.kotlin.descriptors.SourceElement
import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
import org.jetbrains.kotlin.load.java.descriptors.JavaMethodDescriptor
import org.jetbrains.kotlin.load.java.descriptors.JavaPropertyDescriptor
import org.jetbrains.kotlin.js.translate.callTranslator.getReturnType
import org.jetbrains.kotlin.load.java.descriptors.JavaClassConstructorDescriptor
import org.jetbrains.kotlin.load.java.sources.JavaSourceElement
import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryJavaClass
import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryJavaField
import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
import org.jetbrains.kotlin.load.kotlin.KotlinJvmBinarySourceElement
import org.jetbrains.kotlin.load.kotlin.VirtualFileKotlinClass
import org.jetbrains.kotlin.load.kotlin.getContainingKotlinJvmBinaryClass
import org.jetbrains.kotlin.platform.TargetPlatform
import org.jetbrains.kotlin.psi.KtDeclaration
import org.jetbrains.kotlin.psi.KtFile
import org.jetbrains.kotlin.resolve.BindingTrace
import org.jetbrains.kotlin.resolve.FunctionImportedFromObject
import org.jetbrains.kotlin.resolve.PropertyImportedFromObject
import org.jetbrains.kotlin.resolve.ImportedFromObjectCallableDescriptor
import org.jetbrains.kotlin.resolve.calls.checkers.CallChecker
import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
import org.jetbrains.kotlin.resolve.calls.model.ExpressionKotlinCallArgument
import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
import org.jetbrains.kotlin.resolve.calls.util.FakeCallableDescriptorForObject
import org.jetbrains.kotlin.resolve.calls.tower.NewAbstractResolvedCall
import org.jetbrains.kotlin.resolve.calls.tower.PSIKotlinCallArgument
import org.jetbrains.kotlin.resolve.checkers.DeclarationChecker
import org.jetbrains.kotlin.resolve.checkers.DeclarationCheckerContext
import org.jetbrains.kotlin.resolve.descriptorUtil.getImportableDescriptor
import org.jetbrains.kotlin.resolve.jvm.extensions.AnalysisHandlerExtension
import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedTypeAliasDescriptor
import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
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
Expand Down Expand Up @@ -77,7 +84,21 @@ class JdepsGenExtension(
* @return the path corresponding to the JAR where this class was loaded from, or null.
*/
fun getClassCanonicalPath(descriptor: DeclarationDescriptorWithSource): String? {
return when (val sourceElement: SourceElement = descriptor.source) {
// Get the descriptor's source element which may be a type alias
val sourceElement = if (descriptor.source != SourceElement.NO_SOURCE) {
descriptor.source
} else {
when(val declarationDescriptor = descriptor.getImportableDescriptor()) {
is DeserializedTypeAliasDescriptor -> {
declarationDescriptor.containerSource
}
else -> {
null
}
}
}

return when (sourceElement) {
is JavaSourceElement ->
if (sourceElement.javaElement is BinaryJavaClass) {
(sourceElement.javaElement as BinaryJavaClass).virtualFile.canonicalPath
Expand All @@ -94,11 +115,16 @@ class JdepsGenExtension(
}
is KotlinJvmBinarySourceElement ->
(sourceElement.binaryClass as VirtualFileKotlinClass).file.canonicalPath
else -> null
is JvmPackagePartSource -> { // This case is needed to collect type aliases
((sourceElement.knownJvmBinaryClass) as VirtualFileKotlinClass).file.canonicalPath
}
else -> {
null
}
}
}

fun getClassCanonicalPath(typeConstructor: TypeConstructor): String? {
private fun getClassCanonicalPath(typeConstructor: TypeConstructor): String? {
return (typeConstructor.declarationDescriptor as? DeclarationDescriptorWithSource)?.let {
getClassCanonicalPath(
it,
Expand Down Expand Up @@ -130,54 +156,87 @@ class JdepsGenExtension(
reportOn: PsiElement,
context: CallCheckerContext,
) {
when (val resultingDescriptor = resolvedCall.resultingDescriptor) {
is FunctionImportedFromObject -> {
collectTypeReferences(resultingDescriptor.containingObject.defaultType)
}
is PropertyImportedFromObject -> {
collectTypeReferences(resultingDescriptor.containingObject.defaultType)
// First collect type from the Resolved Call
collectExplicitTypes(resolvedCall)

resolvedCall.valueArguments.keys.forEach { valueArgument ->
collectTypeReferences(valueArgument.type, isExplicit = false)
}

// And then collect types from any ResultingDescriptor
val resultingDescriptor = resolvedCall.resultingDescriptor
collectExplicitTypes(resultingDescriptor)

val isExplicitReturnType = resultingDescriptor is JavaClassConstructorDescriptor
collectTypeReferences(resolvedCall.getReturnType(), isExplicit = isExplicitReturnType, collectTypeArguments = false)
resultingDescriptor.returnType?.let {
collectTypeReferences(it, isExplicit = isExplicitReturnType, collectTypeArguments = false)
}

resultingDescriptor.valueParameters.forEach { valueParameter ->
collectTypeReferences(valueParameter.type, isExplicit = false)
}

// Finally, collect types that depend on the type of the ResultingDescriptor and note that
// these descriptors may be composed of multiple classes that we need to extract types from
if (resultingDescriptor is DeclarationDescriptor) {
val containingDeclaration = resultingDescriptor.containingDeclaration
if (containingDeclaration is ClassDescriptor) {
collectTypeReferences(containingDeclaration.defaultType)
}
is JavaMethodDescriptor -> {
getClassCanonicalPath(
(resultingDescriptor.containingDeclaration as ClassDescriptor).typeConstructor,
)?.let { explicitClassesCanonicalPaths.add(it) }

if (resultingDescriptor is PropertyDescriptor) {
(resultingDescriptor.getter
?.correspondingProperty as? SyntheticJavaPropertyDescriptor)
?.let { syntheticJavaPropertyDescriptor ->
collectTypeReferences(syntheticJavaPropertyDescriptor.type, isExplicit = false)

val functionDescriptor = syntheticJavaPropertyDescriptor.getMethod
functionDescriptor.dispatchReceiverParameter?.type?.let { dispatchReceiverType ->
collectTypeReferences(dispatchReceiverType, isExplicit = false)
}
}
}
is FunctionDescriptor -> {
resultingDescriptor.returnType?.let {
collectTypeReferences(it, isExplicit = false, collectTypeArguments = false)
}
resultingDescriptor.valueParameters.forEach { valueParameter ->
collectTypeReferences(valueParameter.type, isExplicit = false)
}

if (resultingDescriptor is ImportedFromObjectCallableDescriptor<*>) {
collectTypeReferences(resultingDescriptor.containingObject.defaultType, isExplicit = false)
}

if (resultingDescriptor is ValueDescriptor) {
collectTypeReferences(resultingDescriptor.type, isExplicit = false)
}
}

private fun collectExplicitTypes(resultingDescriptor: CallableDescriptor) {
val kotlinJvmBinaryClass = resultingDescriptor.getContainingKotlinJvmBinaryClass()
if (kotlinJvmBinaryClass is VirtualFileKotlinClass) {
explicitClassesCanonicalPaths.add(kotlinJvmBinaryClass.file.path)
}
}

private fun collectExplicitTypes(resolvedCall: ResolvedCall<*>) {
val kotlinCallArguments = (resolvedCall as? NewAbstractResolvedCall)
?.resolvedCallAtom?.atom?.argumentsInParenthesis

// note that callArgument can be both a PSIKotlinCallArgument and an ExpressionKotlinCallArgument
kotlinCallArguments?.forEach { callArgument ->
if (callArgument is PSIKotlinCallArgument) {
val dataFlowInfos = listOf(callArgument.dataFlowInfoBeforeThisArgument) +
callArgument.dataFlowInfoAfterThisArgument

dataFlowInfos.forEach { dataFlowInfo ->
dataFlowInfo.completeTypeInfo.values().flatten().forEach { kotlinType ->
collectTypeReferences(kotlinType, isExplicit = true)
}
}
val virtualFileClass =
resultingDescriptor.getContainingKotlinJvmBinaryClass() as? VirtualFileKotlinClass
?: return
explicitClassesCanonicalPaths.add(virtualFileClass.file.path)
}
is ParameterDescriptor -> {
getClassCanonicalPath(resultingDescriptor)?.let { explicitClassesCanonicalPaths.add(it) }
}
is FakeCallableDescriptorForObject -> {
collectTypeReferences(resultingDescriptor.type)
}
is JavaPropertyDescriptor -> {
getClassCanonicalPath(resultingDescriptor)?.let { explicitClassesCanonicalPaths.add(it) }
}
is PropertyDescriptor -> {
when (resultingDescriptor.containingDeclaration) {
is ClassDescriptor -> collectTypeReferences(
(resultingDescriptor.containingDeclaration as ClassDescriptor).defaultType,
)
else -> {
val virtualFileClass =
(resultingDescriptor).getContainingKotlinJvmBinaryClass() as? VirtualFileKotlinClass
?: return
explicitClassesCanonicalPaths.add(virtualFileClass.file.path)
}

if (callArgument is ExpressionKotlinCallArgument) {
callArgument.receiver.allOriginalTypes.forEach { kotlinType ->
collectTypeReferences(kotlinType, isExplicit = true)
}
collectTypeReferences(resultingDescriptor.type, isExplicit = false)
}
else -> return
}
}

Expand All @@ -193,7 +252,7 @@ class JdepsGenExtension(
}
}
is FunctionDescriptor -> {
descriptor.returnType?.let { collectTypeReferences(it) }
descriptor.returnType?.let { collectTypeReferences(it, collectTypeArguments = false ) }
descriptor.valueParameters.forEach { valueParameter ->
collectTypeReferences(valueParameter.type)
}
Expand Down Expand Up @@ -245,6 +304,17 @@ class JdepsGenExtension(
}
}

// Collect type aliases aka abbreviations
// See: https://github.com/Kotlin/KEEP/blob/master/proposals/type-aliases.md#type-alias-expansion
kotlinType.getAbbreviation()?.let { abbreviationType ->
collectTypeReferences(
abbreviationType,
isExplicit = isExplicit,
collectTypeArguments = collectTypeArguments,
visitedKotlinTypes = visitedKotlinTypes,
)
}

kotlinType.supertypes().forEach { supertype ->
collectTypeReferences(
supertype,
Expand Down
14 changes: 13 additions & 1 deletion src/test/kotlin/io/bazel/kotlin/builder/tasks/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,15 @@ java_library(
name = "JdepsParserTestFixtures",
testonly = True,
srcs = glob(["testFixtures/*.java"]),
deps = [
":JdepsParserTestFixtures2",
],
)

java_library(
name = "JdepsParserTestFixtures2",
testonly = True,
srcs = glob(["testFixtures2/*.java"]),
)

kt_rules_test(
Expand All @@ -44,7 +53,10 @@ kt_rules_test(
name = "KotlinBuilderJvmJdepsTest",
size = "large",
srcs = ["jvm/KotlinBuilderJvmJdepsTest.kt"],
data = [":JdepsParserTestFixtures"],
data = [
":JdepsParserTestFixtures",
":JdepsParserTestFixtures2",
],
)

kt_rules_test(
Expand Down
Loading

0 comments on commit 72b1641

Please sign in to comment.