Make ktfmt up to 100x faster via native-image

.gitattributes

@@ -0,0 +1 @@
+core/src/main/native-image/resources/META-INF/native-image/com/facebook/ktfmt/reachability-metadata.json linguist-generated=true

.github/workflows/build_and_test.yml

@@ -4,6 +4,9 @@ on:
   push:
     paths-ignore:
       - '**.md'
+    branches:
+      - main
+      - feat/native-image
   pull_request:
     paths-ignore:
       - '**.md'
@@ -25,5 +28,22 @@ jobs:
       with:
         java-version: ${{ matrix.java }}
         distribution: zulu
+    # set-java-home is intentionally 'false': the main build runs on the Zulu JDK from the matrix
+    # (Java 17/21), so we must not repoint JAVA_HOME at GraalVM 25. setup-graalvm still exports
+    # GRAALVM_HOME, which the GraalVM Gradle plugin uses to locate the native-image toolchain for the
+    # :ktfmt:nativeCompile step. This keeps the JVM build on the matrix JDK and the native build on
+    # GraalVM.
+    - name: Set up GraalVM
+      uses: graalvm/setup-graalvm@v1
+      with:
+        java-version: '25.0.1'
+        distribution: 'graalvm'
+        github-token: ${{ secrets.GITHUB_TOKEN }}
+        set-java-home: 'false'
+        native-image-job-reports: 'true'
     - name: Build ktfmt
       run: ./gradlew build --no-daemon
+    - name: Build Native Image
+      run: ./gradlew :ktfmt:nativeCompile --stacktrace --no-daemon
+    - name: Smoke test Native Image
+      run: ./native_smoke_test.sh

buildSrc/build.gradle.kts

@@ -22,12 +22,18 @@ plugins {
   alias(libs.plugins.ktfmt)
 }
 
+dependencies { implementation(nativeImageLibs.graalvm.gradle.plugin) }
+
 gradlePlugin {
   plugins {
     register("ktfmt-file-generator") {
       id = "ktfmt.ktfmt-file-generator"
       implementationClass = "com.facebook.ktfmt.GenerateKtfmtFilePlugin"
     }
+    register("native-image") {
+      id = "ktfmt.native-image"
+      implementationClass = "com.facebook.ktfmt.NativeImagePlugin"
+    }
   }
 }
 

buildSrc/settings.gradle.kts

@@ -21,7 +21,10 @@ dependencyResolutionManagement {
     gradlePluginPortal()
   }
 
-  versionCatalogs { create("libs") { from(files("../gradle/libs.versions.toml")) } }
+  versionCatalogs {
+    create("libs") { from(files("../gradle/libs.versions.toml")) }
+    create("nativeImageLibs") { from(files("../gradle/native-image.versions.toml")) }
+  }
 }
 
 rootProject.name = "buildSrc"

buildSrc/src/main/kotlin/com/facebook/ktfmt/NativeImagePlugin.kt

