Repository: square/anvil
Pull request #307 (2021-06-05 22:58:49)

Refactor the Gradle plugin and create multiple Anvil configurations.
Rework how the Gradle plugin sets up its configuration, Gradle tasks and patches for other Gradle tasks. Instead of querying the tasks manually and checking which plugins are applied, use the callback methods from the Kotlin Gradle Plugin. With this change we can remove a lot from our custom logic.

This also allows us to easily create an Anvil configuration for each build variant separately. This allows to add custom code generators to specific build types only.

This refactoring is needed for #100.


Patch:
diff --git a/gradle-plugin/src/main/java/com/squareup/anvil/plugin/AnvilPlugin.kt b/gradle-plugin/src/main/java/com/squareup/anvil/plugin/AnvilPlugin.kt
index a701a1c9..cac47a26 100644
--- a/gradle-plugin/src/main/java/com/squareup/anvil/plugin/AnvilPlugin.kt
+++ b/gradle-plugin/src/main/java/com/squareup/anvil/plugin/AnvilPlugin.kt
@@ -1,25 +1,22 @@
 package com.squareup.anvil.plugin
 
 import com.android.build.gradle.AppExtension
-import com.android.build.gradle.AppPlugin
 import com.android.build.gradle.LibraryExtension
-import com.android.build.gradle.LibraryPlugin
 import com.android.build.gradle.TestExtension
 import com.android.build.gradle.TestedExtension
 import com.android.build.gradle.api.BaseVariant
-import com.squareup.anvil.plugin.ProjectType.ANDROID
-import com.squareup.anvil.plugin.ProjectType.JVM
-import com.squareup.anvil.plugin.ProjectType.MULTIPLATFORM_ANDROID
-import com.squareup.anvil.plugin.ProjectType.MULTIPLATFORM_JVM
+import com.android.build.gradle.api.TestVariant
+import com.android.build.gradle.api.UnitTestVariant
+import com.android.build.gradle.internal.tasks.factory.dependsOn
 import org.gradle.api.Action
 import org.gradle.api.GradleException
 import org.gradle.api.Project
 import org.gradle.api.Task
+import org.gradle.api.UnknownTaskException
 import org.gradle.api.artifacts.Configuration
-import org.gradle.api.plugins.AppliedPlugin
-import org.gradle.api.plugins.PluginManager
 import org.gradle.api.provider.Provider
-import org.jetbrains.kotlin.gradle.dsl.KotlinMultiplatformExtension
+import org.gradle.api.tasks.TaskContainer
+import org.gradle.api.tasks.TaskProvider
 import org.jetbrains.kotlin.gradle.internal.KaptGenerateStubsTask
 import org.jetbrains.kotlin.gradle.plugin.FilesSubpluginOption
 import org.jetbrains.kotlin.gradle.plugin.KaptExtension
@@ -27,34 +24,92 @@ import org.jetbrains.kotlin.gradle.plugin.KotlinCompilation
 import org.jetbrains.kotlin.gradle.plugin.KotlinCompilerPluginSupportPlugin
 import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.androidJvm
 import org.jetbrains.kotlin.gradle.plugin.KotlinPlatformType.jvm
-import org.jetbrains.kotlin.gradle.plugin.KotlinPluginWrapper
 import org.jetbrains.kotlin.gradle.plugin.PLUGIN_CLASSPATH_CONFIGURATION_NAME
 import org.jetbrains.kotlin.gradle.plugin.SubpluginArtifact
 import org.jetbrains.kotlin.gradle.plugin.SubpluginOption
+import org.jetbrains.kotlin.gradle.plugin.mpp.KotlinJvmAndroidCompilation
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 import java.io.File
 import java.util.Locale.US
