[internal] kotlin: add kotlin parser for eventual dependency inference (

#15077)

Add a Kotlin parser using the IntelliJ PSI parser built into the Kotlin compiler.

src/python/pants/backend/experimental/kotlin/register.py

@@ -1,6 +1,7 @@
 # Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 from pants.backend.kotlin.compile import kotlinc
+from pants.backend.kotlin.dependency_inference.rules import rules as dep_inf_rules
 from pants.backend.kotlin.goals import check, tailor
 from pants.backend.kotlin.target_types import KotlinSourcesGeneratorTarget, KotlinSourceTarget
 from pants.backend.kotlin.target_types import rules as target_types_rules
@@ -32,6 +33,7 @@ def rules():
         *lockfile.rules(),
         *coursier_fetch.rules(),
         *coursier_setup.rules(),
+        *dep_inf_rules(),
         *jvm_util_rules.rules(),
         *jdk_rules.rules(),
         *target_types_rules(),

src/python/pants/backend/kotlin/dependency_inference/BUILD

@@ -0,0 +1,7 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+python_sources(dependencies=[":kotlin_resources"])
+resources(name="kotlin_resources", sources=["*.kt"])
+
+python_tests(name="tests", timeout=240)

src/python/pants/backend/kotlin/dependency_inference/KotlinParser.kt

@@ -0,0 +1,71 @@
+package org.pantsbuild.backend.kotlin.dependency_inference
+
+import java.nio.file.Files
+import java.nio.file.Paths
+import java.nio.charset.StandardCharsets
+
+import com.google.gson.Gson
+import com.intellij.openapi.util.Disposer
+import com.intellij.psi.PsiManager
+import com.intellij.testFramework.LightVirtualFile
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.psi.KtFile
+
+
+// KtFile: https://github.com/JetBrains/kotlin/blob/8bc29a30111081ee0b0dbe06d1f648a789909a27/compiler/psi/src/org/jetbrains/kotlin/psi/KtFile.kt
+
+data class KotlinImport(
+    val name: String,
+    val alias: String?,
+    val isWildcard: Boolean,
+)
+
+data class KotlinAnalysis(
+    val imports: List<KotlinImport>,
+)
+
+fun parse(code: String): KtFile {
+    val disposable = Disposer.newDisposable()
+    try {
+        val env = KotlinCoreEnvironment.createForProduction(
+            disposable, CompilerConfiguration(), EnvironmentConfigFiles.JVM_CONFIG_FILES)
+        val file = LightVirtualFile("temp.kt", KotlinFileType.INSTANCE, code)
+        return PsiManager.getInstance(env.project).findFile(file) as KtFile
+    } finally {
+        disposable.dispose()
+    }
+}
+
+fun analyze(file: KtFile): KotlinAnalysis {
+    val imports = file.importDirectives.map { importDirective ->
+        val name = importDirective.importedFqName
+        if (name != null) {
+            KotlinImport(
+                name=name.asString(),
+                alias=importDirective.aliasName,
+                isWildcard=importDirective.isAllUnder,
+            )
+        } else {
+            null
+        }
+    }
+
+    return KotlinAnalysis(imports.filterNotNull())
+}
+
+fun main(args: Array<String>) {
+    val analysisOutputPath = args[0]
+    val sourcePath = args[1]
+
+    val sourceContentBytes = Files.readAllBytes(Paths.get(sourcePath))
+    val sourceContent = String(sourceContentBytes, StandardCharsets.UTF_8)
+    val parsed = parse(sourceContent)
+    val analysis = analyze(parsed)
+
+    val gson = Gson()
+    val analysisOutput = gson.toJson(analysis)
+    Files.write(Paths.get(analysisOutputPath), analysisOutput.toByteArray(StandardCharsets.UTF_8))
+}

src/python/pants/backend/kotlin/dependency_inference/kotlin_parser.py

@@ -0,0 +1,241 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
+import json
+import os
+import pkgutil
+from dataclasses import dataclass
+
+from pants.backend.kotlin.subsystems.kotlin import DEFAULT_KOTLIN_VERSION
+from pants.core.util_rules.source_files import SourceFiles
+from pants.engine.fs import CreateDigest, DigestContents, Directory, FileContent
+from pants.engine.internals.native_engine import AddPrefix, Digest, MergeDigests, RemovePrefix
+from pants.engine.internals.selectors import Get, MultiGet
+from pants.engine.process import FallibleProcessResult, ProcessExecutionFailure, ProcessResult
+from pants.engine.rules import collect_rules, rule
+from pants.jvm.compile import ClasspathEntry
+from pants.jvm.jdk_rules import InternalJdk, JdkEnvironment, JdkRequest, JvmProcess
+from pants.jvm.resolve.common import ArtifactRequirements, Coordinate
+from pants.jvm.resolve.coursier_fetch import ToolClasspath, ToolClasspathRequest
+from pants.option.global_options import ProcessCleanupOption
+from pants.util.logging import LogLevel
+
+_KOTLIN_PARSER_ARTIFACT_REQUIREMENTS = ArtifactRequirements.from_coordinates(
+    [
+        Coordinate(
+            group="org.jetbrains.kotlin",
+            artifact="kotlin-compiler",
+            version=DEFAULT_KOTLIN_VERSION,  # TODO: Make this follow the resolve?
+        ),
+        Coordinate(
+            group="org.jetbrains.kotlin",
+            artifact="kotlin-stdlib",
+            version=DEFAULT_KOTLIN_VERSION,  # TODO: Make this follow the resolve?
+        ),
+        Coordinate(
+            group="com.google.code.gson",
+            artifact="gson",
+            version="2.9.0",
+        ),
+    ]
+)
+
+
+@dataclass(frozen=True)
+class KotlinImport:
+    name: str
+    alias: str | None
+    is_wildcard: bool
+
+    @classmethod
+    def from_json_dict(cls, d: dict) -> KotlinImport:
+        return cls(
+            name=d["name"],
+            alias=d.get("alias"),
+            is_wildcard=d["isWildcard"],
+        )
+
+
+@dataclass(frozen=True)
+class KotlinSourceDependencyAnalysis:
+    imports: frozenset[KotlinImport]
+
+    @classmethod
+    def from_json_dict(cls, d: dict) -> KotlinSourceDependencyAnalysis:
+        return cls(
+            imports=frozenset(KotlinImport.from_json_dict(i) for i in d["imports"]),
+        )
+
+
+@dataclass(frozen=True)
+class FallibleKotlinSourceDependencyAnalysisResult:
+    process_result: FallibleProcessResult
+
+
+class KotlinParserCompiledClassfiles(ClasspathEntry):
+    pass
+
+
+@rule(level=LogLevel.DEBUG)
+async def analyze_kotlin_source_dependencies(
+    processor_classfiles: KotlinParserCompiledClassfiles,
+    source_files: SourceFiles,
+) -> FallibleKotlinSourceDependencyAnalysisResult:
+    # Use JDK 8 due to https://youtrack.jetbrains.com/issue/KTIJ-17192 and https://youtrack.jetbrains.com/issue/KT-37446.
+    request = JdkRequest("adopt:8")
+    env = await Get(JdkEnvironment, JdkRequest, request)
+    jdk = InternalJdk(env._digest, env.nailgun_jar, env.coursier, env.jre_major_version)
+
+    if len(source_files.files) > 1:
+        raise ValueError(
+            f"analyze_kotlin_source_dependencies expects sources with exactly 1 source file, but found {len(source_files.snapshot.files)}."
+        )
+    elif len(source_files.files) == 0:
+        raise ValueError(
+            "analyze_kotlin_source_dependencies expects sources with exactly 1 source file, but found none."
+        )
+    source_prefix = "__source_to_analyze"
+    source_path = os.path.join(source_prefix, source_files.files[0])
+    processorcp_relpath = "__processorcp"
+    toolcp_relpath = "__toolcp"
+
+    (tool_classpath, prefixed_source_files_digest,) = await MultiGet(
+        Get(
+            ToolClasspath,
+            ToolClasspathRequest(artifact_requirements=_KOTLIN_PARSER_ARTIFACT_REQUIREMENTS),
+        ),
+        Get(Digest, AddPrefix(source_files.snapshot.digest, source_prefix)),
+    )
+
+    extra_immutable_input_digests = {
+        toolcp_relpath: tool_classpath.digest,
+        processorcp_relpath: processor_classfiles.digest,
+    }
+
+    analysis_output_path = "__source_analysis.json"
+
+    process_result = await Get(
+        FallibleProcessResult,
+        JvmProcess(
+            jdk=jdk,
+            classpath_entries=[
+                *tool_classpath.classpath_entries(toolcp_relpath),
+                processorcp_relpath,
+            ],
+            argv=[
+                "org.pantsbuild.backend.kotlin.dependency_inference.KotlinParserKt",
+                analysis_output_path,
+                source_path,
+            ],
+            input_digest=prefixed_source_files_digest,
+            extra_immutable_input_digests=extra_immutable_input_digests,
+            output_files=(analysis_output_path,),
+            extra_nailgun_keys=extra_immutable_input_digests,
+            description=f"Analyzing {source_files.files[0]}",
+            level=LogLevel.DEBUG,
+        ),
+    )
+
+    return FallibleKotlinSourceDependencyAnalysisResult(process_result=process_result)
+
+
+@rule(level=LogLevel.DEBUG)
+async def resolve_fallible_result_to_analysis(
+    fallible_result: FallibleKotlinSourceDependencyAnalysisResult,
+    process_cleanup: ProcessCleanupOption,
+) -> KotlinSourceDependencyAnalysis:
+    # TODO(#12725): Just convert directly to a ProcessResult like this:
+    # result = await Get(ProcessResult, FallibleProcessResult, fallible_result.process_result)
+    if fallible_result.process_result.exit_code == 0:
+        analysis_contents = await Get(
+            DigestContents, Digest, fallible_result.process_result.output_digest
+        )
+        analysis = json.loads(analysis_contents[0].content)
+        return KotlinSourceDependencyAnalysis.from_json_dict(analysis)
+    raise ProcessExecutionFailure(
+        fallible_result.process_result.exit_code,
+        fallible_result.process_result.stdout,
+        fallible_result.process_result.stderr,
+        "Kotlin source dependency analysis failed.",
+        process_cleanup=process_cleanup.val,
+    )
+
+
+@rule
+async def setup_kotlin_parser_classfiles(jdk: InternalJdk) -> KotlinParserCompiledClassfiles:
+    dest_dir = "classfiles"
+
+    parser_source_content = pkgutil.get_data(
+        "pants.backend.kotlin.dependency_inference", "KotlinParser.kt"
+    )
+    if not parser_source_content:
+        raise AssertionError("Unable to find KotlinParser.kt resource.")
+
+    parser_source = FileContent("KotlinParser.kt", parser_source_content)
+
+    tool_classpath, parser_classpath, source_digest = await MultiGet(
+        Get(
+            ToolClasspath,
+            ToolClasspathRequest(
+                prefix="__toolcp",
+                artifact_requirements=ArtifactRequirements.from_coordinates(
+                    [
+                        Coordinate(
+                            group="org.jetbrains.kotlin",
+                            artifact="kotlin-compiler",
+                            version=DEFAULT_KOTLIN_VERSION,  # TODO: Pull from resolve or hard-code Kotlin version?
+                        ),
+                    ]
+                ),
+            ),
+        ),
+        Get(
+            ToolClasspath,
+            ToolClasspathRequest(
+                prefix="__parsercp", artifact_requirements=_KOTLIN_PARSER_ARTIFACT_REQUIREMENTS
+            ),
+        ),
+        Get(Digest, CreateDigest([parser_source, Directory(dest_dir)])),
+    )
+
+    merged_digest = await Get(
+        Digest,
+        MergeDigests(
+            (
+                tool_classpath.digest,
+                parser_classpath.digest,
+                source_digest,
+            )
+        ),
+    )
+
+    process_result = await Get(
+        ProcessResult,
+        JvmProcess(
+            jdk=jdk,
+            classpath_entries=tool_classpath.classpath_entries(),
+            argv=[
+                "org.jetbrains.kotlin.cli.jvm.K2JVMCompiler",
+                "-classpath",
+                ":".join(parser_classpath.classpath_entries()),
+                "-d",
+                dest_dir,
+                parser_source.path,
+            ],
+            input_digest=merged_digest,
+            output_directories=(dest_dir,),
+            description="Compile Kotlin parser for dependency inference with kotlinc",
+            level=LogLevel.DEBUG,
+            # NB: We do not use nailgun for this process, since it is launched exactly once.
+            use_nailgun=False,
+        ),
+    )
+    stripped_classfiles_digest = await Get(
+        Digest, RemovePrefix(process_result.output_digest, dest_dir)
+    )
+    return KotlinParserCompiledClassfiles(digest=stripped_classfiles_digest)
+
+
+def rules():
+    return collect_rules()