@@ -0,0 +1,251 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.ktfmt
+
+import org.graalvm.buildtools.gradle.dsl.GraalVMExtension
+import org.gradle.api.GradleException
+import org.gradle.api.Plugin
+import org.gradle.api.Project
+import org.gradle.api.artifacts.Configuration
+import org.gradle.api.artifacts.VersionCatalogsExtension
+import org.gradle.api.plugins.JavaApplication
+import org.gradle.api.plugins.JavaPluginExtension
+import org.gradle.jvm.tasks.Jar
+import org.gradle.kotlin.dsl.configure
+import org.gradle.kotlin.dsl.getByType
+import org.gradle.kotlin.dsl.named
+import org.gradle.kotlin.dsl.register
+
+@Suppress("unused")
+class NativeImagePlugin : Plugin<Project> {
+  override fun apply(project: Project) {
+    project.plugins.apply("application")
+    project.plugins.apply("org.graalvm.buildtools.native")
+
+    project.extensions.configure<JavaApplication> { mainClass.set(ENTRYPOINT) }
+
+    project.plugins.withId("java") { configureNativeImage(project) }
+  }
+
+  private fun configureNativeImage(project: Project) {
+    val nativeImageLibs =
+        project.extensions.getByType<VersionCatalogsExtension>().named("nativeImageLibs")
+
+    val nativeImageJavacClasspath: Configuration =
+        project.configurations.create("nativeImageJavacClasspath") {
+          extendsFrom(project.configurations.getByName("implementation"))
+          isCanBeResolved = true
+        }
+
+    project.dependencies.apply {
+      add("nativeImageJavacClasspath", nativeImageLibs.findLibrary("graalvm-nativeimage").get())
+      add("nativeImageClasspath", nativeImageLibs.findLibrary("jline-terminal").get())
+      add("nativeImageClasspath", nativeImageLibs.findLibrary("jline-terminal-jansi").get())
+      add("nativeImageClasspath", nativeImageLibs.findLibrary("jline-terminal-jna").get())
+      add("nativeImageClasspath", nativeImageLibs.findLibrary("jline-terminal-jni").get())
+    }
+
+    val nativeImageDir = project.layout.projectDirectory.dir(NATIVE_IMAGE_SRC_DIR)
+    val javaExtension = project.extensions.getByType<JavaPluginExtension>()
+    val nativeImageSourceSet =
+        javaExtension.sourceSets.create("nativeImage") {
+          java.srcDir(nativeImageDir.dir("java"))
+          resources.srcDir(nativeImageDir.dir("resources"))
+        }
+
+    val compileNativeImageClasses =
+        project.tasks.register(
+            "compileNativeImageClasses",
+            org.gradle.api.tasks.compile.JavaCompile::class,
+        ) {
+          group = "build"
+          description = "Compiles Native Image helper classes"
+          source = nativeImageSourceSet.java
+          classpath = nativeImageJavacClasspath
+          destinationDirectory.set(project.layout.buildDirectory.dir("classes/native-image"))
+          dependsOn(project.tasks.named("compileJava"))
+        }
+
+    val nativeImageJar =
+        project.tasks.register("nativeImageJar", Jar::class) {
+          group = "build"
+          description = "Assembles Native Image jar and resources"
+          dependsOn(compileNativeImageClasses)
+          from(project.layout.buildDirectory.dir("classes/native-image"))
+          from(nativeImageSourceSet.resources)
+          archiveClassifier.set("nativeimage")
+        }
+
+    project.tasks.named("nativeCompile") { dependsOn(compileNativeImageClasses, nativeImageJar) }
+
+    configureGraalvmNative(project, nativeImageJar)
+  }
+
+  private fun configureGraalvmNative(
+      project: Project,
+      nativeImageJar: org.gradle.api.tasks.TaskProvider<Jar>,
+  ) {
+    val nativeRelease = project.findProperty("ktfmt.native.release") == "true"
+    val nativeTarget = project.findProperty("ktfmt.native.target") ?: "compatibility"
+    val nativeGc = project.findProperty("ktfmt.native.gc") ?: "serial"
+    val nativeDebug = project.findProperty("ktfmt.native.debug") == "true"
+    val enableLto = project.findProperty("ktfmt.native.lto") == "true"
+    val muslSysroot =
+        (project.findProperty("ktfmt.native.muslHome") ?: System.getenv("MUSL_HOME"))?.toString()
+    val preferMusl =
+        (project.findProperty("ktfmt.native.musl") == "true").also { enabled ->
+          require(!enabled || muslSysroot != null) {
+            "When `ktfmt.native.musl` is true, -Pktfmt.native.muslHome or MUSL_HOME must be set to the Musl sysroot. " +
+                "See https://www.graalvm.org/latest/reference-manual/native-image/guides/build-static-executables/"
+          }
+        }
+    val preferSmol = project.findProperty("ktfmt.native.smol") == "true"
+    val nativeOpt =
+        when (val opt = project.findProperty("ktfmt.native.opt")) {
+          null ->
+              when {
+                preferSmol -> "s"
+                nativeRelease -> "3"
+                else -> "b"
+              }
+          else -> opt
+        }
+
+    project.extensions.configure<GraalVMExtension>("graalvmNative") {
+      binaries.named("main") {
+        imageName.set("ktfmt")
+        mainClass.set(ENTRYPOINT)
+        classpath(
+            project.files(
+                nativeImageJar.flatMap { it.archiveFile },
+                project.tasks.named("jar", Jar::class).flatMap { it.archiveFile },
+                project.configurations.getByName("compileClasspath"),
+                project.configurations.getByName("runtimeClasspath"),
+                project.configurations.getByName("nativeImageClasspath"),
+            )
+        )
+        buildArgs(
+            buildNativeImageArgs(
+                project,
+                nativeOpt,
+                nativeTarget,
+                nativeDebug,
+                nativeGc,
+                enableLto,
+                preferMusl,
+                muslSysroot,
+            )
+        )
+      }
+    }
+  }
+
+  private fun buildNativeImageArgs(
+      project: Project,
+      nativeOpt: Any,
+      nativeTarget: Any,
+      nativeDebug: Boolean,
+      nativeGc: Any,
+      enableLto: Boolean,
+      preferMusl: Boolean,
+      muslSysroot: String?,
+  ): List<String> = buildList {
+    add("-O$nativeOpt")
+    add("-march=$nativeTarget")
+    if (nativeDebug) {
+      add("-g")
+      add("-H:+SourceLevelDebug")
+    }
+
+    add("--no-fallback")
+    add("--gc=$nativeGc")
+    add("--future-defaults=all")
+    add("--link-at-build-time=com.facebook")
+    add("--initialize-at-build-time=com.facebook")
+    add("--add-opens=java.base/java.util=ALL-UNNAMED")
+    add("--color=always")
+    add("-H:+ReportExceptionStackTraces")
+    add("-H:-UseContainerSupport")
+    add("-R:+InstallSegfaultHandler")
+    add("-H:+UnlockExperimentalVMOptions")
+    add("-H:-ReduceImplicitExceptionStackTraceInformation")
+    add("-H:-UnlockExperimentalVMOptions")
+    add("-J--enable-native-access=ALL-UNNAMED")
+    add("-J--illegal-native-access=allow")
+    add("-J--sun-misc-unsafe-memory-access=allow")
+
+    if (enableLto) {
+      add("--native-compiler-options=-flto")
+      add("-H:NativeLinkerOption=-flto")
+    }
+    if (preferMusl) {
+      add("-H:NativeLinkerOption=-L${muslSysroot}/lib")
+    }
+
+    addLinesFromFile(project, "initialize-at-build-time.txt") { "--initialize-at-build-time=$it" }
+    addLinesFromFile(project, "initialize-at-run-time.txt") { "--initialize-at-run-time=$it" }
+
+    when (System.getProperty("os.name")) {
+      "Linux" ->
+          when (normalizeArch(System.getProperty("os.arch"))) {
+            "x64" ->
+                if (preferMusl) {
+                  addAll(listOf("--static", "--libc=musl", "-H:+StaticLibStdCpp"))
+                } else {
+                  add("--static-nolibc")
+                }
+            else -> add("--static-nolibc")
+          }
+      "Mac OS X" -> add("--static-nolibc")
+    }
+  }
+
+  /** Canonicalizes [java.lang.System.getProperty]("os.arch") values that vary across JVMs. */
+  private fun normalizeArch(arch: String?): String =
+      when (arch) {
+        "amd64",
+        "x86_64" -> "x64"
+        "aarch64",
+        "arm64" -> "aarch64"
+        else -> arch.orEmpty()
+      }
+
+  private fun MutableList<String>.addLinesFromFile(
+      project: Project,
+      fileName: String,
+      mapper: (String) -> String,
+  ) {
+    val file = project.layout.projectDirectory.dir(NATIVE_IMAGE_SRC_DIR).file(fileName).asFile
+    if (!file.exists()) {
+      throw GradleException("Native Image configuration file not found: $file")
+    }
+    file
+        .useLines { lines ->
+          lines
+              .map { it.trim() }
+              .filter { it.isNotEmpty() && !it.startsWith("#") }
+              .map(mapper)
+              .toList()
+        }
+        .also { addAll(it) }
+  }
+
+  private companion object {
+    const val ENTRYPOINT = "com.facebook.ktfmt.cli.Main"
+    const val NATIVE_IMAGE_SRC_DIR = "src/main/native-image"
+  }
+}