-import java.util.concurrent.atomic.AtomicBoolean
 
 @Suppress("unused")
 open class AnvilPlugin : KotlinCompilerPluginSupportPlugin {
 
-  override fun isApplicable(kotlinCompilation: KotlinCompilation<*>): Boolean {
-    return kotlinCompilation.target.project.plugins.hasPlugin(AnvilPlugin::class.java)
-  }
+  override fun apply(target: Project) {
+    val extension = target.extensions.create("anvil", AnvilExtension::class.java)
 
-  override fun getCompilerPluginId(): String = "com.squareup.anvil.compiler"
+    // Create a configuration for collecting CodeGenerator dependencies. We need to create all
+    // configurations eagerly and cannot wait for applyToCompilation(..) below, because this
+    // function is called in an afterEvaluate block by the Kotlin Gradle Plugin. That's too late
+    // to register the configurations.
+    val commonConfiguration = getConfiguration(target, buildType = "")
+
+    // anvilTest is the common test configuration similar to other configurations like kaptTest.
+    // Android build type specific unit test variants will extend this variant below. For the JVM
+    // this single variant is enough.
+    val testConfiguration = getConfiguration(target, buildType = "test").apply {
+      extendsFrom(commonConfiguration)
+    }
 
-  override fun getPluginArtifact(): SubpluginArtifact = SubpluginArtifact(
-    groupId = GROUP,
-    artifactId = "compiler",
-    version = VERSION
-  )
+    agpPlugins.forEach { agpPlugin ->
+      target.pluginManager.withPlugin(agpPlugin) {
+        // This is the common android test variant, similar to anvilTest above.
+        val androidTestVariant = getConfiguration(target, buildType = "androidTest").apply {
+          extendsFrom(commonConfiguration)
+        }
+
+        target.androidVariantsConfigure { variant ->
+          // E.g. "debugAnvil", "testReleaseAnvil", ...
+          val configuration = getConfiguration(target, buildType = variant.name)
+
+          when (variant) {
+            is UnitTestVariant -> configuration.extendsFrom(testConfiguration)
+            is TestVariant -> configuration.extendsFrom(androidTestVariant)
+          }
+        }
+      }
+    }
+
+    target.afterEvaluate {
+      if (!extension.generateDaggerFactories.get() &&
+        extension.generateDaggerFactoriesOnly.get()
+      ) {
+        throw GradleException(
+          "You cannot set generateDaggerFactories to false and generateDaggerFactoriesOnly " +
+            "to true at the same time."
+        )
+      }
+    }
+  }
+
+  override fun isApplicable(kotlinCompilation: KotlinCompilation<*>): Boolean {
+    return when (kotlinCompilation.platformType) {
+      androidJvm, jvm -> true
+      else -> false
+    }
+  }
 
   override fun applyToCompilation(
     kotlinCompilation: KotlinCompilation<*>
   ): Provider<List<SubpluginOption>> {
-    val project = kotlinCompilation.target.project
+    val variant = Variant(kotlinCompilation)
+    val project = variant.project
+
+    // Make the kotlin compiler classpath extend our configurations to pick up our extra
+    // generators.
+    project.configurations.getByName(variant.compilerPluginClasspathName)
+      .extendsFrom(getConfiguration(project, variant.name))
+
+    disableIncrementalKotlinCompilation(variant)
+    disablePreciseJavaTracking(variant)
+
+    if (!variant.extension.generateDaggerFactoriesOnly.get()) {
+      disableCorrectErrorTypes(variant)
+
+      kotlinCompilation.dependencies {
+        implementation("$GROUP:annotations:$VERSION")
+      }
+    }
 
     // Notice that we use the name of the Kotlin compilation as a directory name. Generated code
     // for this specific compile task will be included in the task output. The output of different
@@ -88,233 +143,123 @@ open class AnvilPlugin : KotlinCompilerPluginSupportPlugin {
     }
   }
 
-  override fun apply(target: Project) {
-    val extension = target.extensions.create("anvil", AnvilExtension::class.java)
-
-    // Create a configuration for collecting CodeGenerator dependencies.
-    val anvilConfiguration = target.configurations.maybeCreate("anvil").apply {
-      description = "This configuration is used for dependencies with Anvil CodeGenerator " +
-        "implementations."
-    }
+  override fun getCompilerPluginId(): String = "com.squareup.anvil.compiler"
 
-    val once = AtomicBoolean()
+  override fun getPluginArtifact(): SubpluginArtifact = SubpluginArtifact(
+    groupId = GROUP,
+    artifactId = "compiler",
+    version = VERSION
+  )
 
-    fun PluginManager.withPluginOnce(
-      id: String,
-      action: (AppliedPlugin) -> Unit
-    ) {
-      withPlugin(id) {
-        if (once.compareAndSet(false, true)) {
-          action(it)
-        }
-      }
-    }
+  private fun disablePreciseJavaTracking(variant: Variant) {
+    variant.compileTaskProvider.configure { compileTask ->
+      val result = CheckMixedSourceSet.preparePreciseJavaTrackingCheck(variant)
 
-    // Apply the Anvil plugin after the Kotlin plugin was applied. There could be a timing
-    // issue. Also make sure to apply it only once. A module could accidentally apply the JVM and
-    // Android Kotlin plugin.
-    target.pluginManager.withPluginOnce("org.jetbrains.kotlin.android") {
-      realApply(target, ANDROID, extension, anvilConfiguration)
-    }
-    target.pluginManager.withPluginOnce("org.jetbrains.kotlin.multiplatform") {
-      target.afterEvaluate {
-        // We need to run this in afterEvaluate, because otherwise the multiplatform extension
-        // isn't initialized.
-        val targets =
-          target.extensions.findByType(KotlinMultiplatformExtension::class.java)?.targets
-
-        when {
-          targets?.any { it.platformType == androidJvm } == true ->
-            realApply(target, MULTIPLATFORM_ANDROID, extension, anvilConfiguration)
-          targets?.any { it.platformType == jvm } == true ->
-            realApply(target, MULTIPLATFORM_JVM, extension, anvilConfiguration)
-          else -> throw GradleException(
-            "Multiplatform project $target is not setup for Android or JVM."
-          )
-        }
-      }
-    }
-    target.pluginManager.withPluginOnce("org.jetbrains.kotlin.jvm") {
-      realApply(target, JVM, extension, anvilConfiguration)
-    }
-
-    target.afterEvaluate {
-      if (!once.get()) {
-        throw GradleException(
-          "No supported plugins for Anvil found on project " +
-            "'${target.path}'. Only Android and Java modules are supported for now."
-        )
-      }
+      compileTask.doFirstCompat {
+        // Disable precise java tracking if needed. Note that the doFirst() action only runs
+        // if the task is not up to date. That's ideal, because if nothing needs to be
+        // compiled, then we don't need to disable the flag.
+        //
+        // We also use the doFirst block to walk through the file system at execution time
+        // and minimize the IO at configuration time.
+        CheckMixedSourceSet.disablePreciseJavaTrackingIfNeeded(compileTask, result)
 
-      if (!extension.generateDaggerFactories.get() &&
-        extension.generateDaggerFactoriesOnly.get()
-      ) {
-        throw GradleException(
-          "You cannot set generateDaggerFactories to false and generateDaggerFactoriesOnly " +
-            "to true at the same time."
+        compileTask.logger.info(
+          "Anvil: Use precise java tracking: ${compileTask.usePreciseJavaTracking}"
         )
       }
     }
   }
 
-  private fun realApply(
-    project: Project,
-    projectType: ProjectType,
-    extension: AnvilExtension,
-    anvilConfiguration: Configuration
-  ) {
-    disableIncrementalKotlinCompilation(project, projectType, extension)
-    disablePreciseJavaTracking(project)
-
-    project.afterEvaluate {
-      if (!extension.generateDaggerFactoriesOnly.get()) {
-        project.pluginManager.withPlugin("org.jetbrains.kotlin.kapt") {
-          // This needs to be disabled, otherwise compiler plugins fail in weird ways when
-          // generating stubs, e.g.:
-          //
-          // /anvil/sample/app/build/generated/source/kapt/debug/com/squareup/anvil
-          // /sample/DaggerAppComponent.java:13: error: DaggerAppComponent is not abstract and does
-          // not override abstract method string() in RandomComponent
-          // public final class DaggerAppComponent implements AppComponent {
-          //              ^
-          // 1 error
-          project.extensions.findByType(KaptExtension::class.java)?.correctErrorTypes = false
-        }
-
-        project.dependencies.add("implementation", "$GROUP:annotations:$VERSION")
-      }
+  private fun disableCorrectErrorTypes(variant: Variant) {
+    variant.project.pluginManager.withPlugin("org.jetbrains.kotlin.kapt") {
+      // This needs to be disabled, otherwise compiler plugins fail in weird ways when
+      // generating stubs, e.g.:
+      //
+      // /anvil/sample/app/build/generated/source/kapt/debug/com/squareup/anvil
+      // /sample/DaggerAppComponent.java:13: error: DaggerAppComponent is not abstract and does
+      // not override abstract method string() in RandomComponent
+      // public final class DaggerAppComponent implements AppComponent {
+      //              ^
+      // 1 error
+      variant.project.extensions.findByType(KaptExtension::class.java)?.correctErrorTypes = false
     }
-
-    // Make the kotlin compiler classpath extend our configuration to pick up our extra generators.
-    project.configurations.getByName(PLUGIN_CLASSPATH_CONFIGURATION_NAME)
-      .extendsFrom(anvilConfiguration)
   }
 
-  private fun disablePreciseJavaTracking(
-    project: Project
-  ) {
-    project.tasks
-      .withType(KotlinCompile::class.java)
-      .configureEach { compileTask ->
-        val result = CheckMixedSourceSet.preparePreciseJavaTrackingCheck(compileTask)
-
-        compileTask.doFirstCompat {
-          // Disable precise java tracking if needed. Note that the doFirst() action only runs
-          // if the task is not up to date. That's ideal, because if nothing needs to be
-          // compiled, then we don't need to disable the flag.
-          //
-          // We also use the doFirst block to walk through the file system at execution time
-          // and minimize the IO at configuration time.
-          CheckMixedSourceSet.disablePreciseJavaTrackingIfNeeded(compileTask, result)
-
-          compileTask.logger.info(
-            "Anvil: Use precise java tracking: ${compileTask.usePreciseJavaTracking}"
-          )
-        }
-      }
-  }
-
-  private fun disableIncrementalKotlinCompilation(
-    project: Project,
-    projectType: ProjectType,
-    extension: AnvilExtension
-  ) {
+  private fun disableIncrementalKotlinCompilation(variant: Variant) {
     // Use this signal to share state between DisableIncrementalCompilationTask and the Kotlin
     // compile task. If the plugin classpath changed, then DisableIncrementalCompilationTask sets
     // the signal to false.
     @Suppress("UnstableApiUsage")
-    val incrementalSignal = IncrementalSignal.registerIfAbsent(project)
+    val incrementalSignal = IncrementalSignal.registerIfAbsent(variant.project)
 
-    val taskSuffix = when (projectType) {
-      JVM, ANDROID -> "Kotlin"
-      MULTIPLATFORM_JVM -> "KotlinJvm"
-      MULTIPLATFORM_ANDROID -> "KotlinAndroid"
-    }
+    disableIncrementalCompilationWhenTheCompilerClasspathChanges(
+      variant, incrementalSignal, variant.compileTaskProvider
+    )
 
-    if (extension.generateDaggerFactoriesOnly.get() ||
-      extension.disableComponentMerging.get()
-    ) {
-      // We don't need to disable the incremental compilation for the stub generating task, when we
-      // only generate Dagger factories or contributing modules. That's only needed for merging
-      // Dagger modules.
-      if (projectType.isAndroid) {
-        project.androidVariantsConfigure { variant ->
-          val compileTaskName = "kaptGenerateStubs${variant.name.capitalize(US)}$taskSuffix"
-          disableIncrementalCompilationAction(project, incrementalSignal, compileTaskName)
-        }
-      } else {
-        disableIncrementalCompilationAction(
-          project,
-          incrementalSignal,
-          "kaptGenerateStubs$taskSuffix"
-        )
-        disableIncrementalCompilationAction(
-          project,
-          incrementalSignal,
-          "kaptGenerateStubsTest$taskSuffix"
-        )
-      }
-    } else {
-      project.tasks
-        .withType(KaptGenerateStubsTask::class.java)
-        .configureEach { stubsTask ->
-          // Disable incremental compilation for the stub generating task. Trigger the compiler
-          // plugin if any dependencies in the compile classpath have changed. This will make sure
-          // that we pick up any change from a dependency when merging all the classes. Without
-          // this workaround we could make changes in any library, but these changes wouldn't be
-          // contributed to the Dagger graph, because incremental compilation tricked us.
-          stubsTask.doFirstCompat {
-            stubsTask.incremental = false
-            stubsTask.logger.info(
-              "Anvil: Incremental compilation enabled: ${stubsTask.incremental} (stub)"
+    variant.project.pluginManager.withPlugin("org.jetbrains.kotlin.kapt") {
+      variant.project
+        .namedLazy<KaptGenerateStubsTask>(
+          "kaptGenerateStubs${variant.taskSuffix}"
+        ) { stubsTaskProvider ->
+          if (variant.extension.generateDaggerFactoriesOnly.get() ||
+            variant.extension.disableComponentMerging.get()
+          ) {
+            // We don't need to disable the incremental compilation for the stub generating task
+            // every time, when we only generate Dagger factories or contribute modules. That's only
+            // needed for merging Dagger modules.
+            //
+            // However, we still need to update the generated code when the compiler plugin classpath
+            // changes.
+            disableIncrementalCompilationWhenTheCompilerClasspathChanges(
+              variant, incrementalSignal, stubsTaskProvider
             )
+          } else {
+            stubsTaskProvider.configure { stubsTask ->
+              // Disable incremental compilation for the stub generating task. Trigger the compiler
+              // plugin if any dependencies in the compile classpath have changed. This will make sure
+              // that we pick up any change from a dependency when merging all the classes. Without
+              // this workaround we could make changes in any library, but these changes wouldn't be
+              // contributed to the Dagger graph, because incremental compilation tricked us.
+              stubsTask.doFirstCompat {
+                stubsTask.incremental = false
+                stubsTask.logger.info(
+                  "Anvil: Incremental compilation enabled: ${stubsTask.incremental} (stub)"
+                )
+              }
+            }
           }
         }
     }
-
-    if (projectType.isAndroid) {
-      project.androidVariantsConfigure { variant ->
-        val compileTaskName = "compile${variant.name.capitalize(US)}$taskSuffix"
-        disableIncrementalCompilationAction(project, incrementalSignal, compileTaskName)
-      }
-    } else {
-      // The Java plugin has two Kotlin tasks we care about: compileKotlin and compileTestKotlin.
-      disableIncrementalCompilationAction(project, incrementalSignal, "compile$taskSuffix")
-      disableIncrementalCompilationAction(project, incrementalSignal, "compileTest$taskSuffix")
-    }
   }
 
-  private fun disableIncrementalCompilationAction(
-    project: Project,
+  private fun disableIncrementalCompilationWhenTheCompilerClasspathChanges(
+    variant: Variant,
     incrementalSignal: Provider<IncrementalSignal>,
-    compileTaskName: String
+    compileTaskProvider: TaskProvider<out KotlinCompile>
   ) {
+    val compileTaskName = compileTaskProvider.name
+
     // Disable incremental compilation, if the compiler plugin dependency isn't up-to-date.
     // This will trigger a full compilation of a module using Anvil even though its
     // source files might not have changed. This workaround is necessary, otherwise
     // incremental builds are broken. See https://youtrack.jetbrains.com/issue/KT-38570
-    val disableIncrementalCompilationTaskProvider = project.tasks.register(
+    val disableIncrementalCompilationTaskProvider = variant.project.tasks.register(
       compileTaskName + "CheckIncrementalCompilationAnvil",
       DisableIncrementalCompilationTask::class.java
     ) { task ->
       task.pluginClasspath.from(
-        project.configurations.getByName(PLUGIN_CLASSPATH_CONFIGURATION_NAME)
+        variant.project.configurations.getByName(variant.compilerPluginClasspathName)
       )
       task.incrementalSignal.set(incrementalSignal)
     }
 
-    project.tasks.named(compileTaskName, KotlinCompile::class.java) { compileTask ->
-      compileTask.dependsOn(disableIncrementalCompilationTaskProvider)
-    }
+    compileTaskProvider.dependsOn(disableIncrementalCompilationTaskProvider)
 
     // We avoid a reference to the project in the doFirst.
-    val projectPath = project.path
-
-    // If we merge the block below and the block above, it looks like
-    // the kotlin compiler is generating byte code for
-    // disableIncrementalCompilationTaskProvider to be visible from the doFirst block
+    val projectPath = variant.project.path
 
-    project.tasks.named(compileTaskName, KotlinCompile::class.java) { compileTask ->
+    compileTaskProvider.configure { compileTask ->
       compileTask.doFirstCompat {
         // If the signal is set, then the plugin classpath changed. Apply the setting that
         // DisableIncrementalCompilationTask requested.
@@ -329,77 +274,17 @@ open class AnvilPlugin : KotlinCompilerPluginSupportPlugin {
       }
     }
   }
-}
-
-/**
- * Returns all variants including the androidTest and unit test variants.
- */
-fun Project.androidVariants(): Set<BaseVariant> {
-  return when (val androidExtension = project.extensions.findByName("android")) {
-    is AppExtension ->
-      androidExtension.applicationVariants + androidExtension.testVariants +
-        androidExtension.unitTestVariants
-    is LibraryExtension ->
-      androidExtension.libraryVariants + androidExtension.testVariants +
-        androidExtension.unitTestVariants
-    else -> throw GradleException("Unknown Android module type for project ${project.path}")
-  }
-}
-
-/**
- * Runs the given [action] for each Android variant including androidTest and unit test variants.
- */
-fun Project.androidVariantsConfigure(action: (BaseVariant) -> Unit) {
-  val androidExtension = project.extensions.findByName("android")
 
-  if (androidExtension is AppExtension) {
-    androidExtension.applicationVariants.configureEach(action)
-  }
-  if (androidExtension is LibraryExtension) {
-    androidExtension.libraryVariants.configureEach(action)
-  }
-  if (androidExtension is TestExtension) {
-    androidExtension.applicationVariants.configureEach(action)
-  }
-  if (androidExtension is TestedExtension) {
-    androidExtension.unitTestVariants.configureEach(action)
-    androidExtension.testVariants.configureEach(action)
-  }
-}
-
-fun Collection<BaseVariant>.findVariantForCompileTask(
-  compileTask: KotlinCompile
-): BaseVariant = this
-  .filter { variant ->
-    compileTask.name.contains(variant.name.capitalize(US))
+  private fun getConfiguration(
+    project: Project,
+    buildType: String
+  ): Configuration {
+    val name = if (buildType.isEmpty()) "anvil" else "anvil${buildType.capitalize(US)}"
+    return project.configurations.maybeCreate(name).apply {
+      description = "This configuration is used for dependencies with Anvil CodeGenerator " +
+        "implementations."
+    }
   }
-  .maxByOrNull {
-    // The filter above still returns multiple variants, e.g. for the
-    // "compileDebugUnitTestKotlin" task it returns the variants "debug" and "debugUnitTest".
-    // In this case prefer the variant with the longest matching name, because that's the more
-    // explicit variant that we want.
-    it.name.length
-  }!!
-
-val Project.isKotlinJvmProject: Boolean
-  get() = plugins.hasPlugin(KotlinPluginWrapper::class.java)
-
-val Project.isAndroidProject: Boolean
-  get() = AGP_ON_CLASSPATH &&
-    (plugins.hasPlugin(AppPlugin::class.java) || plugins.hasPlugin(LibraryPlugin::class.java))
-
-@Suppress("SENSELESS_COMPARISON")
-private val AGP_ON_CLASSPATH = try {
-  Class.forName("com.android.build.gradle.AppPlugin") != null
-} catch (t: Throwable) {
-  false
-}
-
-private enum class ProjectType(val isAndroid: Boolean) {
-  JVM(false),
-  ANDROID(true),
-  MULTIPLATFORM_JVM(false),
-  MULTIPLATFORM_ANDROID(true),
 }
 
 /*
@@ -423,3 +308,103 @@ private fun <T : Task> T.doFirstCompat(block: (T) -> Unit) {
     }
   })
 }
+
+/**
+ * Similar to [TaskContainer.named], but waits until the task is registered if it doesn't exist,
+ * yet. If the task is never registered, then this method will throw an error after the
+ * configuration phase.
+ */
+private inline fun <reified T : Task> Project.namedLazy(
+  name: String,
+  crossinline action: (TaskProvider<T>) -> Unit
+) {
+  try {
+    action(tasks.named(name, T::class.java))
+    return
+  } catch (ignored: UnknownTaskException) {
+  }
+
+  var didRun = false
+
+  tasks.whenTaskAdded { task ->
+    if (task.name == name && task is T) {
+      action(tasks.named(name, T::class.java))
+      didRun = true
+    }
+  }
+
+  afterEvaluate {
+    if (!didRun) {
+      throw GradleException("Didn't find task $name with type ${T::class}.")
+    }
+  }
+}
+
+/**
+ * Runs the given [action] for each Android variant including androidTest and unit test variants.
+ */
+private fun Project.androidVariantsConfigure(action: (BaseVariant) -> Unit) {
+  val androidExtension = extensions.findByName("android")
+
+  when (androidExtension) {
+    is AppExtension -> androidExtension.applicationVariants.configureEach(action)
+    is LibraryExtension -> androidExtension.libraryVariants.configureEach(action)
+    is TestExtension -> androidExtension.applicationVariants.configureEach(action)
+  }
+
+  if (androidExtension is TestedExtension) {
+    androidExtension.unitTestVariants.configureEach(action)
+    androidExtension.testVariants.configureEach(action)
+  }
+}
+
+private val agpPlugins = listOf(
+  "com.android.library",
+  "com.android.application",
+  "com.android.test",
+  "com.android.dynamic-feature",
+)
+
+internal class Variant private constructor(
+  val name: String,
+  val project: Project,
+  val extension: AnvilExtension,
+  val compileTaskProvider: TaskProvider<KotlinCompile>,
+  val androidVariant: BaseVariant?,
+  val compilerPluginClasspathName: String,
+) {
+  // E.g. compileKotlin, compileKotlinJvm, compileDebugKotlin.
+  val taskSuffix = compileTaskProvider.name.substringAfter("compile")
+
+  companion object {
+    operator fun invoke(kotlinCompilation: KotlinCompilation<*>): Variant {
+      // Sanity check.
+      require(
+        kotlinCompilation.platformType != androidJvm ||
+          kotlinCompilation is KotlinJvmAndroidCompilation
+      ) {
+        "The KotlinCompilation is KotlinJvmAndroidCompilation, but the platform type " +
+          "is different."
+      }
+
+      val project = kotlinCompilation.target.project
+      val name = kotlinCompilation.name
+
+      @Suppress("UNCHECKED_CAST")
+      return Variant(
+        name = name,
+        project = project,
+        extension = project.extensions.getByType(AnvilExtension::class.java),
+        compileTaskProvider = kotlinCompilation.compileKotlinTaskProvider as
+          TaskProvider<KotlinCompile>,
+        androidVariant = (kotlinCompilation as? KotlinJvmAndroidCompilation)?.androidVariant,
+        compilerPluginClasspathName = PLUGIN_CLASSPATH_CONFIGURATION_NAME +
+          kotlinCompilation.target.targetName.capitalize(US) +
+          name.capitalize(US)
+      ).also {
+        // Sanity check.
+        check(it.compileTaskProvider.name.startsWith("compile"))
+      }
+    }
+  }
+}
diff --git a/gradle-plugin/src/main/java/com/squareup/anvil/plugin/CheckMixedSourceSet.kt b/gradle-plugin/src/main/java/com/squareup/anvil/plugin/CheckMixedSourceSet.kt
index 1a8b45a8..6f472164 100644
--- a/gradle-plugin/src/main/java/com/squareup/anvil/plugin/CheckMixedSourceSet.kt
+++ b/gradle-plugin/src/main/java/com/squareup/anvil/plugin/CheckMixedSourceSet.kt
@@ -17,26 +17,24 @@ import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
  * The workaround for now is to set the kotlin.incremental.usePreciseJavaTracking flag to false for
  * these module using this task.
  */
-object CheckMixedSourceSet {
+internal object CheckMixedSourceSet {
 
-  fun preparePreciseJavaTrackingCheck(compileTask: KotlinCompile): Input {
-    return Input(
-      sources = getAndroidSourceDirs(compileTask)
-        ?: getJvmSourceDirs(compileTask.project)
-        ?: compileTask.project.files(),
-      isKaptApplied = compileTask.project.isKaptApplied()
+  fun preparePreciseJavaTrackingCheck(variant: Variant): Input =
+    Input(
+      sources = getAndroidSourceDirs(variant) ?: getJvmSourceDirs(variant),
+      isKaptApplied = variant.project.isKaptApplied()
     )
-  }
 
   fun disablePreciseJavaTrackingIfNeeded(
     compileTask: KotlinCompile,
     input: Input
   ) {
-    val sourceFiles = input.sources
-      .asSequence()
-      .flatMap { it.walk() }
-      .filter { it.isFile }
-    val hasJavaFile = sourceFiles.any { it.extension == "java" }
+    val hasJavaFile = input.sources
+      ?.asSequence()
+      ?.flatMap { it.walk() }
+      ?.filter { it.isFile }
+      ?.any { it.extension == "java" }
+      ?: false
 
     // If there is Java file, then disable precise Java tracking.
     //
@@ -47,31 +45,18 @@ object CheckMixedSourceSet {
     }
   }
 
-  private fun getAndroidSourceDirs(
-    compileTask: KotlinCompile
-  ): FileCollection? {
-    val project = compileTask.project
-    return if (project.isAndroidProject) {
-      project.androidVariants()
-        .findVariantForCompileTask(compileTask)
-        .sourceSets
-        .flatMap { it.javaDirectories }
-        .let { project.files(it) }
-    } else {
-      null
-    }
+  private fun getAndroidSourceDirs(variant: Variant): FileCollection? {
+    return variant.androidVariant
+      ?.sourceSets
+      ?.flatMap { it.javaDirectories }
+      ?.let { variant.project.files(it) }
   }
 
-  private fun getJvmSourceDirs(project: Project): FileCollection? {
-    return if (project.isKotlinJvmProject) {
-      project.convention.getPlugin(JavaPluginConvention::class.java)
-        .sourceSets
-        // Ignore "test", similar to androidVariants() we ignore unit tests.
-        .single { it.name == "main" }
-        .allJava
-    } else {
-      null
-    }
+  private fun getJvmSourceDirs(variant: Variant): FileCollection? {
+    return variant.project.convention.findPlugin(JavaPluginConvention::class.java)
+      ?.sourceSets
+      ?.single { it.name == variant.name }
+      ?.allJava
   }
 
   private fun Project.isKaptApplied(): Property<Boolean> =
@@ -82,7 +67,7 @@ object CheckMixedSourceSet {
       }
 
   class Input(
-    val sources: FileCollection,
+    val sources: FileCollection?,
     val isKaptApplied: Property<Boolean>
   )
 }


Test patch:
diff --git a/integration-tests/code-generator-tests/build.gradle b/integration-tests/code-generator-tests/build.gradle
index 5da4f283..d9649e1d 100644
--- a/integration-tests/code-generator-tests/build.gradle
+++ b/integration-tests/code-generator-tests/build.gradle
@@ -2,12 +2,8 @@ apply plugin: 'org.jetbrains.kotlin.jvm'
 apply plugin: 'com.squareup.anvil'
 apply plugin: 'org.jetbrains.kotlin.kapt'
 
-anvil {
-  generateDaggerFactories = true
-}
-
 dependencies {
-  anvil project(':integration-tests:code-generator')
+  anvilTest project(':integration-tests:code-generator')
 
   testImplementation testFixtures(project(":compiler-utils"))
   testImplementation deps.dagger2.dagger
diff --git a/integration-tests/mpp/android-module/build.gradle b/integration-tests/mpp/android-module/build.gradle
index 74f5e357..ad35e53a 100644
--- a/integration-tests/mpp/android-module/build.gradle
+++ b/integration-tests/mpp/android-module/build.gradle
@@ -53,4 +53,8 @@ kotlin {
 
 dependencies {
   kapt deps.dagger2.compiler
+
+  // This dependency isn't needed. It's only here for testing purposes (this is still an
+  // integration test).
+  anvilAndroidTest project(':integration-tests:code-generator')
 }