Introduce merge creator annotations (#13)

* Introduce mirror factories

* Support in IR

* Use in KSP, but don't implement yet

Will implement this with contributes subcomponent handler as it has some shared code we wanna reuse

* Formatting

annotations/api/annotations.api

@@ -72,20 +72,48 @@ public abstract interface annotation class com/squareup/anvil/annotations/MergeC
 	public abstract fun scope ()Ljava/lang/Class;
 }
 
+public abstract interface annotation class com/squareup/anvil/annotations/MergeComponent$Builder : java/lang/annotation/Annotation {
+}
+
+public abstract interface annotation class com/squareup/anvil/annotations/MergeComponent$Builder$Container : java/lang/annotation/Annotation {
+	public abstract fun value ()[Lcom/squareup/anvil/annotations/MergeComponent$Builder;
+}
+
 public abstract interface annotation class com/squareup/anvil/annotations/MergeComponent$Container : java/lang/annotation/Annotation {
 	public abstract fun value ()[Lcom/squareup/anvil/annotations/MergeComponent;
 }
 
+public abstract interface annotation class com/squareup/anvil/annotations/MergeComponent$Factory : java/lang/annotation/Annotation {
+}
+
+public abstract interface annotation class com/squareup/anvil/annotations/MergeComponent$Factory$Container : java/lang/annotation/Annotation {
+	public abstract fun value ()[Lcom/squareup/anvil/annotations/MergeComponent$Factory;
+}
+
 public abstract interface annotation class com/squareup/anvil/annotations/MergeSubcomponent : java/lang/annotation/Annotation {
 	public abstract fun exclude ()[Ljava/lang/Class;
 	public abstract fun modules ()[Ljava/lang/Class;
 	public abstract fun scope ()Ljava/lang/Class;
 }
 
+public abstract interface annotation class com/squareup/anvil/annotations/MergeSubcomponent$Builder : java/lang/annotation/Annotation {
+}
+
+public abstract interface annotation class com/squareup/anvil/annotations/MergeSubcomponent$Builder$Container : java/lang/annotation/Annotation {
+	public abstract fun value ()[Lcom/squareup/anvil/annotations/MergeSubcomponent$Builder;
+}
+
 public abstract interface annotation class com/squareup/anvil/annotations/MergeSubcomponent$Container : java/lang/annotation/Annotation {
 	public abstract fun value ()[Lcom/squareup/anvil/annotations/MergeSubcomponent;
 }
 
+public abstract interface annotation class com/squareup/anvil/annotations/MergeSubcomponent$Factory : java/lang/annotation/Annotation {
+}
+
+public abstract interface annotation class com/squareup/anvil/annotations/MergeSubcomponent$Factory$Container : java/lang/annotation/Annotation {
+	public abstract fun value ()[Lcom/squareup/anvil/annotations/MergeSubcomponent$Factory;
+}
+
 public abstract interface annotation class com/squareup/anvil/annotations/compat/MergeInterfaces : java/lang/annotation/Annotation {
 	public abstract fun exclude ()[Ljava/lang/Class;
 	public abstract fun scope ()Ljava/lang/Class;

annotations/src/main/java/com/squareup/anvil/annotations/MergeComponent.kt

@@ -52,4 +52,22 @@ public annotation class MergeComponent(
    * same scope, but should be excluded from the component.
    */
   val exclude: Array<KClass<*>> = [],
-)
+) {
+  /**
+   * Mirrors `dagger.Component.Factory` and just conveys that this annotated factory should be
+   * mirrored in the merged component class.
+   */
+  @Target(CLASS)
+  @Retention(RUNTIME)
+  @Repeatable
+  public annotation class Factory
+
+  /**
+   * Mirrors `dagger.Component.Builder` and just conveys that this annotated builder should be
+   * mirrored in the merged component class.
+   */
+  @Target(CLASS)
+  @Retention(RUNTIME)
+  @Repeatable
+  public annotation class Builder
+}

annotations/src/main/java/com/squareup/anvil/annotations/MergeSubcomponent.kt

@@ -71,4 +71,22 @@ public annotation class MergeSubcomponent(
    * same scope, but should be excluded from the subcomponent.
    */
   val exclude: Array<KClass<*>> = [],
-)
+) {
+  /**
+   * Mirrors `dagger.Subcomponent.Factory` and just conveys that this annotated factory should be
+   * mirrored in the merged component class.
+   */
+  @Target(CLASS)
+  @Retention(RUNTIME)
+  @Repeatable
+  public annotation class Factory
+
+  /**
+   * Mirrors `dagger.Subcomponent.Builder` and just conveys that this annotated builder should be
+   * mirrored in the merged component class.
+   */
+  @Target(CLASS)
+  @Retention(RUNTIME)
+  @Repeatable
+  public annotation class Builder
+}

compiler/src/main/java/com/squareup/anvil/compiler/IrContributionMerger.kt

@@ -65,6 +65,23 @@ internal class IrContributionMerger(
 
           val mergeAnnotatedClass = declaration.symbol.toClassReference(pluginContext)
 
+          // For @Merge*.Factory/Builder annotations, generate the "real" one onto the class
+          // This is for backward compatibility with K1 support
+          for (creatorAnnotation in mergeAnnotatedClass.annotations.findAll(
+            CREATOR_ANNOTATIONS.keys,
+          )) {
+            val daggerAnnotation = CREATOR_ANNOTATIONS.getValue(creatorAnnotation.fqName)
+            pluginContext.irBuiltIns.createIrBuilder(declaration.symbol)
+              .generateCreatorAnnotation(
+                daggerAnnotationFqName = daggerAnnotation,
+                pluginContext = pluginContext,
+                declaration = mergeAnnotatedClass,
+              )
+
+            // Nothing else to do on this class
+            return super.visitClass(declaration)
+          }
+
           val mergeComponentAnnotations = mergeAnnotatedClass.annotations
             .findAll(mergeComponentFqName, mergeSubcomponentFqName)
 
@@ -372,6 +389,23 @@ internal class IrContributionMerger(
     declaration.clazz.owner.annotations += annotationConstructorCall
   }
 
+  private fun IrBuilderWithScope.generateCreatorAnnotation(
+    daggerAnnotationFqName: FqName,
+    pluginContext: IrPluginContext,
+    declaration: ClassReferenceIr,
+  ) {
+    val annotationConstructorCall = irCallConstructor(
+      callee = pluginContext
+        .referenceConstructors(daggerAnnotationFqName.classIdBestGuess())
+        .single { it.owner.isPrimary },
+      typeArguments = emptyList(),
+    )
+
+    // Since we are modifying the state of the code here, this does not need to be reflected in
+    // the associated [ClassReferenceIr] which is more of an initial snapshot.
+    declaration.clazz.owner.annotations += annotationConstructorCall
+  }
+
   private fun checkSameScope(
     contributedClass: ClassReferenceIr,
     classToReplace: ClassReferenceIr,
@@ -651,3 +685,10 @@ private fun ClassReferenceIr.atLeastOneAnnotation(
       )
     }
 }
+
+private val CREATOR_ANNOTATIONS = mapOf(
+  mergeComponentFactoryFqName to daggerComponentFactoryFqName,
+  mergeComponentBuilderFqName to daggerComponentBuilderFqName,
+  mergeSubcomponentFactoryFqName to daggerSubcomponentFactoryFqName,
+  mergeSubcomponentBuilderFqName to daggerSubcomponentBuilderFqName,
+)

compiler/src/main/java/com/squareup/anvil/compiler/KspContributionMerger.kt

@@ -66,7 +66,6 @@ import com.squareup.kotlinpoet.joinToCode
 import com.squareup.kotlinpoet.ksp.toAnnotationSpec
 import com.squareup.kotlinpoet.ksp.toClassName
 import com.squareup.kotlinpoet.ksp.writeTo
-import dagger.Component
 import dagger.Module
 import org.jetbrains.kotlin.name.Name
 
@@ -726,11 +725,14 @@ internal class KspContributionMerger(override val env: SymbolProcessorEnvironmen
           addSuperinterface(contributedInterface)
         }
 
+        // TODO actually generate the creators too
+        //  - Extend the original
+        //  - Annotate with the real dagger annotation
+        //  - Add a binding module binding the extension as the root
         val componentOrFactory = mergeAnnotatedClass.declarations
           .filterIsInstance<KSClassDeclaration>()
           .singleOrNull {
-            // TODO does dagger also use these names? Or are they lowercase versions of the simple class name?
-            if (it.isAnnotationPresent<Component.Factory>()) {
+            if (it.isAnnotationPresent<MergeComponent.Factory>() || it.isAnnotationPresent<MergeSubcomponent.Factory>()) {
               factoryOrBuilderFunSpec = FunSpec.builder("factory")
                 .returns(generatedComponentClassName.nestedClass(it.simpleName.asString()))
                 .addStatement(
@@ -742,7 +744,7 @@ internal class KspContributionMerger(override val env: SymbolProcessorEnvironmen
                 .build()
               return@singleOrNull true
             }
-            if (it.isAnnotationPresent<Component.Builder>()) {
+            if (it.isAnnotationPresent<MergeComponent.Builder>() || it.isAnnotationPresent<MergeSubcomponent.Builder>()) {
               factoryOrBuilderFunSpec = FunSpec.builder("builder")
                 .returns(generatedComponentClassName.nestedClass(it.simpleName.asString()))
                 .addStatement(

compiler/src/main/java/com/squareup/anvil/compiler/Utils.kt

@@ -31,8 +31,12 @@ import javax.inject.Provider
 import javax.inject.Qualifier
 
 internal val mergeComponentFqName = MergeComponent::class.fqName
+internal val mergeComponentFactoryFqName = MergeComponent.Factory::class.fqName
+internal val mergeComponentBuilderFqName = MergeComponent.Builder::class.fqName
 internal val mergeComponentClassName = MergeComponent::class.asClassName()
 internal val mergeSubcomponentFqName = MergeSubcomponent::class.fqName
+internal val mergeSubcomponentFactoryFqName = MergeSubcomponent.Factory::class.fqName
+internal val mergeSubcomponentBuilderFqName = MergeSubcomponent.Builder::class.fqName
 internal val mergeSubcomponentClassName = MergeSubcomponent::class.asClassName()
 internal val mergeInterfacesFqName = MergeInterfaces::class.fqName
 internal val mergeInterfacesClassName = MergeInterfaces::class.asClassName()
@@ -45,6 +49,8 @@ internal val contributesSubcomponentFqName = ContributesSubcomponent::class.fqNa
 internal val contributesSubcomponentFactoryFqName = ContributesSubcomponent.Factory::class.fqName
 internal val internalBindingMarkerFqName = InternalBindingMarker::class.fqName
 internal val daggerComponentFqName = Component::class.fqName
+internal val daggerComponentFactoryFqName = Component.Factory::class.fqName
+internal val daggerComponentBuilderFqName = Component.Builder::class.fqName
 internal val daggerComponentClassName = Component::class.asClassName()
 internal val daggerSubcomponentFqName = Subcomponent::class.fqName
 internal val daggerSubcomponentClassName = Subcomponent::class.asClassName()

compiler/src/main/java/com/squareup/anvil/compiler/codegen/reference/AnnotatedReferenceIr.kt

@@ -21,6 +21,11 @@ internal fun List<AnnotationReferenceIr>.find(
 internal fun List<AnnotationReferenceIr>.findAll(
   vararg annotationNames: FqName,
   scopeName: FqName? = null,
+) = findAll(annotationNames.toSet(), scopeName)
+
+internal fun List<AnnotationReferenceIr>.findAll(
+  annotationNames: Set<FqName>,
+  scopeName: FqName? = null,
 ): List<AnnotationReferenceIr> {
   return filter {
     it.fqName in annotationNames && (scopeName == null || it.scopeOrNull?.fqName == scopeName)