Make ktfmt up to 100x faster via native-image

.github/workflows/build_and_test.yml

@@ -4,6 +4,8 @@ on:
   push:
     paths-ignore:
       - '**.md'
+    branches:
+      - main
   pull_request:
     paths-ignore:
       - '**.md'
@@ -25,9 +27,19 @@ jobs:
       with:
         java-version: ${{ matrix.java }}
         distribution: zulu
+    - 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_idea_plugin
       run: ./gradlew :idea_plugin:build --stacktrace --no-daemon
     - name: Build the Online Formatter
       run: ./gradlew :lambda:build --stacktrace --no-daemon
     - name: Build ktfmt
       run: ./gradlew :ktfmt:build --stacktrace --no-daemon
+    - name: Build Native Image
+      run: ./gradlew :ktfmt:nativeCompile --stacktrace --no-daemon && ./ktfmt/build/native/nativeCompile/ktfmt --version

.gitignore

@@ -23,3 +23,7 @@ release-ktfmt-website/
 .gradle
 **/build/
 !src/**/build/
+
+# Native Image profiles are large
+core/src/main/native-image/profiles/*.iprof
+

build.gradle.kts

@@ -21,6 +21,7 @@ plugins {
   alias(libs.plugins.ktfmt) apply false
   alias(libs.plugins.nexusPublish)
   alias(libs.plugins.shadowJar) apply false
+  alias(libs.plugins.graalvm) apply false
 }
 
 version = providers.gradleProperty("ktfmt.version").get()

core/build.gradle.kts

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import java.nio.file.Paths
 import kotlin.io.path.writeText
 import org.jetbrains.intellij.platform.gradle.utils.asPath
 
@@ -22,15 +23,28 @@ plugins {
   id("com.gradleup.shadow")
   id("com.ncorti.ktfmt.gradle")
   id("maven-publish")
+  id("org.graalvm.buildtools.native")
   id("org.jetbrains.dokka")
   id("signing")
+  application
 }
 
+val entrypoint = "com.facebook.ktfmt.cli.Main"
+
+application { mainClass = entrypoint }
+
 repositories {
   mavenLocal()
   mavenCentral()
 }
 
+// Configuration for building `src/native-image/java`.
+val nativeImageJavacClasspath by
+    configurations.creating {
+      extendsFrom(configurations.implementation.get())
+      isCanBeResolved = true
+    }
+
 dependencies {
   api(libs.googleJavaformat)
   api(libs.guava)
@@ -40,6 +54,12 @@ dependencies {
   api(libs.kotlin.compilerEmbeddable)
   testImplementation(libs.googleTruth)
   testImplementation(libs.junit)
+
+  nativeImageJavacClasspath(libs.graalvm.nativeimage)
+  nativeImageClasspath(libs.jline.terminal)
+  nativeImageClasspath(libs.jline.terminal.jansi)
+  nativeImageClasspath(libs.jline.terminal.jna)
+  nativeImageClasspath(libs.jline.terminal.jni)
 }
 
 val generateSources by
@@ -116,6 +136,36 @@ tasks {
   }
 }
 
+sourceSets {
+  create("nativeImage") {
+    java { srcDir("src/main/native-image/java") }
+    resources { srcDir("src/main/native-image/resources") }
+  }
+}
+
+val compileNativeImageClasses by
+    tasks.registering(JavaCompile::class) {
+      group = "build"
+      description = "Compiles Native Image helper classes"
+      source = sourceSets["nativeImage"].java
+      classpath = nativeImageJavacClasspath
+      destinationDirectory = layout.buildDirectory.dir("classes/native-image")
+      dependsOn(tasks.named("compileJava"))
+    }
+
+// Native Image artifacts jar (local only, not published)
+val nativeImageJar by
+    tasks.registering(Jar::class) {
+      group = "build"
+      description = "Assembles Native Image jar and resources"
+      dependsOn(compileNativeImageClasses)
+      from(layout.buildDirectory.dir("classes/native-image"))
+      from(sourceSets["nativeImage"].resources)
+      archiveClassifier = "nativeimage"
+    }
+
+tasks.nativeCompile.configure { dependsOn(compileNativeImageClasses, nativeImageJar) }
+
 kotlin {
   val javaVersion: String = rootProject.libs.versions.java.get()
   jvmToolchain(javaVersion.toInt())
@@ -136,6 +186,181 @@ ktfmt {
   )
 }
 
+object DefaultArchitectureTarget {
+  val amd64 = "x86-64-v4"
+  val arm64 = "armv8.4-a+crypto+sve"
+}
+
+// Pass `-Pktfmt.native.release=true` to enable release mode for Native Image.
+val nativeRelease = findProperty("ktfmt.native.release") == "true"
+
+// Pass `-Pktfmt.native.target=xx` to pass `-march=xx` to Native Image.
+val nativeTarget =
+    findProperty("ktfmt.native.target")
+        ?: when (val hostArch = System.getProperty("os.arch")) {
+          "amd64",
+          "x86_64" -> DefaultArchitectureTarget.amd64
+          "aarch64",
+          "arm64" -> DefaultArchitectureTarget.arm64
+          else -> error("Unrecognized host architecture: '$hostArch'")
+        }
+
+// Pass `-Pktfmt.native.gc=xx` to select a garbage collector; options include `serial`, `G1`, and
+// `epsilon`.
+val nativeGc = findProperty("ktfmt.native.gc") ?: "G1"
+
+// Pass `-Pktfmt.native.gc=xx` to select a garbage collector; options include `serial`, `G1`, and
+// `epsilon`.
+val nativeDebug = findProperty("ktfmt.native.debug") == "true"
+
+// Pass `-Pktfmt.native.lto=true` to enable LTO for the Native Image binary.
+val enableLto = findProperty("ktfmt.native.lto") == "true"
+
+// Pass `-Pktfmt.native.muslHome=xx` or set MUSL_HOME to point to the Musl sysroot when building for
+// Musl Libc.
+val muslSysroot = (findProperty("ktfmt.native.muslHome") ?: System.getenv("MUSL_HOME"))?.toString()
+
+// Pass `-Pktfmt.native.musl=true` to build a fully-static binary against Musl Libc.
+val preferMusl =
+    (findProperty("ktfmt.native.musl") == "true").also { preferMusl ->
+      require(!preferMusl || 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/"
+      }
+    }
+
+// Pass `-Pktfmt.native.smol=true` to build a small, instead of a fast, binary.
+val preferSmol = (findProperty("ktfmt.native.smol") == "true")
+
+// Pass `-Pktfmt.native.opt=s` to pass e.g. `-Os` to Native Image.
+val nativeOpt =
+    when (val opt = findProperty("ktfmt.native.opt")) {
+      null ->
+          when {
+            preferSmol -> "s"
+            nativeRelease -> "3"
+            else -> "b" // prefer build speed
+          }
+      else -> opt
+    }
+
+// List of PGO profiles, which are held in `src/main/native-image/profiles`.
+val pgoProfiles =
+    listOf("default.iprof")
+        .map { profileName ->
+          layout.projectDirectory.file(
+              Paths.get("src", "main", "native-image", "profiles", profileName).toString()
+          )
+        }
+        .let { allProfiles -> listOf("--pgo=${allProfiles.joinToString(",")}") }
+
+// Pass `-Pktfmt.native.pgo=true` to build with PGO; pass `train` to enable instrumentation.
+val pgoArgs =
+    when (val pgo = findProperty("ktfmt.native.pgo")) {
+      null -> if (nativeRelease) pgoProfiles else emptyList()
+      "true" -> pgoProfiles
+      "false" -> emptyList()
+      "train" -> listOf("--pgo-instrument")
+      else -> error("Unrecognized `ktfmt.native.pgo` argument: '$pgo'")
+    }
+
+graalvmNative {
+  binaries {
+    named("main") {
+      imageName = "ktfmt"
+      mainClass = entrypoint
+      classpath =
+          files(
+              nativeImageJar.get().archiveFile,
+              tasks.jar.get().archiveFile,
+              configurations["compileClasspath"],
+              configurations["runtimeClasspath"],
+              configurations["nativeImageClasspath"],
+          )
+
+      buildArgs(
+          buildList {
+            // If PGO flags are present, add them first; if not, add `-Ox`.
+            when (pgoArgs.isEmpty()) {
+              true -> add("-O$nativeOpt")
+              false -> addAll(pgoArgs)
+            }
+
+            // Common flags for Native Image.
+            addAll(
+                buildList {
+                  add("-march=$nativeTarget")
+                  if (nativeDebug) {
+                    add("-g")
+                    add("-H:+SourceLevelDebug")
+                  }
+                  // --
+                  add("--gc=G1")
+                  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("--emit=build-report")
+                  add("--color=always")
+                  add("--enable-sbom=cyclonedx,embed")
+                  // -- ▼ SVM Hosted Options
+                  add("-H:+UseCompressedReferences")
+                  add("-H:+ReportExceptionStackTraces")
+                  // -- ▼ SVM Runtime Options
+                  add("-R:+InstallSegfaultHandler")
+                  // -- ▼ Experimental Options
+                  add("-H:+UnlockExperimentalVMOptions")
+                  add("-H:-ReduceImplicitExceptionStackTraceInformation")
+                  add("-H:+ReportDynamicAccess")
+                  add("-H:-UnlockExperimentalVMOptions")
+                  // -- ▼ VM flags
+                  add("-J--enable-native-access=ALL-UNNAMED")
+                  add("-J--illegal-native-access=allow")
+                  add("-J--sun-misc-unsafe-memory-access=allow")
+                  // -- ▼ C Compiler / Linker Flags
+                  if (enableLto) {
+                    add("--native-compiler-options=-flto")
+                    add("-H:NativeLinkerOption=-flto")
+                  }
+                  if (preferMusl) {
+                    add("-H:NativeLinkerOption=-L${muslSysroot}/lib")
+                  }
+                }
+            )
+
+            // Mark what should be initialized at build-time, i.e. persisted to the heap image.
+            // See `src/main/native-image/initialize-at-build-time.txt` for a list of such classes.
+            addLinesFromFile("src", "main", "native-image", "initialize-at-build-time.txt") {
+              "--initialize-at-build-time=$it"
+            }
+
+            // Still other classes must be initialized at runtime only.
+            // See `src/main/native-image/initialize-at-run-time.txt` for a list of such classes.
+            addLinesFromFile("src", "main", "native-image", "initialize-at-run-time.txt") {
+              "--initialize-at-run-time=$it"
+            }
+
+            // Here, we prefer static linking, for startup performance and release simplicity.
+            // On Linux amd64, we target musl to avoid linking conflicts with older glibc.
+            // On macOS, pass `--static-nolibc` for the closest option available.
+            when (System.getProperty("os.name")) {
+              "Linux" ->
+                  when (System.getProperty("os.arch")) {
+                    "amd64" ->
+                        when (preferMusl) {
+                          true -> addAll(listOf("--static", "--libc=musl", "-H:+StaticLibStdCpp"))
+                          false -> add("--static-nolibc")
+                        }
+                    else -> add("--static-nolibc")
+                  }
+              "Mac OS X" -> add("--static-nolibc")
+            }
+          }
+      )
+    }
+  }
+}
+
 group = "com.facebook"
 
 version = rootProject.version
@@ -180,3 +405,17 @@ if (System.getenv("SIGN_BUILD") != null) {
     sign(publishing.publications["maven"])
   }
 }
+
+fun MutableList<String>.addLinesFromFile(vararg path: String, mapper: (String) -> String) {
+  file(Paths.get(path.first(), *path.drop(1).toTypedArray()).toString())
+      .useLines { lines ->
+        lines
+            .filter { line ->
+              // filter empty lines
+              line.isNotEmpty()
+            }
+            .map(mapper)
+            .toList()
+      }
+      .also { addAll(it) }
+}

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,12 +52,11 @@ object Parser {
         CommonConfigurationKeys.MESSAGE_COLLECTOR_KEY,
         PrintingMessageCollector(System.err, PLAIN_RELATIVE_PATHS, false),
     )
-    env =
-        KotlinCoreEnvironment.createForProduction(
-            disposable,
-            configuration,
-            EnvironmentConfigFiles.JVM_CONFIG_FILES,
-        )
+    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,8 @@
+com.google.common.collect.SingletonImmutableList
+kotlin.Function
+kotlin.KotlinVersion
+kotlin.SynchronizedLazyImpl
+kotlin.UNINITIALIZED_VALUE
+kotlin.collections.EmptyMap
+kotlin.jvm.functions.Function1
+kotlin.text.Regex

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

@@ -0,0 +1,7 @@
+kotlin.random.AbstractPlatformRandom
+kotlin.random.Random
+kotlin.random.Random$Default
+kotlin.random.RandomKt
+kotlin.random.XorWowRandom
+kotlin.random.jdk8.PlatformThreadLocalRandom
+kotlin.uuid.SecureRandomHolder