Generate Dagger modules for @ContributesBinding with basic configuration

compiler-k2/api/compiler-k2.api

@@ -8,9 +8,43 @@ public final class com/squareup/anvil/compiler/k2/fir/AnvilFirExtensionRegistrar
 	public fun <init> (Lorg/jetbrains/kotlin/cli/common/messages/MessageCollector;)V
 }
 
+public final class com/squareup/anvil/compiler/k2/fir/contributions/BindingModuleData {
+	public fun <init> (Lorg/jetbrains/kotlin/name/ClassId;Lorg/jetbrains/kotlin/fir/symbols/impl/FirClassSymbol;Lorg/jetbrains/kotlin/fir/extensions/FirExtension;Lorg/jetbrains/kotlin/fir/FirSession;)V
+	public final fun getBoundType ()Lorg/jetbrains/kotlin/fir/types/ConeKotlinType;
+	public final fun getCallableName ()Lorg/jetbrains/kotlin/name/Name;
+	public final fun getContributesBindingAnnotation ()Lorg/jetbrains/kotlin/fir/expressions/FirAnnotation;
+	public final fun getGeneratedClassId ()Lorg/jetbrains/kotlin/name/ClassId;
+	public final fun getGeneratedClassSymbol ()Lorg/jetbrains/kotlin/fir/symbols/impl/FirClassLikeSymbol;
+	public final fun getMatchedClassSymbol ()Lorg/jetbrains/kotlin/fir/symbols/impl/FirClassSymbol;
+}
+
+public final class com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingFirExtension : com/squareup/anvil/compiler/k2/fir/AnvilFirDeclarationGenerationExtension {
+	public fun <init> (Lcom/squareup/anvil/compiler/k2/fir/AnvilFirContext;Lorg/jetbrains/kotlin/fir/FirSession;)V
+	public fun generateFunctions (Lorg/jetbrains/kotlin/name/CallableId;Lorg/jetbrains/kotlin/fir/extensions/DeclarationGenerationContext$Member;)Ljava/util/List;
+	public fun generateTopLevelClassLikeDeclaration (Lorg/jetbrains/kotlin/name/ClassId;)Lorg/jetbrains/kotlin/fir/symbols/impl/FirClassLikeSymbol;
+	public fun getCallableNamesForClass (Lorg/jetbrains/kotlin/fir/symbols/impl/FirClassSymbol;Lorg/jetbrains/kotlin/fir/extensions/DeclarationGenerationContext$Member;)Ljava/util/Set;
+	public fun getTopLevelClassIds ()Ljava/util/Set;
+	public fun registerPredicates (Lorg/jetbrains/kotlin/fir/extensions/FirDeclarationPredicateRegistrar;)V
+}
+
+public final class com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingSessionComponent : com/squareup/anvil/compiler/k2/fir/AnvilFirExtensionSessionComponent {
+	public fun <init> (Lcom/squareup/anvil/compiler/k2/fir/AnvilFirContext;Lorg/jetbrains/kotlin/fir/FirSession;)V
+	public final fun getBindingModuleCache ()Lorg/jetbrains/kotlin/fir/caches/FirCache;
+	public final fun getGeneratedIdsToMatchedSymbols ()Ljava/util/Map;
+	public fun registerPredicates (Lorg/jetbrains/kotlin/fir/extensions/FirDeclarationPredicateRegistrar;)V
+}
+
+public final class com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingSessionComponentKt {
+	public static final fun getContributesBindingSessionComponent (Lorg/jetbrains/kotlin/fir/FirSession;)Lcom/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingSessionComponent;
+}
+
 public final class com/squareup/anvil/compiler/k2/ir/GeneratedDeclarationsIrBodyFiller : org/jetbrains/kotlin/backend/common/extensions/IrGenerationExtension {
 	public fun <init> ()V
 	public fun generate (Lorg/jetbrains/kotlin/ir/declarations/IrModuleFragment;Lorg/jetbrains/kotlin/backend/common/extensions/IrPluginContext;)V
 	public fun resolveSymbol (Lorg/jetbrains/kotlin/ir/symbols/IrSymbol;Lorg/jetbrains/kotlin/ir/builders/TranslationPluginContext;)Lorg/jetbrains/kotlin/ir/declarations/IrDeclaration;
 }
 
+public final class com/squareup/anvil/compiler/k2/util/UtilsKt {
+	public static final fun classId (Lorg/jetbrains/kotlin/name/FqName;)Lorg/jetbrains/kotlin/name/ClassId;
+}
+

compiler-k2/build.gradle.kts

