[BT] Implement classpath snapshotting in the Build Tools API

#KT-57565 In Progress

compiler/build-tools/kotlin-build-tools-api/api/kotlin-build-tools-api.api

@@ -9,7 +9,7 @@ public final class org/jetbrains/kotlin/buildtools/api/CompilationResult : java/
 
 public abstract interface class org/jetbrains/kotlin/buildtools/api/CompilationService {
 	public static final field Companion Lorg/jetbrains/kotlin/buildtools/api/CompilationService$Companion;
-	public abstract fun calculateClasspathSnapshot (Ljava/io/File;)Lorg/jetbrains/kotlin/buildtools/api/jvm/ClasspathEntrySnapshot;
+	public abstract fun calculateClasspathSnapshot (Ljava/io/File;Lorg/jetbrains/kotlin/buildtools/api/jvm/ClassSnapshotGranularity;)Lorg/jetbrains/kotlin/buildtools/api/jvm/ClasspathEntrySnapshot;
 	public abstract fun compileJvm (Lorg/jetbrains/kotlin/buildtools/api/ProjectId;Lorg/jetbrains/kotlin/buildtools/api/CompilerExecutionStrategyConfiguration;Lorg/jetbrains/kotlin/buildtools/api/jvm/JvmCompilationConfiguration;Ljava/util/List;Ljava/util/List;)Lorg/jetbrains/kotlin/buildtools/api/CompilationResult;
 	public abstract fun finishProjectCompilation (Lorg/jetbrains/kotlin/buildtools/api/ProjectId;)V
 	public static fun loadImplementation (Ljava/lang/ClassLoader;)Lorg/jetbrains/kotlin/buildtools/api/CompilationService;
@@ -82,6 +82,10 @@ public final class org/jetbrains/kotlin/buildtools/api/SourcesChanges$Unknown :
 	public static final field INSTANCE Lorg/jetbrains/kotlin/buildtools/api/SourcesChanges$Unknown;
 }
 
+public abstract interface class org/jetbrains/kotlin/buildtools/api/jvm/AccessibleClassSnapshot : org/jetbrains/kotlin/buildtools/api/jvm/ClassSnapshot {
+	public abstract fun getClassAbiHash ()J
+}
+
 public abstract interface class org/jetbrains/kotlin/buildtools/api/jvm/ClassSnapshot {
 }
 
@@ -93,7 +97,7 @@ public final class org/jetbrains/kotlin/buildtools/api/jvm/ClassSnapshotGranular
 }
 
 public abstract interface class org/jetbrains/kotlin/buildtools/api/jvm/ClasspathEntrySnapshot {
-	public abstract fun getClassSnapshots ()Ljava/util/LinkedHashMap;
+	public abstract fun getClassSnapshots ()Ljava/util/Map;
 	public abstract fun saveSnapshot (Ljava/io/File;)V
 }
 
@@ -115,6 +119,9 @@ public abstract interface class org/jetbrains/kotlin/buildtools/api/jvm/Classpat
 	public abstract fun useProjectDir (Ljava/io/File;)Lorg/jetbrains/kotlin/buildtools/api/jvm/ClasspathSnapshotBasedIncrementalJvmCompilationConfiguration;
 }
 
+public abstract interface class org/jetbrains/kotlin/buildtools/api/jvm/InaccessibleClassSnapshot : org/jetbrains/kotlin/buildtools/api/jvm/ClassSnapshot {
+}
+
 public abstract interface class org/jetbrains/kotlin/buildtools/api/jvm/IncrementalCompilationApproachParameters {
 }
 

compiler/build-tools/kotlin-build-tools-api/src/main/kotlin/org/jetbrains/kotlin/buildtools/api/CompilationService.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.buildtools.api
 
+import org.jetbrains.kotlin.buildtools.api.jvm.ClassSnapshotGranularity
 import org.jetbrains.kotlin.buildtools.api.jvm.ClasspathEntrySnapshot
 import org.jetbrains.kotlin.buildtools.api.jvm.JvmCompilationConfiguration
 import java.io.File
