Add experimental Kotlin incremental compilation

erikkerber · erikkerber · commit bcb43a9f9341 · 2025-06-09T22:48:52.000-05:00
diff --git a/MODULE.bazel b/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",
diff --git a/kotlin/internal/jvm/compile.bzl b/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)
diff --git a/kotlin/internal/toolchains.bzl b/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"],
diff --git a/kotlin/settings/BUILD.bazel b/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"],
+)
diff --git a/kotlin_rules_maven_install.json b/kotlin_rules_maven_install.json
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/BUILD.bazel b/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",
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/KotlinBuilder.kt b/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
       }
 
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/ClasspathSnapshotGenerator.kt b/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/ClasspathSnapshotGenerator.kt
@@ -0,0 +1,31 @@
+package io.bazel.kotlin.builder.tasks.jvm
+
+import java.nio.file.Path
+import kotlin.system.measureTimeMillis
+import org.jetbrains.kotlin.buildtools.api.CompilationService
+import org.jetbrains.kotlin.buildtools.api.ExperimentalBuildToolsApi
+import org.jetbrains.kotlin.buildtools.api.KotlinLogger
+import java.io.File
+import java.util.logging.Logger
+import kotlin.io.path.exists
+
+@OptIn(ExperimentalBuildToolsApi::class)
+class ClasspathSnapshotGenerator(
+    private val inputJar: Path,
+    private val outputSnapshot: Path,
+    private val granularity: SnapshotGranularity
+) {
+
+    fun run() {
+        if (outputSnapshot.exists()) { return }
+        val timeSpent = measureTimeMillis {
+            val compilationService = CompilationService.loadImplementation(this.javaClass.classLoader!!)
+            // val snapshot =
+            //     compilationService.calculateClasspathSnapshot(
+            //         inputJar.toFile(), granularity.toClassSnapshotGranularity)
+            // snapshot.saveSnapshot(outputSnapshot.toFile())
+        }
+        // TODO: Log impl
+        // LOG.info("$timeSpent ms for input jar: $inputJar")
+    }
+}
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/CompilationTask.kt b/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 {
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/KotlinJvmTaskExecutor.kt b/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(
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/SnapshotGranularity.kt b/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
+            }
+}
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/toolchain/CompilerModes.kt b/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"
+} 
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/toolchain/KotlinToolchain.kt b/src/main/kotlin/io/bazel/kotlin/builder/toolchain/KotlinToolchain.kt
@@ -286,6 +286,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)
diff --git a/src/main/kotlin/io/bazel/kotlin/compiler/BuildToolsAPICompiler.kt b/src/main/kotlin/io/bazel/kotlin/compiler/BuildToolsAPICompiler.kt
@@ -19,9 +19,20 @@ import org.jetbrains.kotlin.buildtools.api.CompilationResult
 import org.jetbrains.kotlin.buildtools.api.CompilationService
 import org.jetbrains.kotlin.buildtools.api.ExperimentalBuildToolsApi
 import org.jetbrains.kotlin.buildtools.api.ProjectId
+import org.jetbrains.kotlin.buildtools.api.SourcesChanges
+import org.jetbrains.kotlin.buildtools.api.jvm.ClasspathSnapshotBasedIncrementalCompilationApproachParameters
 import org.jetbrains.kotlin.cli.common.ExitCode
+import org.jetbrains.kotlin.incremental.createDirectory
+import java.io.File
+import java.nio.file.Paths
 import java.util.UUID
 