@@ -32,4 +32,5 @@ dependencies {
   testImplementation(libs.kase)
   testImplementation(libs.kotest.assertions.api)
   testImplementation(libs.kotest.assertions.core.jvm)
+  testImplementation(libs.truth)
 }

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/contributions/BindingModuleData.kt

@@ -0,0 +1,84 @@
+package com.squareup.anvil.compiler.k2.fir.contributions
+
+import com.squareup.anvil.compiler.k2.names.Names
+import com.squareup.anvil.compiler.k2.util.classId
+import com.squareup.anvil.compiler.k2.util.toFirAnnotation
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.getKClassArgument
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirNamedArgumentExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
+import org.jetbrains.kotlin.fir.expressions.impl.FirAnnotationArgumentMappingImpl
+import org.jetbrains.kotlin.fir.extensions.ExperimentalTopLevelDeclarationsGenerationApi
+import org.jetbrains.kotlin.fir.extensions.FirExtension
+import org.jetbrains.kotlin.fir.plugin.createTopLevelClass
+import org.jetbrains.kotlin.fir.resolve.fqName
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.types.constructClassLikeType
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+
+@OptIn(ExperimentalTopLevelDeclarationsGenerationApi::class)
+public class BindingModuleData(
+  public val generatedClassId: ClassId,
+  public val matchedClassSymbol: FirClassSymbol<*>,
+  firExtension: FirExtension,
+  session: FirSession,
+) {
+  public val generatedClassSymbol: FirClassLikeSymbol<*> by lazy {
+    firExtension.createTopLevelClass(
+      classId = generatedClassId,
+      key = GeneratedBindingDeclarationKey,
+      classKind = ClassKind.INTERFACE,
+    ).apply {
+      replaceAnnotations(
+        listOf(
+          buildContributesToAnnotation(),
+          Names.daggerModule.toFirAnnotation(),
+        ),
+      )
+    }.symbol
+  }
+
+  public val contributesBindingAnnotation: FirAnnotation by lazy {
+    matchedClassSymbol.annotations.single {
+      it.fqName(session)?.equals(Names.anvilContributesBinding) ?: false
+    }
+  }
+
+  public val boundType: ConeKotlinType by lazy {
+    contributesBindingAnnotation.getKClassArgument(Name.identifier("boundType"), session)!!
+  }
+
+  public val callableName: Name by lazy {
+    "bind${boundType.classId!!.shortClassName.asString()}"
+      .let(Name::identifier)
+  }
+
+  private fun buildContributesToAnnotation(): FirAnnotation = buildAnnotation {
+    annotationTypeRef = buildResolvedTypeRef {
+      coneType = Names.anvilContributesTo.classId().constructClassLikeType()
+    }
+    // Argument mapping may be empty if merging is also happening
+    val newArgMapping = if (contributesBindingAnnotation.argumentMapping.mapping.isEmpty()) {
+      val scopeArg = (contributesBindingAnnotation as FirAnnotationCall).argumentList
+        .arguments
+        .single { (it as FirNamedArgumentExpression).name.toString() == "scope" }
+
+      mapOf(Name.identifier("scope") to scopeArg)
+    } else {
+      contributesBindingAnnotation.argumentMapping.mapping
+        .filter { (key, _) ->
+          key.asString() == "scope"
+        }
+    }
+
+    argumentMapping = FirAnnotationArgumentMappingImpl(null, newArgMapping)
+  }
+}

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingFirExtension.kt

