Move KSP to an independent Gradle task

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/KotlinSymbolProcessingExtension.kt

@@ -59,8 +59,6 @@ class KotlinSymbolProcessingExtension(
 
 abstract class AbstractKotlinSymbolProcessingExtension(val options: KspOptions, val logger: KSPLogger, val testMode: Boolean) :
     AnalysisHandlerExtension {
-    private var completed = false
-
     override fun doAnalysis(
         project: Project,
         module: ModuleDescriptor,
@@ -69,9 +67,6 @@ abstract class AbstractKotlinSymbolProcessingExtension(val options: KspOptions,
         bindingTrace: BindingTrace,
         componentProvider: ComponentProvider
     ): AnalysisResult? {
-        if (completed)
-            return null
-
         val psiManager = PsiManager.getInstance(project)
         val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
         val javaFiles = options.javaSourceRoots
@@ -101,7 +96,7 @@ abstract class AbstractKotlinSymbolProcessingExtension(val options: KspOptions,
 
         KSObjectCacheManager.clear()
 
-        return AnalysisResult.success(BindingContext.EMPTY, module, shouldGenerateCode = false)
+        return AnalysisResult.EMPTY
     }
 
     abstract fun loadProcessors(): List<SymbolProcessor>
@@ -112,21 +107,7 @@ abstract class AbstractKotlinSymbolProcessingExtension(val options: KspOptions,
         bindingTrace: BindingTrace,
         files: Collection<KtFile>
     ): AnalysisResult? {
-        if (completed)
-            return null
-
-        completed = true
-
-        if (testMode)
-            return null
-
-        return AnalysisResult.RetryWithAdditionalRoots(
-            BindingContext.EMPTY,
-            module,
-            listOf(options.javaOutputDir),
-            listOf(options.kotlinOutputDir),
-            addToEnvironment = true
-        )
+        return AnalysisResult.success(BindingContext.EMPTY, module, shouldGenerateCode = false)
     }
 
     private var annotationProcessingComplete = false

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/KotlinSymbolProcessingPlugin.kt

@@ -65,17 +65,6 @@ class KotlinSymbolProcessingCommandLineProcessor : CommandLineProcessor {
 
 class KotlinSymbolProcessingComponentRegistrar : ComponentRegistrar {
     override fun registerProjectComponents(project: MockProject, configuration: CompilerConfiguration) {
-        // Don't bother with KAPT tasks.
-        // There is no way to pass KSP options to compileKotlin only. Have to workaround here.
-        val outputDir = configuration[JVMConfigurationKeys.OUTPUT_DIRECTORY]
-        val kaptOutputDirs = listOf(
-            listOf("tmp", "kapt3", "stubs"),
-            listOf("tmp", "kapt3", "incrementalData"),
-            listOf("tmp", "kapt3", "incApCache")
-        ).map { File(it.joinToString(File.separator)) }
-        if (kaptOutputDirs.any { outputDir?.parentFile?.endsWith(it) == true })
-            return
-
         val contentRoots = configuration[CLIConfigurationKeys.CONTENT_ROOTS] ?: emptyList()
         val options = configuration[KSP_OPTIONS]?.apply {
             javaSourceRoots.addAll(contentRoots.filterIsInstance<JavaSourceRoot>().map { it.file })

gradle-plugin/build.gradle.kts

@@ -16,6 +16,8 @@ plugins {
 
 dependencies {
     implementation("org.jetbrains.kotlin:kotlin-gradle-plugin-api:$kotlinBaseVersion")
+    implementation("org.jetbrains.kotlin:kotlin-gradle-plugin:$kotlinBaseVersion")
+    compileOnly("org.jetbrains.kotlin:kotlin-compiler:$kotlinBaseVersion")
 
     compileOnly(gradleApi())
 

gradle-plugin/src/main/java/com/google/devtools/ksp/gradle/InternalTrampoline.java

@@ -0,0 +1,11 @@
+package com.google.devtools.ksp.gradle;
+
+import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation;
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompileTaskData;
+
+// FIXME: Ask upstream to open these API
+public class InternalTrampoline {
+    public static void KotlinCompileTaskData_register(String taskName, KotlinCompilation<?> kotlinCompilation) {
+        KotlinCompileTaskData kotlinCompileTaskData = KotlinCompileTaskData.Companion.register(taskName, kotlinCompilation);
+    }
+}

gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspSubplugin.kt

@@ -32,6 +32,7 @@ import com.google.devtools.ksp.gradle.KspGradleSubplugin.Companion.getKspKotlinO
 import com.google.devtools.ksp.gradle.KspGradleSubplugin.Companion.getKspResourceOutputDir
 import java.io.File
 import javax.inject.Inject
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 
 class KspGradleSubplugin @Inject internal constructor(private val registry: ToolingModelBuilderRegistry) : Plugin<Project> {
     companion object {
@@ -67,6 +68,8 @@ class KspGradleSubplugin @Inject internal constructor(private val registry: Tool
 class KspKotlinGradleSubplugin : KotlinGradleSubplugin<AbstractCompile> {
     companion object {
         const val KSP_ARTIFACT_NAME = "symbol-processing"
+        const val KSP_PLUGIN_ID = "com.google.devtools.ksp.symbol-processing"
+
     }
 
     override fun isApplicable(project: Project, task: AbstractCompile) = KspGradleSubplugin.isEnabled(project)
@@ -88,8 +91,6 @@ class KspKotlinGradleSubplugin : KotlinGradleSubplugin<AbstractCompile> {
 
         kotlinCompile.dependsOn(kspConfiguration.buildDependencies)
 
-        kotlinCompile.setProperty("incremental", false)
-
         val options = mutableListOf<SubpluginOption>()
 
         options += FilesSubpluginOption("apclasspath", kspConfiguration)
@@ -114,10 +115,47 @@ class KspKotlinGradleSubplugin : KotlinGradleSubplugin<AbstractCompile> {
             javaCompile.source(generatedJavaSources)
         }
 
-        return options
+        val kspTaskName = kotlinCompile.name.replaceFirst("compile", "ksp")
+        InternalTrampoline.KotlinCompileTaskData_register(kspTaskName, kotlinCompilation)
+
+        val kspTaskProvider = project.tasks.register(kspTaskName, KspTask::class.java) { kspTask ->
+            kspTask.setDestinationDir(File(project.buildDir, "generated/ksp"))
+            kspTask.mapClasspath { kotlinCompile.classpath }
+            kspTask.options = options
+        }.apply {
+            configure {
+                kotlinCompilation?.allKotlinSourceSets?.forEach { sourceSet -> it.source(sourceSet.kotlin) }
+            }
+        }
+
+        kotlinCompile.dependsOn(kspTaskProvider)
+        kotlinCompile.source(kotlinOutputDir, javaOutputDir)
+
+        return emptyList()
     }
 
-    override fun getCompilerPluginId() = "com.google.devtools.ksp.symbol-processing"
+    override fun getCompilerPluginId() = KSP_PLUGIN_ID
     override fun getPluginArtifact(): SubpluginArtifact =
         SubpluginArtifact(groupId = "com.google.devtools.ksp", artifactId = KSP_ARTIFACT_NAME, version = javaClass.`package`.implementationVersion)
 }
+
+open class KspTask : KotlinCompile() {
+    lateinit var options: List<SubpluginOption>
+
+    init {
+        // kotlinc's incremental compilation isn't compatible with symbol processing in a few ways:
+        // * It doesn't consider private / internal changes when computing dirty sets.
+        // * It compiles iteratively; Sources can be compiled in different rounds.
+        incremental = false
+    }
+
+    override fun setupCompilerArgs(
+        args: org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments,
+        defaultsOnly: Boolean,
+        ignoreClasspathResolutionErrors: Boolean
+    ) {
+        fun SubpluginOption.toArg() = "plugin:${KspKotlinGradleSubplugin.KSP_PLUGIN_ID}:${key}=${value}"
+        super.setupCompilerArgs(args, defaultsOnly, ignoreClasspathResolutionErrors)
+        args.pluginOptions = (options.map { it.toArg() } + args.pluginOptions!!).toTypedArray()
+    }
+}