(WIP) Support Dagger KSP

compiler-utils/src/main/java/com/squareup/anvil/compiler/internal/reference/AnnotationReference.kt

@@ -172,7 +172,7 @@ public sealed class AnnotationReference {
   public class Psi internal constructor(
     public val annotation: KtAnnotationEntry,
     override val classReference: ClassReference,
-    override val declaringClass: ClassReference.Psi?,
+    override val declaringClass: ClassReference?,
   ) : AnnotationReference() {
 
     override val arguments: List<AnnotationArgumentReference.Psi> by lazy(NONE) {
@@ -198,7 +198,7 @@ public sealed class AnnotationReference {
   public class Descriptor internal constructor(
     public val annotation: AnnotationDescriptor,
     override val classReference: ClassReference,
-    override val declaringClass: ClassReference.Descriptor?,
+    override val declaringClass: ClassReference?,
   ) : AnnotationReference() {
 
     override val arguments: List<AnnotationArgumentReference.Descriptor> by lazy(NONE) {
@@ -221,10 +221,21 @@ public sealed class AnnotationReference {
 public fun KtAnnotationEntry.toAnnotationReference(
   declaringClass: ClassReference.Psi?,
   module: ModuleDescriptor,
+): Psi {
+  return toAnnotationReference(
+    classReference = requireFqName(module).toClassReference(module),
+    declaringClass = declaringClass,
+  )
+}
+
+@ExperimentalAnvilApi
+public fun KtAnnotationEntry.toAnnotationReference(
+  declaringClass: ClassReference?,
+  classReference: ClassReference,
 ): Psi {
   return Psi(
     annotation = this,
-    classReference = requireFqName(module).toClassReference(module),
+    classReference = classReference,
     declaringClass = declaringClass,
   )
 }
@@ -237,10 +248,17 @@ public fun AnnotationDescriptor.toAnnotationReference(
   val annotationClass = annotationClass ?: throw AnvilCompilationException(
     message = "Couldn't find the annotation class for $fqName",
   )
+  return toAnnotationReference(declaringClass, annotationClass.toClassReference(module))
+}
 
+@ExperimentalAnvilApi
+public fun AnnotationDescriptor.toAnnotationReference(
+  declaringClass: ClassReference?,
+  classReference: ClassReference,
+): Descriptor {
   return Descriptor(
     annotation = this,
-    classReference = annotationClass.toClassReference(module),
+    classReference = classReference,
     declaringClass = declaringClass,
   )
 }

compiler-utils/src/testFixtures/java/com/squareup/anvil/compiler/internal/testing/AnvilCompilation.kt

@@ -16,6 +16,7 @@ import com.tschuchort.compiletesting.kspArgs
 import com.tschuchort.compiletesting.kspWithCompilation
 import com.tschuchort.compiletesting.symbolProcessorProviders
 import dagger.internal.codegen.ComponentProcessor
+import dagger.internal.codegen.KspComponentProcessor
 import org.intellij.lang.annotations.Language
 import org.jetbrains.kotlin.config.JvmTarget
 import java.io.File
@@ -39,7 +40,7 @@ public class AnvilCompilation internal constructor(
   @Suppress("SuspiciousCollectionReassignment")
   @ExperimentalAnvilApi
   public fun configureAnvil(
-    enableDaggerAnnotationProcessor: Boolean = false,
+    daggerAnnotationProcessingMode: DaggerAnnotationProcessingMode = DaggerAnnotationProcessingMode.NONE,
     generateDaggerFactories: Boolean = false,
     generateDaggerFactoriesOnly: Boolean = false,
     disableComponentMerging: Boolean = false,
@@ -59,8 +60,25 @@ public class AnvilCompilation internal constructor(
       // Deprecation tracked in https://github.com/square/anvil/issues/672
       @Suppress("DEPRECATION")
       componentRegistrars = listOf(anvilComponentRegistrar)
-      if (enableDaggerAnnotationProcessor) {
-        annotationProcessors = listOf(ComponentProcessor(), AutoAnnotationProcessor())
+
+      when (daggerAnnotationProcessingMode) {
+        DaggerAnnotationProcessingMode.KAPT -> {
+          annotationProcessors = listOf(ComponentProcessor(), AutoAnnotationProcessor())
+        }
+        DaggerAnnotationProcessingMode.KSP -> {
+          @Suppress("invisible_reference", "invisible_member")
+          symbolProcessorProviders +=
+            listOf(
+              com.squareup.anvil.compiler.ksp.MergeComponentSymbolProcessor.Provider(),
+              KspComponentProcessor.Provider(),
+            )
+          // Run KSP in a single-pass
+          // https://kotlinlang.slack.com/archives/C013BA8EQSE/p1639462548225400?thread_ts=1639433474.224900&cid=C013BA8EQSE
+          kspWithCompilation = true
+        }
+        DaggerAnnotationProcessingMode.NONE -> {
+          // Do nothing
+        }
       }
 
       val anvilCommandLineProcessor = AnvilCommandLineProcessor()
@@ -104,10 +122,13 @@ public class AnvilCompilation internal constructor(
         is Ksp -> {
           symbolProcessorProviders += buildList {
             addAll(
-              ServiceLoader.load(
-                SymbolProcessorProvider::class.java,
-                SymbolProcessorProvider::class.java.classLoader,
-              ),
+              ServiceLoader
+                .load(
+                  SymbolProcessorProvider::class.java,
+                  SymbolProcessorProvider::class.java.classLoader,
+                )
+                // Exclude the Dagger KSP processor, we have special handling for that as we decorate it
+                .filterNot { it is KspComponentProcessor.Provider },
             )
             addAll(mode.symbolProcessorProviders)
           }
@@ -219,6 +240,13 @@ public class AnvilCompilation internal constructor(
   }
 }
 
+/** Available Dagger annotation processing modes. */
+enum class DaggerAnnotationProcessingMode {
+  KAPT,
+  KSP,
+  NONE,
+}
+
 /**
  * Helpful for testing code generators in unit tests end to end.
  *
@@ -229,7 +257,7 @@ public class AnvilCompilation internal constructor(
 @ExperimentalAnvilApi
 public fun compileAnvil(
   @Language("kotlin") vararg sources: String,
-  enableDaggerAnnotationProcessor: Boolean = false,
+  daggerAnnotationProcessingMode: DaggerAnnotationProcessingMode = DaggerAnnotationProcessingMode.NONE,
   generateDaggerFactories: Boolean = false,
   generateDaggerFactoriesOnly: Boolean = false,
   disableComponentMerging: Boolean = false,
@@ -243,7 +271,8 @@ public fun compileAnvil(
   jvmTarget: JvmTarget? = null,
   block: JvmCompilationResult.() -> Unit = { },
 ): JvmCompilationResult {
-  return AnvilCompilation()
+  val compilation = AnvilCompilation()
+  return compilation
     .apply {
       kotlinCompilation.apply {
         this.allWarningsAsErrors = allWarningsAsErrors
@@ -265,13 +294,15 @@ public fun compileAnvil(
       }
     }
     .configureAnvil(
-      enableDaggerAnnotationProcessor = enableDaggerAnnotationProcessor,
+      daggerAnnotationProcessingMode = daggerAnnotationProcessingMode,
       generateDaggerFactories = generateDaggerFactories,
       generateDaggerFactoriesOnly = generateDaggerFactoriesOnly,
       disableComponentMerging = disableComponentMerging,
       enableExperimentalAnvilApis = enableExperimentalAnvilApis,
       mode = mode,
     )
     .compile(*sources)
-    .also(block)
+    .also { result ->
+      result.block()
+    }
 }

compiler/build.gradle.kts

@@ -39,6 +39,7 @@ dependencies {
   implementation(project(":compiler-api"))
   implementation(project(":compiler-utils"))
   implementation(libs.dagger2)
+  implementation(libs.dagger2.compiler)
   implementation(libs.jsr250)
   implementation(libs.kotlinpoet)
   implementation(libs.kotlinpoet.ksp)
@@ -47,16 +48,21 @@ dependencies {
   compileOnly(libs.kotlin.compiler)
   compileOnly(libs.ksp.compilerPlugin)
   compileOnly(libs.ksp.api)
+  // just for reference to get underlying APIs
+  compileOnly(libs.ksp.compilerPlugin)
 
   kapt(libs.auto.service.processor)
 
   testImplementation(testFixtures(project(":compiler-utils")))
-  testImplementation(libs.dagger2.compiler)
   testImplementation(libs.kotlin.annotationProcessingEmbeddable)
   testImplementation(libs.kotlin.compileTesting)
   testImplementation(libs.kotlin.compileTesting.ksp)
   testImplementation(libs.ksp.compilerPlugin)
   testImplementation(libs.kotlin.compiler)
   testImplementation(libs.kotlin.test)
+  testImplementation(libs.ksp.api)
+  testImplementation(libs.ksp.compilerPlugin)
   testImplementation(libs.truth)
+  // Force latest guava version for alignment with dagger's versions
+  testImplementation(libs.guava)
 }