Fix returning super type in Anvil. Both in embedded compiler and KSP

square · Nov 27, 2024 · 800ce3b · 800ce3b
1 parent 6b3d029
commit 800ce3b
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 5 deletions.
diff --git a/compiler/src/main/java/com/squareup/anvil/compiler/codegen/ContributesSubcomponentCodeGen.kt b/compiler/src/main/java/com/squareup/anvil/compiler/codegen/ContributesSubcomponentCodeGen.kt
@@ -208,9 +208,13 @@ internal object ContributesSubcomponentCodeGen : AnvilApplicabilityChecker {
         .filter { it.isAbstract }
         .toList()
 
-      if (functions.size != 1 || functions[0].returnType?.resolve()
-          ?.resolveKSClassDeclaration() != this
-      ) {
+      val returnType = functions.singleOrNull()?.returnType?.resolve()?.resolveKSClassDeclaration()
+      if (returnType != this) {
+
+        val isReturnSuperType = returnType != null && this.superTypes
+          .any { type -> type.resolve().resolveKSClassDeclaration() == returnType }
+        if (isReturnSuperType) return
+
         throw KspAnvilException(
           node = factory,
           message = "A factory must have exactly one abstract function returning the " +
@@ -393,6 +397,11 @@ internal object ContributesSubcomponentCodeGen : AnvilApplicabilityChecker {
         ?.asClassReference()
 
       if (returnType != this) {
+
+        val isReturnSuperType = returnType != null && this.directSuperTypeReferences()
+          .any { it.asClassReference() == returnType }
+        if (isReturnSuperType) return
+
         throw AnvilCompilationExceptionClassReference(
           classReference = factory,
           message = "A factory must have exactly one abstract function returning the " +

diff --git a/.../main/java/com/squareup/anvil/compiler/codegen/ContributesSubcomponentHandlerGenerator.kt b/.../main/java/com/squareup/anvil/compiler/codegen/ContributesSubcomponentHandlerGenerator.kt
@@ -340,6 +340,10 @@ internal class ContributesSubcomponentHandlerGenerator(
           )
         }
 
+        val superTypes by lazy {
+          contribution.clazz.directSuperTypeReferences()
+            .map { it.asClassReference() }
+        }
         val createComponentFunctions = factory.memberFunctions
           // filter by `isAbstract` even for interfaces,
           // otherwise we get `toString()`, `equals()`, and `hashCode()`.
@@ -349,7 +353,7 @@ internal class ContributesSubcomponentHandlerGenerator(
               ?.asClassReference()
               ?: return@filter false
 
-            returnType.fqName == contributionFqName
+            returnType.fqName == contributionFqName || superTypes.any { it == returnType }
           }
 
         if (createComponentFunctions.size != 1) {

diff --git a/...src/test/java/com/squareup/anvil/compiler/codegen/ContributesSubcomponentGeneratorTest.kt b/...src/test/java/com/squareup/anvil/compiler/codegen/ContributesSubcomponentGeneratorTest.kt
@@ -110,7 +110,7 @@ class ContributesSubcomponentGeneratorTest(
     }
   }
 
-  @Test fun `there is a hint for contributed subcomponents with an interace factory`() {
+  @Test fun `there is a hint for contributed subcomponents with an interface factory`() {
     compile(
       """
       package com.squareup.test
@@ -571,6 +571,46 @@ class ContributesSubcomponentGeneratorTest(
     }
   }
 
+  @Test fun `a factory function may returns the component super type`() {
+    compile(
+      """
+        package com.squareup.test
+
+        import com.squareup.anvil.annotations.ContributesSubcomponent
+        import com.squareup.anvil.annotations.ContributesSubcomponent.Factory
+        import com.squareup.anvil.annotations.ContributesTo
+        import com.squareup.anvil.annotations.MergeComponent
+
+        interface BaseSubcomponentInterface {
+            interface Factory {
+                fun createComponent(): BaseSubcomponentInterface
+            }
+        }
+
+        @ContributesSubcomponent(Any::class, parentScope = Unit::class)
+        interface SubcomponentInterface : BaseSubcomponentInterface {
+          @Factory
+          interface ComponentFactory: BaseSubcomponentInterface.Factory
+
+          @ContributesTo(Unit::class)
+          interface ParentComponent {
+            fun createFactory(): ComponentFactory
+          }
+        }
+
+        @MergeComponent(Unit::class)
+        interface ComponentInterface
+      """,
+      mode = mode,
+    ) {
+      assertThat(subcomponentInterface.hintSubcomponent?.java).isEqualTo(subcomponentInterface)
+      assertThat(subcomponentInterface.hintSubcomponentParentScope).isEqualTo(Unit::class)
+
+      assertThat(subcomponentInterface.componentFactoryInterface.methods.map { it.name })
+        .containsExactly("createComponent")
+    }
+  }
+
   @Test
   fun `using Dagger's @Subcomponent_Factory is an error`() {
     compile(