From cfadb0326300359feeaac23d2713b8401af96a0f Mon Sep 17 00:00:00 2001
From: Ting-Yuan Huang <laszio@google.com>
Date: Mon, 11 Nov 2024 17:24:00 -0800
Subject: [PATCH] Introduce KspAATask.commandLineArgumentProviders

so that each KSP task can configured differently.
---
 .../com/google/devtools/ksp/gradle/KspAATask.kt   | 15 ++++++++++++---
 .../devtools/ksp/gradle/GradleCompilationTest.kt  | 12 +++++++++++-
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
index 2cec471e17..1b68e92cce 100644
--- a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
+++ b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
@@ -24,11 +24,13 @@ import org.gradle.api.artifacts.Configuration
 import org.gradle.api.file.ConfigurableFileCollection
 import org.gradle.api.file.DirectoryProperty
 import org.gradle.api.logging.LogLevel
+import org.gradle.api.provider.ListProperty
 import org.gradle.api.provider.MapProperty
 import org.gradle.api.provider.Property
 import org.gradle.api.provider.SetProperty
 import org.gradle.api.tasks.*
 import org.gradle.api.tasks.Optional
+import org.gradle.process.CommandLineArgumentProvider
 import org.gradle.work.ChangeType
 import org.gradle.work.Incremental
 import org.gradle.work.InputChanges
@@ -63,6 +65,9 @@ abstract class KspAATask @Inject constructor(
     @get:Nested
     abstract val kspConfig: KspGradleConfig
 
+    @get:Nested
+    abstract val commandLineArgumentProviders: ListProperty<CommandLineArgumentProvider>
+
     @TaskAction
     fun execute(inputChanges: InputChanges) {
         // FIXME: Create a class loader with clean classpath instead of shadowing existing ones. It'll require either:
@@ -220,8 +225,9 @@ abstract class KspAATask @Inject constructor(
                         )
                     )
                     cfg.processorOptions.putAll(kspExtension.apOptions)
-                    cfg.processorOptions.putAll(
-                        kspExtension.commandLineArgumentProviders.map { providers ->
+
+                    fun ListProperty<CommandLineArgumentProvider>.mapArgProviders() =
+                        map { providers ->
                             buildMap {
                                 for (provider in providers) {
                                     provider.asArguments().forEach { argument ->
@@ -234,7 +240,10 @@ abstract class KspAATask @Inject constructor(
                                 }
                             }
                         }
-                    )
+
+                    cfg.processorOptions.putAll(kspExtension.commandLineArgumentProviders.mapArgProviders())
+                    cfg.processorOptions.putAll(kspAATask.commandLineArgumentProviders.mapArgProviders())
+
                     val logLevel = LogLevel.entries.first {
                         project.logger.isEnabled(it)
                     }
diff --git a/gradle-plugin/src/test/kotlin/com/google/devtools/ksp/gradle/GradleCompilationTest.kt b/gradle-plugin/src/test/kotlin/com/google/devtools/ksp/gradle/GradleCompilationTest.kt
index 22f17be1c8..8a766215f3 100644
--- a/gradle-plugin/src/test/kotlin/com/google/devtools/ksp/gradle/GradleCompilationTest.kt
+++ b/gradle-plugin/src/test/kotlin/com/google/devtools/ksp/gradle/GradleCompilationTest.kt
@@ -329,7 +329,6 @@ class GradleCompilationTest(val useKSP2: Boolean) {
     @Test
     fun commandLineArgumentIsIncludedInApoptionsWhenAddedInKspTask() {
         Assume.assumeFalse(System.getProperty("os.name").startsWith("Windows", ignoreCase = true))
-        Assume.assumeFalse(useKSP2)
         testRule.setupAppAsAndroidApp()
         testRule.appModule.dependencies.addAll(
             listOf(
@@ -361,6 +360,17 @@ class GradleCompilationTest(val useKSP2: Boolean) {
                        println("commandLine=${'$'}{commandLine.asArguments()}")
                      }
                    }
+                   tasks.withType<com.google.devtools.ksp.gradle.KspAATask>().configureEach {
+                     val destination = project.layout.projectDirectory.dir("schemas-${'$'}{this.name}")
+                     commandLineArgumentProviders.add(Provider(destination.asFile))
+
+                     kspConfig.processorOptions.get().forEach { (key, value) ->
+                         println("apoption=${'$'}key=${'$'}value")
+                     }
+                     commandLineArgumentProviders.get().forEach { commandLine ->
+                       println("commandLine=${'$'}{commandLine.asArguments()}")
+                     }
+                   }
                  }
             """.trimIndent()
         )