@@ -0,0 +1,103 @@
+package com.squareup.anvil.compiler.k2.fir.contributions
+
+import com.squareup.anvil.annotations.ContributesBinding
+import com.squareup.anvil.compiler.k2.fir.AnvilFirContext
+import com.squareup.anvil.compiler.k2.fir.AnvilFirDeclarationGenerationExtension
+import com.squareup.anvil.compiler.k2.names.Names
+import com.squareup.anvil.compiler.k2.util.AnvilPredicates.hasContributesBindingAnnotation
+import com.squareup.anvil.compiler.k2.util.toFirAnnotation
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.contains
+import org.jetbrains.kotlin.fir.extensions.ExperimentalTopLevelDeclarationsGenerationApi
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
+import org.jetbrains.kotlin.fir.plugin.createMemberFunction
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.types.constructType
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+
+/**
+ * Generates binding modules for every [ContributesBinding]-annotated class.
+ *
+ * Given a Kotlin source file like:
+ * ```
+ * @ContributesBinding(Any::class)
+ * class Foo @Inject constructor() : Bar
+ * ```
+ *
+ * We will generate the FIR equivalent of:
+ * ```
+ * @ContributesTo(Any::class)
+ * @Module
+ * abstract class Foo_BindingModule {
+ *  @Binds
+ *  abstract fun bindBar(impl: Foo): Bar
+ * }
+ * ```
+ */
+@OptIn(ExperimentalTopLevelDeclarationsGenerationApi::class)
+public class ContributesBindingFirExtension(
+  anvilFirContext: AnvilFirContext,
+  session: FirSession,
+) : AnvilFirDeclarationGenerationExtension(anvilFirContext, session) {
+
+  override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+    register(hasContributesBindingAnnotation)
+  }
+
+  override fun getTopLevelClassIds(): Set<ClassId> {
+    return session.contributesBindingSessionComponent.generatedIdsToMatchedSymbols.keys
+  }
+
+  override fun generateTopLevelClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*> {
+    return session.contributesBindingSessionComponent
+      .bindingModuleCache
+      .getValue(classId, session)
+      .generatedClassSymbol
+  }
+
+  override fun getCallableNamesForClass(
+    classSymbol: FirClassSymbol<*>,
+    context: MemberGenerationContext,
+  ): Set<Name> {
+    val cache = session.contributesBindingSessionComponent.bindingModuleCache
+    if (!cache.contains(classSymbol.classId)) {
+      return emptySet()
+    }
+    val bindingData = cache.getValue(classSymbol.classId, session)
+
+    return setOf(bindingData.callableName)
+  }
+
+  override fun generateFunctions(
+    callableId: CallableId,
+    context: MemberGenerationContext?,
+  ): List<FirNamedFunctionSymbol> {
+    val bindingData = session.contributesBindingSessionComponent.bindingModuleCache.getValue(
+      context!!.owner.classId,
+      session,
+    )
+
+    return listOf(
+      createMemberFunction(
+        owner = context.owner,
+        key = GeneratedBindingDeclarationKey,
+        name = callableId.callableName,
+        returnType = bindingData.boundType,
+      ) {
+        modality = Modality.ABSTRACT
+        valueParameter(
+          name = Name.identifier("concreteType"),
+          type = bindingData.matchedClassSymbol.constructType(),
+        )
+      }.apply {
+        replaceAnnotations(listOf(Names.daggerBinds.toFirAnnotation()))
+      }.symbol,
+    )
+  }
+}

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingFirExtensionFactory.kt

@@ -0,0 +1,17 @@
+package com.squareup.anvil.compiler.k2.fir.contributions
+
+import com.google.auto.service.AutoService
+import com.squareup.anvil.compiler.k2.fir.AnvilFirContext
+import com.squareup.anvil.compiler.k2.fir.AnvilFirDeclarationGenerationExtension
+import com.squareup.anvil.compiler.k2.fir.AnvilFirExtensionFactory
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+
+@AutoService(AnvilFirExtensionFactory::class)
+internal class ContributesBindingFirExtensionFactory : AnvilFirDeclarationGenerationExtension.Factory {
+
+  override fun create(anvilFirContext: AnvilFirContext): FirDeclarationGenerationExtension.Factory {
+    return FirDeclarationGenerationExtension.Factory {
+      ContributesBindingFirExtension(anvilFirContext = anvilFirContext, session = it)
+    }
+  }
+}

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingSessionComponent.kt

@@ -0,0 +1,53 @@
+package com.squareup.anvil.compiler.k2.fir.contributions
+
+import com.squareup.anvil.compiler.k2.fir.AnvilFirContext
+import com.squareup.anvil.compiler.k2.fir.AnvilFirExtensionSessionComponent
+import com.squareup.anvil.compiler.k2.fir.utils.joinSimpleNames
+import com.squareup.anvil.compiler.k2.util.AnvilPredicates.hasContributesBindingAnnotation
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.name.ClassId
+
+/**
+ * Responsible for tracking the classes annotated with @ContributesBinding and creating + caching
+ * their generated Dagger module metadata.
+ */
+public class ContributesBindingSessionComponent(
+  anvilFirContext: AnvilFirContext,
+  session: FirSession,
+) : AnvilFirExtensionSessionComponent(anvilFirContext, session) {
+  /**
+   * A map to help us track the original annotated classes' bindings, and their
+   * generated module IDs.
+   * E.g. Key: "Foo_BindingModule", Value: ClassSymbol<Foo>
+   */
+  public val generatedIdsToMatchedSymbols: Map<ClassId, FirClassSymbol<*>> by lazy {
+    session.predicateBasedProvider.getSymbolsByPredicate(hasContributesBindingAnnotation)
+      .filterIsInstance<FirClassSymbol<*>>()
+      .associateBy {
+        it.classId.joinSimpleNames(suffix = "BindingModule")
+      }
+  }
+
+  public val bindingModuleCache: FirCache<ClassId, BindingModuleData, FirSession> =
+    session.firCachesFactory
+      .createCache<ClassId, BindingModuleData, FirSession> { key: ClassId, context ->
+        BindingModuleData(
+          key,
+          generatedIdsToMatchedSymbols[key] as FirClassSymbol<*>,
+          this@ContributesBindingSessionComponent,
+          session,
+        )
+      }
+
+  override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+    register(hasContributesBindingAnnotation)
+  }
+}
+
+public val FirSession.contributesBindingSessionComponent: ContributesBindingSessionComponent
+  by FirSession.sessionComponentAccessor()

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/contributions/ContributesBindingSessionComponentFactory.kt