core/build.gradle.kts

@@ -28,6 +28,7 @@ plugins {
   id("maven-publish")
   id("signing")
   id("ktfmt.ktfmt-file-generator")
+  id("ktfmt.native-image")
 }
 
 repositories {

core/src/main/java/com/facebook/ktfmt/format/Parser.kt

@@ -43,9 +43,7 @@ object Parser {
    * from [KotlinCoreEnvironment.createForProduction]:
    * https://github.com/JetBrains/kotlin/blob/master/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt#L544
    */
-  val env: KotlinCoreEnvironment
-
-  init {
+  val env: KotlinCoreEnvironment by lazy {
     // To hide annoying warning on Windows
     System.setProperty("idea.use.native.fs.for.win", "false")
     val disposable = Disposer.newDisposable()
@@ -54,13 +52,12 @@ object Parser {
         CommonConfigurationKeys.MESSAGE_COLLECTOR_KEY,
         PrintingMessageCollector(System.err, PLAIN_RELATIVE_PATHS, false),
     )
-    env =
-        @Suppress("OPT_IN_USAGE_ERROR") // KotlinCoreEnvironment.createForProduction
-        KotlinCoreEnvironment.createForProduction(
-            disposable,
-            configuration,
-            EnvironmentConfigFiles.JVM_CONFIG_FILES,
-        )
+    @Suppress("OPT_IN_USAGE_ERROR") // KotlinCoreEnvironment.createForProduction
+    KotlinCoreEnvironment.createForProduction(
+        disposable,
+        configuration,
+        EnvironmentConfigFiles.JVM_CONFIG_FILES,
+    )
   }
 
   fun parse(code: String): KtFile {

core/src/main/native-image/initialize-at-build-time.txt

@@ -0,0 +1,10 @@
+com.google.common.collect.SingletonImmutableList
+kotlin.Function
+kotlin.KotlinVersion
+kotlin.SynchronizedLazyImpl
+kotlin.UNINITIALIZED_VALUE
+kotlin.collections.AbstractList$Companion
+kotlin.collections.EmptyMap
+kotlin.enums.EnumEntriesList
+kotlin.jvm.functions.Function1
+kotlin.text.Regex

core/src/main/native-image/initialize-at-run-time.txt

@@ -0,0 +1,12 @@
+# Some ktfmt classes build run-time-only object graphs in their <clinit> (ec4j config types, Kotlin
+# compiler PSI stub element types); keep them out of the build-time image heap. The rest of
+# com.facebook is initialized at build time for startup speed.
+com.facebook.ktfmt.cli.EditorConfigResolver
+com.facebook.ktfmt.util.CompatibilityUtilsKt
+kotlin.random.AbstractPlatformRandom
+kotlin.random.Random
+kotlin.random.Random$Default
+kotlin.random.RandomKt
+kotlin.random.XorWowRandom
+kotlin.random.jdk8.PlatformThreadLocalRandom
+kotlin.uuid.SecureRandomHolder