Support replaces attribute in @ContributesBinding and `@Contribut…

…esTo`

This change adds the `replaces` attribute in `@ContributesBinding` and `@ContributesTo`. It allows one contributed type to replace other contributed types.

Fixes #79

compiler-utils/src/main/kotlin/software/amazon/lastmile/kotlin/inject/anvil/ContextAware.kt

@@ -226,6 +226,36 @@ interface ContextAware {
         return classNames.replace(".", separator)
     }
 
+    fun KSAnnotation.getReplaces(): List<KSClassDeclaration> {
+        val argument = arguments.firstOrNull { it.name?.asString() == "replaces" }
+            ?: return emptyList()
+
+        @Suppress("UNCHECKED_CAST")
+        return (argument.value as? List<KSType>)
+            ?.map { it.declaration as KSClassDeclaration }
+            ?: emptyList()
+    }
+
+    fun checkReplacesHasSameScope(
+        clazz: KSClassDeclaration,
+        annotations: List<KSAnnotation>,
+    ) {
+        val scope = clazz.scope()
+
+        annotations
+            .flatMap { it.getReplaces() }
+            .map { replacedClass ->
+                checkHasScope(replacedClass)
+
+                check(scope == replacedClass.scope(), clazz) {
+                    "Replaced types must use the same scope. ${clazz.requireQualifiedName()} " +
+                        "uses scope ${scope.type}, but tries to replace " +
+                        "${replacedClass.requireQualifiedName()} using scope " +
+                        "${replacedClass.scope().type}."
+                }
+            }
+    }
+
     /**
      * Return `software.amazon.Test` into `SoftwareAmazonTest`.
      */

compiler-utils/src/testFixtures/kotlin/software/amazon/lastmile/kotlin/inject/anvil/CommonSourceCode.kt

@@ -17,6 +17,12 @@ import java.lang.reflect.Modifier
 val JvmCompilationResult.componentInterface: Class<*>
     get() = classLoader.loadClass("software.amazon.test.ComponentInterface")
 
+val Class<*>.kotlinInjectComponent: Class<*>
+    get() = classLoader.loadClass(
+        "$packageName.KotlinInject" +
+            canonicalName.substring(packageName.length + 1).replace(".", ""),
+    )
+
 val Class<*>.inner: Class<*>
     get() = classes.single { it.simpleName == "Inner" }
 

compiler/src/main/kotlin/software/amazon/lastmile/kotlin/inject/anvil/processor/ContributesBindingProcessor.kt