@@ -0,0 +1,17 @@
+package com.squareup.anvil.compiler.k2.fir.contributions
+
+import com.google.auto.service.AutoService
+import com.squareup.anvil.compiler.k2.fir.AnvilFirContext
+import com.squareup.anvil.compiler.k2.fir.AnvilFirExtensionFactory
+import com.squareup.anvil.compiler.k2.fir.AnvilFirExtensionSessionComponent
+import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent
+
+@AutoService(AnvilFirExtensionFactory::class)
+internal class ContributesBindingSessionComponentFactory : AnvilFirExtensionSessionComponent.Factory {
+
+  override fun create(anvilFirContext: AnvilFirContext): FirExtensionSessionComponent.Factory {
+    return FirExtensionSessionComponent.Factory {
+      ContributesBindingSessionComponent(anvilFirContext = anvilFirContext, session = it)
+    }
+  }
+}

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/contributions/GeneratedBindingDeclarationKey.kt

@@ -0,0 +1,5 @@
+package com.squareup.anvil.compiler.k2.fir.contributions
+
+import org.jetbrains.kotlin.GeneratedDeclarationKey
+
+internal object GeneratedBindingDeclarationKey : GeneratedDeclarationKey()

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/fir/utils/ClassIdUtils.kt

@@ -14,3 +14,27 @@ internal fun ClassId.append(suffix: String): ClassId {
 internal val FqName.classId: ClassId get() {
   return ClassId.topLevel(this)
 }
+
+/**
+ * Joins the simple names of a class with the given [separator], [prefix], and [suffix].
+ *
+ * ```
+ * val normalName = ClassName("com.example", "Outer", "Middle", "Inner")
+ * val joinedName = normalName.joinSimpleNames(separator = "_", suffix = "Factory")
+ *
+ * println(joinedName) // com.example.Outer_Middle_InnerFactory
+ * ```
+ * @throws IllegalArgumentException if the resulting class name is too long to be a valid file name.
+ */
+internal fun ClassId.joinSimpleNames(
+  separator: String = "_",
+  prefix: String = "",
+  suffix: String = "",
+): ClassId = ClassId.topLevel(
+  packageFqName.child(
+    Name.identifier(
+      relativeClassName.pathSegments()
+        .joinToString(separator = separator, prefix = prefix, postfix = suffix),
+    ),
+  ),
+)

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/names/Names.kt

@@ -3,7 +3,12 @@ package com.squareup.anvil.compiler.k2.names
 import org.jetbrains.kotlin.name.FqName
 
 internal object Names {
+  val anvilContributesBinding = FqName("com.squareup.anvil.annotations.ContributesBinding")
+  val anvilContributesTo = FqName("com.squareup.anvil.annotations.ContributesTo")
+
+  val daggerBinds = FqName("dagger.Binds")
   val daggerFactory = FqName("dagger.internal.Factory")
+  val daggerModule = FqName("dagger.Module")
 
   val kotlinJvmStatic = FqName("kotlin.jvm.JvmStatic")
 

compiler-k2/src/main/kotlin/com/squareup/anvil/compiler/k2/util/AnvilPredicates.kt

@@ -0,0 +1,13 @@
+package com.squareup.anvil.compiler.k2.util
+
+import com.squareup.anvil.compiler.k2.names.Names
+import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
+import org.jetbrains.kotlin.name.FqName
+
+internal object AnvilPredicates {
+  val hasContributesBindingAnnotation
+    get() = Names.anvilContributesBinding.toLookupPredicate()
+
+  private fun FqName.toLookupPredicate(): LookupPredicate =
+    LookupPredicate.create { annotated(this@toLookupPredicate) }
+}