+fun stripNonKotlinArgument(arg: String, args: MutableList<String>): String? {
+    val index = args.indexOf(arg)
+    if (index < 0 || index + 1 >= args.size) return null
+    return args.removeAt(index + 1).also { args.removeAt(index) }
+}
+
 @Suppress("unused")
 class BuildToolsAPICompiler {
   @OptIn(ExperimentalBuildToolsApi::class)
@@ -31,23 +42,82 @@ class BuildToolsAPICompiler {
   ): ExitCode {
     System.setProperty("zip.handler.uses.crc.instead.of.timestamp", "true")
 
+    var mutableArgs = args.toMutableList()
+
+    // Set of vararg arguments into BuildToolsAPICompiler not to be passed to the compiler.
+    // Allows keeping the reflection-based invocation consistent with BazelK2KVMCompiler
+    val snapshots = stripNonKotlinArgument("-snapshot", mutableArgs)
+      ?.split(":")
+      ?.map { File(it) }
+      .orEmpty()
+
+    // A manufactured strategy for delineating running KSP plugins and compiling Kotlin.
+    // label: The Bazel label
+    // incrementalId: Used to delineate Kotlin tasks within a label (i.e. KSP vs Kotlin vs Kapt)
+    val label = stripNonKotlinArgument("-label", mutableArgs)!!
+    val incrementalId = stripNonKotlinArgument("-incremental_id", mutableArgs)!!
+
+    // The incremental directory is roughly analogous to a Gradle transform cache directory. Must be unique per compilation target.
+    val incrementalDirectory = Paths.get("").toAbsolutePath().parent.resolve("_kotlin_incremental/$label/$incrementalId").toFile()
+
     val kotlinService = CompilationService.loadImplementation(this.javaClass.classLoader!!)
     // The execution configuration. Controls in-process vs daemon execution strategies. Default is in-process.
     val executionConfig = kotlinService.makeCompilerExecutionStrategyConfiguration()
     // The compilation configuration. Controls everything related to incremental compilation.
-    val compilationConfig =
-      kotlinService.makeJvmCompilationConfiguration().apply {
-        // useLogger(BasicKotlinLogger(true, "/tmp/kotlin_log/$label.log"))
-      }
-
-    val result =
-      kotlinService.compileJvm(
-        ProjectId.ProjectUUID(UUID.randomUUID()),
-        executionConfig,
-        compilationConfig,
-        emptyList(),
-        args.toList(),
+    val compilationConfig = kotlinService.makeJvmCompilationConfiguration().apply {
+      useIncrementalCompilation(
+        // The working directory is where Kotlin stores incremental data. This must be unique per compilation target.
+        workingDirectory = incrementalDirectory,
+        // For Bazel, this will always be ToBeCalculated. We don't have a way to get the sources changes from the last build.
+        sourcesChanges = SourcesChanges.ToBeCalculated,
+        approachParameters = ClasspathSnapshotBasedIncrementalCompilationApproachParameters(
+          /**
+           * The classpath snapshots files actual at the moment of compilation
+           */
+          // TODO: Track this with Bazel
+          newClasspathSnapshotFiles = snapshots,
+
+          /**
+           * The shrunk classpath snapshot, a result of the previous compilation. Could point to a non-existent file.
+           * At the successful end of the compilation, the shrunk version of the [newClasspathSnapshotFiles] will be stored at this path.
+           */
+          // TODO: Track this with Bazel
+          shrunkClasspathSnapshot = incrementalDirectory.resolve("shrunk-classpath-snapshot.bin").apply { parentFile?.createDirectory() }
+        ),
+
+        options = makeClasspathSnapshotBasedIncrementalCompilationConfiguration().apply {
+
+          // NOTE: The following settings are to enable "relocatable" compilation caches. I was never able to
+          //       determine how to work with these without resulting in sporadic failures.
+          //       Notably, turning these on requires all input files to be passed with absolute, not relative paths.
+          // setRootProjectDir(incrementalDir.resolve("_main").toFile())
+          // setBuildDir(incrementalDir.resolve("_kotlin_incremental").toFile())
+          // useOutputDirs(emptyList())
+
+          // An indicator whether incremental compilation will analyze Java files precisely for better changes detection
+          usePreciseJavaTracking(true)
+
+          // Incremental compilation uses the PersistentHashMap of the intellij platform for storing caches.
+          // An indicator whether the changes should remain in memory and not being flushed to the disk until we could mark the compilation as successful.
+          keepIncrementalCompilationCachesInMemory(true)
+
+          // I don't believe we will ever need to force non-incremental
+          forceNonIncrementalMode(false)
+
+          // Classpath snapshots and changes will be tracked by Bazel
+          assureNoClasspathSnapshotsChanges(false)
+        }
       )
+      // useLogger(BasicKotlinLogger(true, "/tmp/kotlin_log/$label.log"))
+      useKotlinScriptFilenameExtensions(listOf("kts"))
+    }
+    val result = kotlinService.compileJvm(
+      ProjectId.ProjectUUID(UUID.randomUUID()),
+      executionConfig,
+      compilationConfig,
+      emptyList(),
+      mutableArgs.toList()
+    )
 
     // BTAPI returns a different type than K2JVMCompiler (CompilationResult vs ExitCode).
     return when (result) {
@@ -58,3 +128,4 @@ class BuildToolsAPICompiler {
     }
   }
 }
+
diff --git a/src/main/protobuf/kotlin_model.proto b/src/main/protobuf/kotlin_model.proto
diff --git a/src/main/starlark/core/repositories/initialize.release.bzl b/src/main/starlark/core/repositories/initialize.release.bzl
diff --git a/src/main/starlark/core/repositories/versions.bzl b/src/main/starlark/core/repositories/versions.bzl

Original file line number	Diff line number	Diff line change
`@@ -88,6 +88,7 @@ class KotlinBuilder`
`88`	`88`	`INSTRUMENT_COVERAGE("--instrument_coverage"),`
`89`	`89`	`KSP_GENERATED_JAVA_SRCJAR("--ksp_generated_java_srcjar"),`
`90`	`90`	`BUILD_TOOLS_API("--build_tools_api"),`
	`91`	`+ INCREMENTAL_COMPILATION("--incremental_compilation"),`
`91`	`92`	`}`
`92`	`93`	`}`
`93`	`94`
`@@ -173,6 +174,9 @@ class KotlinBuilder`
`173`	`174`	`argMap.optionalSingle(KotlinBuilderFlags.BUILD_TOOLS_API)?.let {`
`174`	`175`	`buildToolsApi = it == "true"`
`175`	`176`	`}`
	`177`	`+ argMap.optionalSingle(KotlinBuilderFlags.INCREMENTAL_COMPILATION)?.let {`
	`178`	`+ incrementalCompilation = it == "true"`
	`179`	`+ }`
`176`	`180`	`this`
`177`	`181`	`}`
`178`	`182`