Update ComponentDetectorCheck for ksp (#753)

compiler/src/main/java/com/squareup/anvil/compiler/codegen/dagger/ComponentDetectorCheck.kt

@@ -1,38 +1,72 @@
 package com.squareup.anvil.compiler.codegen.dagger
 
 import com.google.auto.service.AutoService
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
+import com.google.devtools.ksp.processing.SymbolProcessorProvider
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.squareup.anvil.compiler.api.AnvilApplicabilityChecker
 import com.squareup.anvil.compiler.api.AnvilContext
 import com.squareup.anvil.compiler.api.CodeGenerator
 import com.squareup.anvil.compiler.codegen.PrivateCodeGenerator
+import com.squareup.anvil.compiler.codegen.ksp.AnvilSymbolProcessor
+import com.squareup.anvil.compiler.codegen.ksp.AnvilSymbolProcessorProvider
+import com.squareup.anvil.compiler.codegen.ksp.KspAnvilException
 import com.squareup.anvil.compiler.daggerComponentFqName
 import com.squareup.anvil.compiler.internal.reference.AnvilCompilationExceptionClassReference
 import com.squareup.anvil.compiler.internal.reference.classAndInnerClassReferences
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
-@AutoService(CodeGenerator::class)
-internal class ComponentDetectorCheck : PrivateCodeGenerator() {
-
+internal object ComponentDetectorCheck : AnvilApplicabilityChecker {
+  private const val MESSAGE =
+    "Anvil cannot generate the code for Dagger components or subcomponents. In " +
+      "these cases the Dagger annotation processor is required. Enabling the Dagger " +
+      "annotation processor and turning on Anvil to generate Dagger factories is " +
+      "redundant. Set 'generateDaggerFactories' to false."
+  private val ANNOTATIONS_TO_CHECK = setOf(daggerComponentFqName)
   override fun isApplicable(context: AnvilContext) = context.generateFactories
 
-  override fun generateCodePrivate(
-    codeGenDir: File,
-    module: ModuleDescriptor,
-    projectFiles: Collection<KtFile>
-  ) {
-    val component = projectFiles
-      .classAndInnerClassReferences(module)
-      .firstOrNull { it.isAnnotatedWith(daggerComponentFqName) }
-
-    if (component != null) {
-      throw AnvilCompilationExceptionClassReference(
-        message = "Anvil cannot generate the code for Dagger components or subcomponents. In " +
-          "these cases the Dagger annotation processor is required. Enabling the Dagger " +
-          "annotation processor and turning on Anvil to generate Dagger factories is " +
-          "redundant. Set 'generateDaggerFactories' to false.",
-        classReference = component
-      )
+  internal class KspGenerator(
+    override val env: SymbolProcessorEnvironment,
+  ) : AnvilSymbolProcessor() {
+    @AutoService(SymbolProcessorProvider::class)
+    class Provider : AnvilSymbolProcessorProvider(ComponentDetectorCheck, ::KspGenerator)
+
+    override fun processChecked(resolver: Resolver): List<KSAnnotated> {
+      val clazz =
+        ANNOTATIONS_TO_CHECK.flatMap { resolver.getSymbolsWithAnnotation(it.asString()) }
+          .firstOrNull()
+          ?: return emptyList()
+
+      throw KspAnvilException(message = MESSAGE, node = clazz)
+    }
+  }
+
+  @AutoService(CodeGenerator::class)
+  internal class EmbeddedGenerator : PrivateCodeGenerator() {
+
+    override fun isApplicable(context: AnvilContext) =
+      ComponentDetectorCheck.isApplicable(context)
+
+    override fun generateCodePrivate(
+      codeGenDir: File,
+      module: ModuleDescriptor,
+      projectFiles: Collection<KtFile>
+    ) {
+      val clazz =
+        projectFiles.classAndInnerClassReferences(module).firstOrNull {
+          clazz ->
+          ANNOTATIONS_TO_CHECK.any { clazz.isAnnotatedWith(it) }
+        }
+
+      if (clazz != null) {
+        throw AnvilCompilationExceptionClassReference(
+          message = MESSAGE,
+          classReference = clazz
+        )
+      }
     }
   }
 }

compiler/src/test/java/com/squareup/anvil/compiler/dagger/ComponentDetectorCheckTest.kt

@@ -2,14 +2,28 @@ package com.squareup.anvil.compiler.dagger
 
 import com.google.common.truth.Truth.assertThat
 import com.squareup.anvil.compiler.WARNINGS_AS_ERRORS
+import com.squareup.anvil.compiler.api.CodeGenerator
+import com.squareup.anvil.compiler.internal.testing.AnvilCompilationMode
 import com.squareup.anvil.compiler.internal.testing.compileAnvil
 import com.squareup.anvil.compiler.isError
 import com.tschuchort.compiletesting.JvmCompilationResult
 import com.tschuchort.compiletesting.KotlinCompilation.ExitCode.OK
 import org.intellij.lang.annotations.Language
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
 
-class ComponentDetectorCheckTest {
+@RunWith(Parameterized::class)
+class ComponentDetectorCheckTest(
+  private val mode: AnvilCompilationMode
+) {
+
+  companion object {
+    @Parameterized.Parameters(name = "{0}")
+    @JvmStatic fun modes(): Collection<Any> {
+      return listOf(AnvilCompilationMode.Embedded(), AnvilCompilationMode.Ksp())
+    }
+  }
 
   @Test fun `a Dagger component causes an error`() {
     compile(
@@ -20,11 +34,12 @@ class ComponentDetectorCheckTest {
       
       @Component
       interface ComponentInterface
-      """
+      """,
+      mode = mode
     ) {
       assertThat(exitCode).isError()
-      // Position to the class.
-      assertThat(messages).contains("Source0.kt:6:11")
+      // Position to the class. )
+      assertThat(messages).contains("Source0.kt:6")
       assertThat(messages).contains(
         "Anvil cannot generate the code for Dagger components or subcomponents. In these " +
           "cases the Dagger annotation processor is required. Enabling the Dagger " +
@@ -43,7 +58,8 @@ class ComponentDetectorCheckTest {
       
       @Subcomponent
       interface ComponentInterface
-      """
+      """,
+      mode = mode
     ) {
       assertThat(exitCode).isEqualTo(OK)
     }
@@ -60,11 +76,12 @@ class ComponentDetectorCheckTest {
           @Component
           interface ComponentInterface
         }
-        """
+        """,
+      mode = mode
     ) {
       assertThat(exitCode).isError()
       // Position to the class.
-      assertThat(messages).contains("Source0.kt:7:13")
+      assertThat(messages).contains("Source0.kt:7")
       assertThat(messages).contains(
         "Anvil cannot generate the code for Dagger components or subcomponents. In these " +
           "cases the Dagger annotation processor is required. Enabling the Dagger " +
@@ -85,19 +102,23 @@ class ComponentDetectorCheckTest {
         @Subcomponent
         interface ComponentInterface
       }
-      """
+      """,
+      mode = mode
     ) {
       assertThat(exitCode).isEqualTo(OK)
     }
   }
 
   private fun compile(
     @Language("kotlin") vararg sources: String,
+    codeGenerators: List<CodeGenerator> = emptyList(),
+    mode: AnvilCompilationMode = AnvilCompilationMode.Embedded(codeGenerators),
     block: JvmCompilationResult.() -> Unit = { }
   ): JvmCompilationResult = compileAnvil(
     sources = sources,
     generateDaggerFactories = true,
     allWarningsAsErrors = WARNINGS_AS_ERRORS,
-    block = block
+    block = block,
+    mode = mode
   )
 }