From 76b581999a81a4a6e91d5f5fc6b51639095575b2 Mon Sep 17 00:00:00 2001
From: Daniel Santiago <danysantiago@google.com>
Date: Fri, 20 Dec 2024 05:31:12 -0800
Subject: [PATCH] Update Hilt Gradle Plugin KSP configuration to support KSP2

The plugin does per-variant / per-platform configuration by adding a CommandLineArgumentProvider to the KSP task, with KSP2 the task class is different and does not implement the KSP1 base class. This CL updates the plugin integration to find and also configure the KSP2 task.

Fixes https://github.com/google/dagger/issues/4303

RELNOTES=Upgrade Hilt Gradle Plugin to support KSP2 configuration.
PiperOrigin-RevId: 708292642
---
 .../dagger/hilt/android/plugin/util/Tasks.kt  | 47 ++++++++++++++-----
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/java/dagger/hilt/android/plugin/main/src/main/kotlin/dagger/hilt/android/plugin/util/Tasks.kt b/java/dagger/hilt/android/plugin/main/src/main/kotlin/dagger/hilt/android/plugin/util/Tasks.kt
index 388c90418cc..0489de7a661 100644
--- a/java/dagger/hilt/android/plugin/main/src/main/kotlin/dagger/hilt/android/plugin/util/Tasks.kt
+++ b/java/dagger/hilt/android/plugin/main/src/main/kotlin/dagger/hilt/android/plugin/util/Tasks.kt
@@ -17,7 +17,9 @@
 package dagger.hilt.android.plugin.util
 
 import com.android.build.api.variant.ComponentIdentity
-import com.google.devtools.ksp.gradle.KspTaskJvm
+import com.google.devtools.ksp.gradle.KspAATask
+import com.google.devtools.ksp.gradle.KspTask
+import kotlin.reflect.KClass
 import org.gradle.api.Project
 import org.gradle.api.Task
 import org.gradle.api.tasks.compile.JavaCompile
@@ -73,7 +75,7 @@ internal fun addKspTaskProcessorOptions(
   variantIdentity: ComponentIdentity,
   produceArgProvider: (Task) -> CommandLineArgumentProvider
 ) = project.plugins.withId("com.google.devtools.ksp") {
-  checkClass("com.google.devtools.ksp.gradle.KspTaskJvm") {
+  check(kspOneTaskClass != null || kspTwoTaskClass != null) {
     """
     The KSP plugin was detected to be applied but its task class could not be found.
 
@@ -85,13 +87,25 @@ internal fun addKspTaskProcessorOptions(
     See https://github.com/google/dagger/issues/3965 for more details.
     """.trimIndent()
   }
-  project.tasks.withType(KspTaskJvm::class.java).configureEach { task ->
-    if (task.name == "ksp${variantIdentity.name.capitalize()}Kotlin" ||
+  fun <T : Task> configureEach(
+    kclass: KClass<T>,
+    block: T.(CommandLineArgumentProvider) -> Unit
+  ) {
+    project.tasks.withType(kclass.java).configureEach { task ->
+      if (task.name == "ksp${variantIdentity.name.capitalize()}Kotlin" ||
         // Task names in shared/src/AndroidMain in KMP projects has a platform suffix.
         task.name == "ksp${variantIdentity.name.capitalize()}KotlinAndroid") {
-      task.commandLineArgumentProviders.add(produceArgProvider.invoke(task))
+        val argProvider = produceArgProvider.invoke(task)
+        task.block(argProvider)
+      }
     }
   }
+  if (kspOneTaskClass != null) {
+    configureEach(KspTask::class) { commandLineArgumentProviders.add(it) }
+  }
+  if (kspTwoTaskClass != null) {
+    configureEach(KspAATask::class) { commandLineArgumentProviders.add(it) }
+  }
 }
 
 private inline fun checkClass(fqn: String, msg: () -> String) {
@@ -102,10 +116,21 @@ private inline fun checkClass(fqn: String, msg: () -> String) {
   }
 }
 
-internal fun Task.isKspTask(): Boolean = try {
-  val kspTaskClass = Class.forName("com.google.devtools.ksp.gradle.KspTask")
-  kspTaskClass.isAssignableFrom(this::class.java)
-} catch (ex: ClassNotFoundException) {
-  false
-}
+private val kspOneTaskClass =
+  try {
+    Class.forName("com.google.devtools.ksp.gradle.KspTask")
+  } catch (ex: ClassNotFoundException) {
+    null
+  }
+
+private val kspTwoTaskClass =
+  try {
+    Class.forName("com.google.devtools.ksp.gradle.KspAATask")
+  } catch (ex: ClassNotFoundException) {
+    null
+  }
+
+internal fun Task.isKspTask() =
+  kspOneTaskClass?.isAssignableFrom(this::class.java) == true ||
+    kspTwoTaskClass?.isAssignableFrom(this::class.java) == true