Add experimental Kotlin incremental compilation

MODULE.bazel

@@ -29,7 +29,9 @@ use_repo(
     "com_github_jetbrains_kotlin",
     "com_github_jetbrains_kotlin_git",
     "com_github_pinterest_ktlint",
+    "kotlin_build_tools_api",
     "kotlin_build_tools_impl",
+    "kotlin_compiler_embeddable",
     "kotlinx_serialization_core_jvm",
     "kotlinx_serialization_json",
     "kotlinx_serialization_json_jvm",

kotlin/internal/jvm/compile.bzl

@@ -463,6 +463,7 @@ def _run_kt_builder_action(
     args.add_all("--classpath", compile_deps.compile_jars)
     args.add("--reduced_classpath_mode", toolchains.kt.experimental_reduce_classpath_mode)
     args.add("--build_tools_api", toolchains.kt.experimental_build_tools_api)
+    args.add("--incremental_compilation", toolchains.kt.experimental_incremental_compilation)
     args.add_all("--sources", srcs.all_srcs, omit_if_empty = True)
     args.add_all("--source_jars", srcs.src_jars + generated_src_jars, omit_if_empty = True)
     args.add_all("--deps_artifacts", deps_artifacts, omit_if_empty = True)

kotlin/internal/toolchains.bzl

@@ -92,6 +92,7 @@ def _kotlin_toolchain_impl(ctx):
         experimental_report_unused_deps = ctx.attr.experimental_report_unused_deps,
         experimental_reduce_classpath_mode = ctx.attr.experimental_reduce_classpath_mode,
         experimental_build_tools_api = ctx.attr.experimental_build_tools_api,
+        experimental_incremental_compilation = ctx.attr.experimental_incremental_compilation,
         javac_options = ctx.attr.javac_options[JavacOptions] if ctx.attr.javac_options else None,
         kotlinc_options = ctx.attr.kotlinc_options[KotlincOptions] if ctx.attr.kotlinc_options else None,
         empty_jar = ctx.file._empty_jar,
@@ -252,6 +253,10 @@ _kt_toolchain = rule(
             doc = "Enables experimental support for Build Tools API integration",
             default = False,
         ),
+        "experimental_incremental_compilation": attr.bool(
+            doc = "TODO",
+            default = False,
+        ),
         "javac_options": attr.label(
             doc = "Compiler options for javac",
             providers = [JavacOptions],
@@ -327,6 +332,7 @@ def define_kt_toolchain(
         experimental_reduce_classpath_mode = None,
         experimental_multiplex_workers = None,
         experimental_build_tools_api = None,
+        experimental_incremental_compilation = None,
         javac_options = Label("//kotlin/internal:default_javac_options"),
         kotlinc_options = Label("//kotlin/internal:default_kotlinc_options"),
         jvm_stdlibs = None,
@@ -356,6 +362,7 @@ def define_kt_toolchain(
         experimental_report_unused_deps = experimental_report_unused_deps,
         experimental_reduce_classpath_mode = experimental_reduce_classpath_mode,
         experimental_build_tools_api = experimental_build_tools_api,
+        experimental_incremental_compilation = experimental_incremental_compilation,
         javac_options = javac_options,
         kotlinc_options = kotlinc_options,
         visibility = ["//visibility:public"],

kotlin/settings/BUILD.bazel

@@ -48,3 +48,9 @@ bool_flag(
     build_setting_default = False,
     visibility = ["//visibility:public"],
 )
+
+bool_flag(
+    name = "experimental_incremental_compilation",
+    build_setting_default = False,
+    visibility = ["//visibility:public"],
+)

src/main/kotlin/io/bazel/kotlin/builder/cmd/BUILD.bazel

@@ -27,6 +27,8 @@ kt_bootstrap_binary(
         "//src/main/kotlin/io/bazel/kotlin/compiler:compiler.jar",
         "@com_github_jetbrains_kotlin//:home",
         "@kotlin_build_tools_impl//jar",
+        "@kotlin_compiler_embeddable//jar",
+        "@kotlin_rules_maven//:org_jetbrains_kotlin_kotlin_daemon_client",
         "@kotlinx_serialization_core_jvm//jar",
         "@kotlinx_serialization_json//jar",
         "@kotlinx_serialization_json_jvm//jar",
@@ -36,10 +38,12 @@ kt_bootstrap_binary(
         "-D@com_github_jetbrains_kotlinx...serialization-json=$(rlocationpath @kotlinx_serialization_json//jar)",
         "-D@com_github_jetbrains_kotlinx...serialization-json-jvm=$(rlocationpath @kotlinx_serialization_json_jvm//jar)",
         "-D@com_github_jetbrains_kotlin...build-tools-impl=$(rlocationpath @kotlin_build_tools_impl//jar)",
+        "-D@kotlin_compiler_embeddable...kotlin-embedded=$(rlocationpath @kotlin_compiler_embeddable//jar)",
         "-D@com_github_jetbrains_kotlin...jvm-abi-gen=$(rlocationpath //kotlin/compiler:jvm-abi-gen)",
         "-D@com_github_jetbrains_kotlin...kotlin-compiler=$(rlocationpath //kotlin/compiler:kotlin-compiler)",
         "-D@com_github_jetbrains_kotlin...kapt=$(rlocationpath //kotlin/compiler:kotlin-annotation-processing)",
         "-D@rules_kotlin...jdeps-gen=$(rlocationpath //src/main/kotlin:jdeps-gen)",
+        "-D@com_github_jetbrains_kotlin...daemon-client=$(rlocationpath @kotlin_rules_maven//:org_jetbrains_kotlin_kotlin_daemon_client)",
         "-D@rules_kotlin...skip-code-gen=$(rlocationpath //src/main/kotlin:skip-code-gen)",
         "-D@rules_kotlin...compiler=$(rlocationpath //src/main/kotlin/io/bazel/kotlin/compiler:compiler.jar)",
         "-D@com_github_google_ksp...symbol-processing-api=$(rlocationpath //kotlin/compiler:symbol-processing-api)",

src/main/kotlin/io/bazel/kotlin/builder/tasks/BUILD.bazel

@@ -24,6 +24,7 @@ kt_bootstrap_library(
     ],
     visibility = ["//src:__subpackages__"],
     deps = [
+        "//kotlin/compiler:kotlin-compiler",
         "//kotlin/compiler:kotlin-preloader",
         "//src/main/kotlin/io/bazel/kotlin/builder/toolchain",
         "//src/main/kotlin/io/bazel/kotlin/builder/utils",
@@ -32,6 +33,9 @@ kt_bootstrap_library(
         "//src/main/protobuf:deps_java_proto",
         "//src/main/protobuf:kotlin_model_java_proto",
         "//src/main/protobuf:worker_protocol_java_proto",
+        "@kotlin_build_tools_api//jar",
+        "@kotlin_build_tools_impl//jar",
+        "@kotlin_compiler_embeddable//jar",
         "@kotlin_rules_maven//:com_google_protobuf_protobuf_java",
         "@kotlin_rules_maven//:com_google_protobuf_protobuf_java_util",
         "@kotlin_rules_maven//:javax_inject_javax_inject",

src/main/kotlin/io/bazel/kotlin/builder/tasks/KotlinBuilder.kt

@@ -88,6 +88,7 @@ class KotlinBuilder
         INSTRUMENT_COVERAGE("--instrument_coverage"),
         KSP_GENERATED_JAVA_SRCJAR("--ksp_generated_java_srcjar"),
         BUILD_TOOLS_API("--build_tools_api"),
+        INCREMENTAL_COMPILATION("--incremental_compilation"),
       }
     }
 
@@ -173,6 +174,9 @@ class KotlinBuilder
         argMap.optionalSingle(KotlinBuilderFlags.BUILD_TOOLS_API)?.let {
           buildToolsApi = it == "true"
         }
+        argMap.optionalSingle(KotlinBuilderFlags.INCREMENTAL_COMPILATION)?.let {
+          incrementalCompilation = it == "true"
+        }
         this
       }
 
@@ -288,8 +292,17 @@ class KotlinBuilder
             argMap.optional(KotlinBuilderFlags.COMPILER_PLUGIN_CLASS_PATH) ?: emptyList(),
           )
 
+          // Kotlin compiler always requires absolute path for source input in incremental mode
+          val useAbsolutePath = argMap.optionalSingle(KotlinBuilderFlags.INCREMENTAL_COMPILATION) == "true"
           argMap
             .optional(KotlinBuilderFlags.SOURCES)
+            ?.map {
+              if (useAbsolutePath) {
+                FileSystems.getDefault().getPath(it).toAbsolutePath().toString()
+              } else {
+                it
+              }
+            }
             ?.iterator()
             ?.partitionJvmSources(
               { addKotlinSources(it) },

src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/ClasspathSnapshotGenerator.kt

@@ -0,0 +1,65 @@
+package io.bazel.kotlin.builder.tasks.jvm
+
+import java.io.FileInputStream
+import java.nio.file.Files
+import java.nio.file.Path
+import java.security.MessageDigest
+import kotlin.system.measureTimeMillis
+import org.jetbrains.kotlin.buildtools.api.CompilationService
+import org.jetbrains.kotlin.buildtools.api.ExperimentalBuildToolsApi
+import kotlin.io.path.exists
+
+@OptIn(ExperimentalBuildToolsApi::class)
+class ClasspathSnapshotGenerator(
+    private val inputJar: Path,
+    private val outputSnapshot: Path,
+    private val granularity: SnapshotGranularity
+) {
+
+    private val hashPath: Path by lazy {
+      outputSnapshot.resolveSibling(outputSnapshot.fileName.toString() + ".hash")
+    }
+
+    fun run() {
+        if (!isSnapshotOutdated()) {
+          return
+        }
+
+        val timeSpent = measureTimeMillis {
+          val compilationService =
+            CompilationService.loadImplementation(this.javaClass.classLoader!!)
+          val snapshot =
+            compilationService.calculateClasspathSnapshot(
+              inputJar.toFile(), granularity.toClassSnapshotGranularity
+            )
+          // TODO : make things atomic / avoid race conditions
+          val hash = hashFile(inputJar)
+          snapshot.saveSnapshot(outputSnapshot.toFile())
+          hashPath.toFile().writeText(hash)
+        }
+
+        // TODO: Log impl
+        // LOG.info("$timeSpent ms for input jar: $inputJar")
+    }
+
+    private fun isSnapshotOutdated(): Boolean {
+      if (!outputSnapshot.exists() || !hashPath.exists()) {
+        return true
+      }
+      val storedHash = Files.readAllLines(hashPath).firstOrNull()?.trim()
+      val currentHash = hashFile(inputJar)
+      return storedHash == null || storedHash != currentHash
+    }
+
+    private fun hashFile(path: Path): String {
+      val digest = MessageDigest.getInstance("SHA-256")
+      FileInputStream(path.toFile()).use { fis ->
+        val buffer = ByteArray(8192)
+        var bytesRead: Int
+        while (fis.read(buffer).also { bytesRead = it } != -1) {
+          digest.update(buffer, 0, bytesRead)
+        }
+      }
+      return digest.digest().joinToString("") { "%02x".format(it) }
+    }
+}

src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/CompilationTask.kt

@@ -35,6 +35,7 @@ import java.io.BufferedInputStream
 import java.io.ByteArrayOutputStream
 import java.io.File
 import java.io.ObjectOutputStream
+import java.nio.file.FileSystems
 import java.nio.file.Files
 import java.nio.file.Files.isDirectory
 import java.nio.file.Files.walk
@@ -91,6 +92,11 @@ fun JvmCompilationTask.baseArgs(overrides: Map<String, String> = emptyMap()): Co
       overrides[LANGUAGE_VERSION_ARG] ?: info.toolchainInfo.common.languageVersion,
     ).flag("-jvm-target", info.toolchainInfo.jvm.jvmTarget)
     .flag("-module-name", info.moduleName)
+    .apply {
+      if (info.buildToolsApi) {
+        flag("-label", info.label)
+      }
+    }
 }
 
 internal fun JvmCompilationTask.plugins(
@@ -308,8 +314,13 @@ private fun JvmCompilationTask.runKspPlugin(
     baseArgs(overrides)
       .plus(kspArgs(plugins))
       .flag("-d", directories.generatedClasses)
-      .values(inputs.kotlinSourcesList)
       .values(inputs.javaSourcesList)
+      .values(inputs.kotlinSourcesList)
+      .apply {
+        if (info.incrementalCompilation) {
+          flag("-incremental_id", "ksp")
+        }
+      }
       .list()
       .let { args ->
         context.executeCompilerTask(
@@ -422,6 +433,28 @@ internal fun JvmCompilationTask.createGeneratedKspKotlinSrcJar() {
   }
 }
 
+internal fun JvmCompilationTask.createClasspathSnapshots() {
+  inputs.classpathList.forEach {
+    ClasspathSnapshotGenerator(Paths.get(it), Paths.get("$it.snapshot"), SnapshotGranularity.CLASS_MEMBER_LEVEL).run()
+  }
+}
+
+internal fun JvmCompilationTask.createClasspathSnapshotsPaths(): List<String> {
+  return inputs.classpathList.map { it: String ->
+    "$it.snapshot"
+  }
+}
+
+
+val ROOT: String by lazy {
+  FileSystems
+    .getDefault()
+    .getPath("")
+    .toAbsolutePath()
+    .toString() + File.separator
+}
+
+
 /**
  * Compiles Kotlin sources to classes. Does not compile Java sources.
  */
@@ -441,9 +474,23 @@ fun JvmCompilationTask.compileKotlin(
           options = inputs.compilerPluginOptionsList,
           classpath = inputs.compilerPluginClasspathList,
         )
-    ).values(inputs.javaSourcesList)
+      )
+      .values(inputs.javaSourcesList)
       .values(inputs.kotlinSourcesList)
       .flag("-d", directories.classes)
+      .apply {
+        if (info.incrementalCompilation) {
+          flag("-incremental_id", "kotlin")
+          flag("-snapshot")
+          paths(
+            createClasspathSnapshotsPaths(),
+          ) {
+            it
+              .map(Path::toString)
+              .joinToString(File.pathSeparator)
+          }
+        }
+      }
       .list()
       .let {
         context.whenTracing {

src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/KotlinJvmTaskExecutor.kt

@@ -65,6 +65,7 @@ class KotlinJvmTaskExecutor
         preprocessedTask.apply {
           sequenceOf(
             runCatching {
+              context.execute("create classpath snapshots", ::createClasspathSnapshots)
               context.execute("kotlinc") {
                 if (compileKotlin) {
                   compileKotlin(

src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/SnapshotGranularity.kt

@@ -0,0 +1,15 @@
+package io.bazel.kotlin.builder.tasks.jvm
+
+import org.jetbrains.kotlin.buildtools.api.jvm.ClassSnapshotGranularity
+
+enum class SnapshotGranularity {
+    CLASS_LEVEL,
+    CLASS_MEMBER_LEVEL;
+
+    val toClassSnapshotGranularity: ClassSnapshotGranularity
+        get() =
+            when (this) {
+                CLASS_LEVEL -> ClassSnapshotGranularity.CLASS_LEVEL
+                CLASS_MEMBER_LEVEL -> ClassSnapshotGranularity.CLASS_MEMBER_LEVEL
+            }
+}

src/main/kotlin/io/bazel/kotlin/builder/toolchain/CompilerModes.kt

@@ -0,0 +1,21 @@
+package io.bazel.kotlin.builder.toolchain
+
+/**
+ * Constants for build tools API mode values
+ */
+object CompilerModes {
+  /**
+   * Use the Kotlin K2 compiler
+   */
+  const val K2 = "K2"
+
+  /**
+   * Use the BuildTools API compiler (default)
+   */
+  const val BUILD_TOOLS = "BUILD_TOOLS"
+
+  /**
+   * Use incremental compilation with BuildTools
+   */
+  const val INCREMENTAL = "INCREMENTAL"
+} 

src/main/kotlin/io/bazel/kotlin/builder/toolchain/KotlinToolchain.kt

@@ -132,6 +132,20 @@ class KotlinToolchain private constructor(
         ).toPath()
     }
 
+    private val KOTLIN_EMBEDDED by lazy {
+      BazelRunFiles
+        .resolveVerifiedFromProperty(
+          "@kotlin_compiler_embeddable...kotlin-embedded",
+        ).toPath()
+    }
+
+    private val KOTLIN_DAEMON by lazy {
+      BazelRunFiles
+        .resolveVerifiedFromProperty(
+          "@com_github_jetbrains_kotlin...daemon-client",
+        ).toPath()
+    }
+
     private val JAVA_HOME by lazy {
       FileSystems
         .getDefault()
@@ -161,6 +175,8 @@ class KotlinToolchain private constructor(
         KOTLINX_SERIALIZATION_CORE_JVM.toFile(),
         KOTLINX_SERIALIZATION_JSON.toFile(),
         KOTLINX_SERIALIZATION_JSON_JVM.toFile(),
+        KOTLIN_EMBEDDED.toFile(),
+        KOTLIN_DAEMON.toFile(),
       )
 
     @JvmStatic
@@ -178,6 +194,8 @@ class KotlinToolchain private constructor(
       kotlinxSerializationCoreJvm: File,
       kotlinxSerializationJson: File,
       kotlinxSerializationJsonJvm: File,
+      kotlinEmbedded: File,
+      kotlinDaemon: File,
     ): KotlinToolchain =
       KotlinToolchain(
         listOf(
@@ -195,6 +213,8 @@ class KotlinToolchain private constructor(
           kotlinxSerializationCoreJvm,
           kotlinxSerializationJson,
           kotlinxSerializationJsonJvm,
+          kotlinEmbedded,
+          kotlinDaemon,
         ),
         jvmAbiGen =
           CompilerPlugin(
@@ -286,6 +306,7 @@ class KotlinToolchain private constructor(
       val exitCodeClass =
         toolchain.classLoader.loadClass("org.jetbrains.kotlin.cli.common.ExitCode")
 
+      toolchain.classLoader.loadClass("org.jetbrains.kotlin.buildtools.internal.CompilationServiceProxy")
       compiler = compilerClass.getConstructor().newInstance()
       execMethod =
         compilerClass.getMethod("exec", PrintStream::class.java, Array<String>::class.java)