[BT] Implement classpath snapshotting in the Build Tools API

#KT-57565 In Progress

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)
+        }
+    }
+
 }