@@ -81,6 +81,7 @@ internal class ContributesBindingProcessor(
 
         val annotations = clazz.findAnnotationsAtLeastOne(ContributesBinding::class)
         checkNoDuplicateBoundTypes(clazz, annotations)
+        checkReplacesHasSameScope(clazz, annotations)
 
         val boundTypes = annotations
             .map {

compiler/src/main/kotlin/software/amazon/lastmile/kotlin/inject/anvil/processor/ContributesToProcessor.kt

@@ -61,6 +61,8 @@ internal class ContributesToProcessor(
     private fun generateComponentInterface(clazz: KSClassDeclaration) {
         val componentClassName = ClassName(LOOKUP_PACKAGE, clazz.safeClassName)
 
+        checkReplacesHasSameScope(clazz, listOf(clazz.findAnnotation(ContributesTo::class)))
+
         val fileSpec = FileSpec.builder(componentClassName)
             .addType(
                 TypeSpec

compiler/src/main/kotlin/software/amazon/lastmile/kotlin/inject/anvil/processor/MergeComponentProcessor.kt

@@ -140,6 +140,7 @@ internal class MergeComponentProcessor(
         return emptyList()
     }
 
+    @Suppress("LongMethod")
     private fun generateComponentInterface(
         resolver: Resolver,
         clazz: KSClassDeclaration,
@@ -178,6 +179,50 @@ internal class MergeComponentProcessor(
                     it.contributedSubcomponent().requireQualifiedName() !in excludeNames
             }
             .toList()
+            .let { componentInterfaces ->
+                // This block removes all contributed classes that were replaced. For this we
+                // need the full list contributed classes. For each class we get all replaced
+                // classes. Contributed classes that are not part of the replaced classes should
+                // be merged. Contributed classes that are part of the replaced class list should
+                // be removed.
+
+                val replacedClasses = componentInterfaces
+                    .flatMap { it.originChain() }
+                    .mapNotNull { contributedClass ->
+                        val contributeAnnotations = contributedClass
+                            .findAnnotations(ContributesBinding::class)
+                            .plus(contributedClass.findAnnotations(ContributesTo::class))
+
+                        val replaceClasses = contributeAnnotations.flatMap { it.getReplaces() }
+                        if (replaceClasses.isEmpty()) {
+                            return@mapNotNull null
+                        }
+
+                        contributedClass to replaceClasses
+                    }
+                    .toMap()
+                    .let { originToReplacedClasses ->
+                        // originToReplacedClasses contains contributed classes (key) with all
+                        // the classes (value) that it replaces.
+
+                        val allReplacedClasses = originToReplacedClasses.values.flatten()
+
+                        // A class that replaces other class but is being replaced itself should
+                        // not replace the other classes. That's what this filter ensures.
+                        originToReplacedClasses.filter { (contributedClass, _) ->
+                            contributedClass !in allReplacedClasses
+                        }
+                    }
+                    .values
+                    .flatten()
+                    .map { it.requireQualifiedName() }
+
+                componentInterfaces.filter { componentInterface ->
+                    componentInterface.originChain().none { origin ->
+                        origin.requireQualifiedName() in replacedClasses
+                    }
+                }
+            }
 
         val generatedSubcomponents = contributesSubcomponentProcessor.generateFinalComponents(
             parentScopeComponent = clazz,

compiler/src/test/kotlin/software/amazon/lastmile/kotlin/inject/anvil/processor/ContributesBindingProcessorTest.kt

@@ -314,6 +314,66 @@ class ContributesBindingProcessorTest {
         }
     }
 
+    @Test
+    fun `a replaced binding must use the same scope`() {
+        compile(
+            """
+            package software.amazon.test
+                            
+            import me.tatarka.inject.annotations.Inject
+            import software.amazon.lastmile.kotlin.inject.anvil.AppScope
+            import software.amazon.lastmile.kotlin.inject.anvil.ContributesBinding
+
+            interface Base
+
+            @Inject
+            @ContributesBinding(AppScope::class)
+            class Impl : Base
+
+            @Inject
+            @ContributesBinding(Unit::class, replaces = [Impl::class])
+            class Fake : Base
+            """,
+            exitCode = COMPILATION_ERROR,
+        ) {
+            assertThat(messages).contains(
+                "Replaced types must use the same scope. software.amazon.test.Fake " +
+                    "uses scope Unit, but tries to replace software.amazon.test.Impl using " +
+                    "scope AppScope.",
+            )
+        }
+    }
+
+    @Test
+    fun `a replaced binding must use the same scope without named parameter`() {
+        compile(
+            """
+            package software.amazon.test
+                            
+            import me.tatarka.inject.annotations.Inject
+            import software.amazon.lastmile.kotlin.inject.anvil.AppScope
+            import software.amazon.lastmile.kotlin.inject.anvil.ContributesBinding
+
+            interface Base
+
+            @Inject
+            @ContributesBinding(AppScope::class)
+            class Impl : Base
+
+            @Inject
+            @ContributesBinding(Unit::class, Unit::class, false, [Impl::class])
+            class Fake : Base
+            """,
+            exitCode = COMPILATION_ERROR,
+        ) {
+            assertThat(messages).contains(
+                "Replaced types must use the same scope. software.amazon.test.Fake " +
+                    "uses scope Unit, but tries to replace software.amazon.test.Impl using " +
+                    "scope AppScope.",
+            )
+        }
+    }
+
     private val JvmCompilationResult.base: Class<*>
         get() = classLoader.loadClass("software.amazon.test.Base")
 

compiler/src/test/kotlin/software/amazon/lastmile/kotlin/inject/anvil/processor/ContributesToProcessorTest.kt

@@ -93,4 +93,60 @@ class ContributesToProcessorTest {
             assertThat(messages).contains("Only interfaces can be contributed.")
         }
     }
+
+    @Test
+    fun `a replaced component must use the same scope`() {
+        compile(
+            """
+            package software.amazon.test
+                            
+            import me.tatarka.inject.annotations.Inject
+            import software.amazon.lastmile.kotlin.inject.anvil.AppScope
+            import software.amazon.lastmile.kotlin.inject.anvil.ContributesTo
+
+            interface Base
+
+            @ContributesTo(AppScope::class)
+            interface Component1
+
+            @ContributesTo(Unit::class, replaces = [Component1::class])
+            interface Component2
+            """,
+            exitCode = COMPILATION_ERROR,
+        ) {
+            assertThat(messages).contains(
+                "Replaced types must use the same scope. software.amazon.test." +
+                    "Component2 uses scope Unit, but tries to replace software.amazon.test." +
+                    "Component1 using scope AppScope.",
+            )
+        }
+    }
+
+    @Test
+    fun `a replaced component must use the same scope without named parameter`() {
+        compile(
+            """
+            package software.amazon.test
+                            
+            import me.tatarka.inject.annotations.Inject
+            import software.amazon.lastmile.kotlin.inject.anvil.AppScope
+            import software.amazon.lastmile.kotlin.inject.anvil.ContributesTo
+
+            interface Base
+
+            @ContributesTo(AppScope::class)
+            interface Component1
+
+            @ContributesTo(Unit::class, [Component1::class])
+            interface Component2
+            """,
+            exitCode = COMPILATION_ERROR,
+        ) {
+            assertThat(messages).contains(
+                "Replaced types must use the same scope. software.amazon.test." +
+                    "Component2 uses scope Unit, but tries to replace software.amazon.test." +
+                    "Component1 using scope AppScope.",
+            )
+        }
+    }
 }

compiler/src/test/kotlin/software/amazon/lastmile/kotlin/inject/anvil/processor/GenerateKotlinInjectComponentProcessorTest.kt

@@ -13,6 +13,7 @@ import org.junit.jupiter.api.Test
 import software.amazon.lastmile.kotlin.inject.anvil.compile
 import software.amazon.lastmile.kotlin.inject.anvil.componentInterface
 import software.amazon.lastmile.kotlin.inject.anvil.inner
+import software.amazon.lastmile.kotlin.inject.anvil.kotlinInjectComponent
 import software.amazon.lastmile.kotlin.inject.anvil.newComponent
 import java.lang.reflect.Field
 import java.lang.reflect.Method
@@ -475,12 +476,6 @@ class GenerateKotlinInjectComponentProcessorTest {
     private val JvmCompilationResult.impl: Class<*>
         get() = classLoader.loadClass("software.amazon.test.Impl")
 
-    private val Class<*>.kotlinInjectComponent: Class<*>
-        get() = classLoader.loadClass(
-            "$packageName.KotlinInject" +
-                canonicalName.substring(packageName.length + 1).replace(".", ""),
-        )
-
     private val Class<*>.createFunction: Method
         get() = classLoader.loadClass("${canonicalName}Kt").methods.single { it.name == "create" }