@@ -28,7 +29,7 @@ public interface CompilationService {
     /**
      * TODO KT-57565
      */
-    public fun calculateClasspathSnapshot(classpathEntry: File): ClasspathEntrySnapshot
+    public fun calculateClasspathSnapshot(classpathEntry: File, granularity: ClassSnapshotGranularity): ClasspathEntrySnapshot
 
     /**
      * Provides a default [CompilerExecutionStrategyConfiguration] allowing to use it as is or customizing for specific requirements.

compiler/build-tools/kotlin-build-tools-api/src/main/kotlin/org/jetbrains/kotlin/buildtools/api/jvm/ClasspathEntrySnapshot.kt

@@ -15,7 +15,7 @@ import java.io.File
  */
 @ExperimentalBuildToolsApi
 public interface ClasspathEntrySnapshot {
-    public val classSnapshots: LinkedHashMap<String, ClassSnapshot>
+    public val classSnapshots: Map<String, ClassSnapshot>
 
     public fun saveSnapshot(path: File)
 }
@@ -26,6 +26,18 @@ public interface ClasspathEntrySnapshot {
  * This interface is not intended to be implemented by the API consumers.
  */
 @ExperimentalBuildToolsApi
-public interface ClassSnapshot {
-    // ... TODO: KT-57565, it will expose some part of org.jetbrains.kotlin.incremental.classpathDiff.ClassSnapshot hierarchy
+public sealed interface ClassSnapshot
+
+/**
+ * [ClassSnapshot] of an inaccessible class.
+ *
+ * A class is inaccessible if it can't be referenced from other source files (and therefore any changes in an inaccessible class will not
+ * require recompilation of other source files).
+ */
+@ExperimentalBuildToolsApi
+public interface InaccessibleClassSnapshot : ClassSnapshot
+
+@ExperimentalBuildToolsApi
+public interface AccessibleClassSnapshot : ClassSnapshot {
+    public val classAbiHash: Long
 }

compiler/build-tools/kotlin-build-tools-impl/build.gradle.kts

@@ -13,13 +13,18 @@ dependencies {
     compileOnly(project(":kotlin-daemon-client"))
     compileOnly(project(":compiler:incremental-compilation-impl"))
     compileOnly(project(":kotlin-compiler-runner-unshaded"))
+    compileOnly(intellijCore())
     runtimeOnly(project(":kotlin-compiler-embeddable"))
     runtimeOnly(project(":kotlin-compiler-runner"))
 }
 
 publish()
 
-standardPublicJars()
+runtimeJar(rewriteDefaultJarDepsToShadedCompiler {
+    from(mainSourceSet.output)
+})
+sourcesJar()
+javadocJar()
 
 kotlin {
     compilerOptions {

compiler/build-tools/kotlin-build-tools-impl/src/main/kotlin/org/jetbrains/kotlin/buildtools/internal/ClassSnapshotImpl.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.buildtools.internal
+
+import org.jetbrains.kotlin.buildtools.api.jvm.AccessibleClassSnapshot
+import org.jetbrains.kotlin.buildtools.api.jvm.InaccessibleClassSnapshot
+
+class AccessibleClassSnapshotImpl(override val classAbiHash: Long) : AccessibleClassSnapshot
+
+object InaccessibleClassSnapshotImpl : InaccessibleClassSnapshot

compiler/build-tools/kotlin-build-tools-impl/src/main/kotlin/org/jetbrains/kotlin/buildtools/internal/ClasspathEntrySnapshotImpl.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.buildtools.internal
+
+import org.jetbrains.kotlin.buildtools.api.jvm.ClassSnapshot
+import org.jetbrains.kotlin.buildtools.api.jvm.ClasspathEntrySnapshot
+import org.jetbrains.kotlin.incremental.classpathDiff.AccessibleClassSnapshot
+import org.jetbrains.kotlin.incremental.classpathDiff.ClasspathEntrySnapshotExternalizer
+import org.jetbrains.kotlin.incremental.classpathDiff.InaccessibleClassSnapshot
+import org.jetbrains.kotlin.incremental.storage.saveToFile
+import java.io.File
+
+class ClasspathEntrySnapshotImpl(
+    private val origin: org.jetbrains.kotlin.incremental.classpathDiff.ClasspathEntrySnapshot,
+) :
+    ClasspathEntrySnapshot {
+    override val classSnapshots: Map<String, ClassSnapshot>
+        get() = origin.classSnapshots.mapValues {
+            when (val snapshot = it.value) {
+                is AccessibleClassSnapshot -> AccessibleClassSnapshotImpl(snapshot.classAbiHash)
+                is InaccessibleClassSnapshot -> InaccessibleClassSnapshotImpl
+            }
+        }
+
+    override fun saveSnapshot(path: File) {
+        ClasspathEntrySnapshotExternalizer.saveToFile(path, origin)
+    }
+}

compiler/build-tools/kotlin-build-tools-impl/src/main/kotlin/org/jetbrains/kotlin/buildtools/internal/CompilationServiceImpl.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.buildtools.internal
 
+import org.jetbrains.kotlin.build.report.metrics.DoNothingBuildMetricsReporter
 import org.jetbrains.kotlin.buildtools.api.*
 import org.jetbrains.kotlin.buildtools.api.jvm.*
 import org.jetbrains.kotlin.cli.common.ExitCode
@@ -18,6 +19,7 @@ import org.jetbrains.kotlin.daemon.client.BasicCompilerServicesWithResultsFacade
 import org.jetbrains.kotlin.daemon.common.CompilerId
 import org.jetbrains.kotlin.daemon.common.configureDaemonJVMOptions
 import org.jetbrains.kotlin.daemon.common.filterExtractProps
+import org.jetbrains.kotlin.incremental.classpathDiff.ClasspathEntrySnapshotter
 import java.io.File
 import java.net.URLClassLoader
 import java.util.concurrent.ConcurrentHashMap
@@ -36,9 +38,8 @@ private fun getCurrentClasspath() = (CompilationServiceImpl::class.java.classLoa
 internal object CompilationServiceImpl : CompilationService {
     private val buildIdToSessionFlagFile: MutableMap<ProjectId, File> = ConcurrentHashMap()
 
-    override fun calculateClasspathSnapshot(classpathEntry: File): ClasspathEntrySnapshot {
-        TODO("Calculating classpath snapshots via the Build Tools API is not yet implemented: KT-57565")
-    }
+    override fun calculateClasspathSnapshot(classpathEntry: File, granularity: ClassSnapshotGranularity) =
+        ClasspathEntrySnapshotImpl(ClasspathEntrySnapshotter.snapshot(classpathEntry, granularity, DoNothingBuildMetricsReporter))
 
     override fun makeCompilerExecutionStrategyConfiguration() = CompilerExecutionStrategyConfigurationImpl()
 

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/compilerRunner/btapi/BuildToolsApiCompilationWork.kt

@@ -142,7 +142,7 @@ internal abstract class BuildToolsApiCompilationWork : WorkAction<BuildToolsApiC
 }
 
 @OptIn(ExperimentalBuildToolsApi::class)
-private object SharedApiClassesClassLoaderProvider : ParentClassLoaderProvider {
+internal object SharedApiClassesClassLoaderProvider : ParentClassLoaderProvider {
     override fun getClassLoader() = SharedApiClassesClassLoader()
 
     override fun hashCode() = SharedApiClassesClassLoaderProvider::class.hashCode()

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/internal/transforms/BuildToolsApiClasspathEntrySnapshotTransform.kt

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.gradle.internal.transforms
+
+import org.gradle.api.artifacts.transform.*
+import org.gradle.api.file.ConfigurableFileCollection
+import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.FileSystemLocation
+import org.gradle.api.provider.Property
+import org.gradle.api.provider.Provider
+import org.gradle.api.tasks.Classpath
+import org.gradle.api.tasks.Internal
+import org.jetbrains.kotlin.buildtools.api.CompilationService
+import org.jetbrains.kotlin.buildtools.api.ExperimentalBuildToolsApi
+import org.jetbrains.kotlin.buildtools.api.jvm.ClassSnapshotGranularity
+import org.jetbrains.kotlin.buildtools.api.jvm.ClassSnapshotGranularity.CLASS_LEVEL
+import org.jetbrains.kotlin.buildtools.api.jvm.ClassSnapshotGranularity.CLASS_MEMBER_LEVEL
+import org.jetbrains.kotlin.compilerRunner.btapi.SharedApiClassesClassLoaderProvider
+import org.jetbrains.kotlin.gradle.internal.ClassLoadersCachingBuildService
+import java.io.File
+
+/** Transform to create a snapshot of a classpath entry (directory or jar). */
+@CacheableTransform
+@OptIn(ExperimentalBuildToolsApi::class)
+abstract class BuildToolsApiClasspathEntrySnapshotTransform : TransformAction<BuildToolsApiClasspathEntrySnapshotTransform.Parameters> {
+
+    abstract class Parameters : TransformParameters {
+        @get:Internal
+        abstract val gradleUserHomeDir: DirectoryProperty
+
+        @get:Internal
+        internal abstract val classLoadersCachingService: Property<ClassLoadersCachingBuildService>
+
+        @get:Classpath
+        internal abstract val classpath: ConfigurableFileCollection
+    }
+
+    @get:Classpath
+    @get:InputArtifact
+    abstract val inputArtifact: Provider<FileSystemLocation>
+
+    override fun transform(outputs: TransformOutputs) {
+        val classpathEntryInputDirOrJar = inputArtifact.get().asFile
+        val snapshotOutputFile = outputs.file(classpathEntryInputDirOrJar.name.replace('.', '_') + "-snapshot.bin")
+
+        val granularity = getClassSnapshotGranularity(classpathEntryInputDirOrJar, parameters.gradleUserHomeDir.get().asFile)
+
+        try {
+            val classLoader = parameters.classLoadersCachingService.get()
+                .getClassLoader(parameters.classpath, SharedApiClassesClassLoaderProvider)
+            val compilationService = CompilationService.loadImplementation(classLoader)
+            val snapshot = compilationService.calculateClasspathSnapshot(classpathEntryInputDirOrJar, granularity)
+            snapshot.saveSnapshot(snapshotOutputFile)
+        } catch (e: Throwable) {
+            throw e
+        }
+    }
+
+    /**
+     * Determines the [ClassSnapshotGranularity] when taking a snapshot of the given [classpathEntryDirOrJar].
+     *
+     * As mentioned in [ClassSnapshotGranularity]'s kdoc, we will take [CLASS_LEVEL] snapshots for classes that are infrequently changed
+     * (e.g., external libraries which are typically stored/transformed inside the Gradle user home, or a few hard-coded cases), and take
+     * [CLASS_MEMBER_LEVEL] snapshots for the others.
+     */
+    private fun getClassSnapshotGranularity(classpathEntryDirOrJar: File, gradleUserHomeDir: File): ClassSnapshotGranularity {
+        return if (
+            classpathEntryDirOrJar.startsWith(gradleUserHomeDir) ||
+            classpathEntryDirOrJar.name == "android.jar"
+        ) CLASS_LEVEL
+        else CLASS_MEMBER_LEVEL
+    }
+}

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/tasks/configuration/KotlinCompileConfig.kt

@@ -6,11 +6,17 @@
 package org.jetbrains.kotlin.gradle.tasks.configuration
 
 import org.gradle.api.Project
+import org.gradle.api.artifacts.transform.TransformSpec
 import org.gradle.api.attributes.Attribute
+import org.gradle.api.file.FileCollection
+import org.gradle.api.provider.Provider
 import org.jetbrains.kotlin.gradle.dsl.KotlinAndroidProjectExtension
 import org.jetbrains.kotlin.gradle.dsl.KotlinJvmProjectExtension
 import org.jetbrains.kotlin.gradle.dsl.KotlinTopLevelExtension
+import org.jetbrains.kotlin.gradle.internal.ClassLoadersCachingBuildService
+import org.jetbrains.kotlin.gradle.internal.transforms.BuildToolsApiClasspathEntrySnapshotTransform
 import org.jetbrains.kotlin.gradle.internal.transforms.ClasspathEntrySnapshotTransform
+import org.jetbrains.kotlin.gradle.plugin.BUILD_TOOLS_API_CLASSPATH_CONFIGURATION_NAME
 import org.jetbrains.kotlin.gradle.plugin.KotlinCompilationInfo
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmAndroidCompilation
 import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmCompilation
@@ -18,6 +24,7 @@ import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinWithJavaCompilation
 import org.jetbrains.kotlin.gradle.utils.markResolvable
 import org.jetbrains.kotlin.gradle.plugin.tcsOrNull
 import org.jetbrains.kotlin.gradle.report.BuildMetricsService
+import org.jetbrains.kotlin.gradle.tasks.DefaultKotlinJavaToolchain
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 import org.jetbrains.kotlin.gradle.utils.providerWithLazyConvention
 
@@ -29,7 +36,9 @@ internal open class BaseKotlinCompileConfig<TASK : KotlinCompile> : AbstractKotl
         configureTaskProvider { taskProvider ->
             val useClasspathSnapshot = propertiesProvider.useClasspathSnapshot
             val classpathConfiguration = if (useClasspathSnapshot) {
-                registerTransformsOnce(project)
+                val runKotlinCompilerViaBuildToolsApi = propertiesProvider.runKotlinCompilerViaBuildToolsApi
+                val jvmToolchain = taskProvider.flatMap { it.defaultKotlinJavaToolchain }
+                registerTransformsOnce(project, jvmToolchain, runKotlinCompilerViaBuildToolsApi)
                 // Note: Creating configurations should be done during build configuration, not task configuration, to avoid issues with
                 // composite builds (e.g., https://issuetracker.google.com/183952598).
                 project.configurations.detachedConfiguration(
@@ -108,12 +117,24 @@ internal open class BaseKotlinCompileConfig<TASK : KotlinCompile> : AbstractKotl
         const val CLASSPATH_ENTRY_SNAPSHOT_ARTIFACT_TYPE = "classpath-entry-snapshot"
     }
 
-    private fun registerTransformsOnce(project: Project) {
+    private fun registerTransformsOnce(
+        project: Project,
+        jvmToolchain: Provider<DefaultKotlinJavaToolchain>,
+        runKotlinCompilerViaBuildToolsApi: Boolean,
+    ) {
         if (project.extensions.extraProperties.has(TRANSFORMS_REGISTERED)) {
             return
         }
         project.extensions.extraProperties[TRANSFORMS_REGISTERED] = true
 
+        if (runKotlinCompilerViaBuildToolsApi) {
+            registerBuildToolsApiTransformations(project, jvmToolchain)
+        } else {
+            registerOldTransformations(project)
+        }
+    }
+
+    private fun registerOldTransformations(project: Project) {
         val buildMetricsService = BuildMetricsService.registerIfAbsent(project)
         project.dependencies.registerTransform(ClasspathEntrySnapshotTransform::class.java) {
             it.from.attribute(ARTIFACT_TYPE_ATTRIBUTE, JAR_ARTIFACT_TYPE)
@@ -129,4 +150,37 @@ internal open class BaseKotlinCompileConfig<TASK : KotlinCompile> : AbstractKotl
         }
     }
 
+    private fun TransformSpec<BuildToolsApiClasspathEntrySnapshotTransform.Parameters>.configureCommonParameters(
+        classLoadersCachingService: Provider<ClassLoadersCachingBuildService>,
+        classpath: Provider<out FileCollection>,
+        jvmToolchain: Provider<DefaultKotlinJavaToolchain>,
+    ) {
+        parameters.gradleUserHomeDir.set(project.gradle.gradleUserHomeDir)
+        parameters.classLoadersCachingService.set(classLoadersCachingService)
+        parameters.classpath.from(classpath)
+        // add some tools.jar in order to reuse some of the classloaders required for compilation
+        parameters.classpath.from(jvmToolchain.map { toolchain ->
+            if (toolchain.currentJvmJdkToolsJar.isPresent) {
+                setOf(toolchain.currentJvmJdkToolsJar.get())
+            } else {
+                emptySet()
+            }
+        })
+    }
+
+    private fun registerBuildToolsApiTransformations(project: Project, jvmToolchain: Provider<DefaultKotlinJavaToolchain>) {
+        val classLoadersCachingService = ClassLoadersCachingBuildService.registerIfAbsent(project)
+        val classpath = project.configurations.named(BUILD_TOOLS_API_CLASSPATH_CONFIGURATION_NAME)
+        project.dependencies.registerTransform(BuildToolsApiClasspathEntrySnapshotTransform::class.java) {
+            it.from.attribute(ARTIFACT_TYPE_ATTRIBUTE, JAR_ARTIFACT_TYPE)
+            it.to.attribute(ARTIFACT_TYPE_ATTRIBUTE, CLASSPATH_ENTRY_SNAPSHOT_ARTIFACT_TYPE)
+            it.configureCommonParameters(classLoadersCachingService, classpath, jvmToolchain)
+        }
+        project.dependencies.registerTransform(BuildToolsApiClasspathEntrySnapshotTransform::class.java) {
+            it.from.attribute(ARTIFACT_TYPE_ATTRIBUTE, DIRECTORY_ARTIFACT_TYPE)
+            it.to.attribute(ARTIFACT_TYPE_ATTRIBUTE, CLASSPATH_ENTRY_SNAPSHOT_ARTIFACT_TYPE)
+            it.configureCommonParameters(classLoadersCachingService, classpath, jvmToolchain)
+        }
+    }
+
 }