ContributesMultibindingScopedMetroExtension.kt

package com.squareup.metro.extensions.scoped

import com.squareup.metro.extensions.fir.extractScopeClassId
import dev.zacsweers.metro.compiler.MetroOptions
import dev.zacsweers.metro.compiler.api.fir.MetroContributionExtension
import dev.zacsweers.metro.compiler.compat.CompatContext
import dev.zacsweers.metro.compiler.fir.MetroFirTypeResolver
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
import org.jetbrains.kotlin.fir.resolve.defaultType
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
import org.jetbrains.kotlin.fir.scopes.getSingleClassifier
import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
import org.jetbrains.kotlin.name.ClassId

/**
 * Implements [MetroContributionExtension] to tell Metro's `ContributedInterfaceSupertypeGenerator`
 * about contributions from our generated `MultibindingScopedContribution` interfaces.
 *
 * This is necessary because Metro's predicate-based provider only sees source declarations and
 * declarations generated by generators that registered the predicate for `@ContributesTo`. Our
 * `MultibindingScopedContribution` is generated by [ContributesMultibindingScopedFir] (which
 * registered a predicate for `@ContributesMultibindingScoped`), so Metro's in-compilation discovery
 * doesn't see it.
 */
public class ContributesMultibindingScopedMetroExtension(private val session: FirSession) :
  MetroContributionExtension {

  private val predicate = ContributesMultibindingScopedIds.PREDICATE

  private val annotatedClasses by lazy {
    session.predicateBasedProvider
      .getSymbolsByPredicate(predicate)
      .filterIsInstance<FirRegularClassSymbol>()
      .toList()
  }

  override fun FirDeclarationPredicateRegistrar.registerPredicates() {
    register(predicate)
  }

  override fun getContributions(
    scopeClassId: ClassId,
    typeResolverFactory: MetroFirTypeResolver.Factory,
  ): List<MetroContributionExtension.Contribution> {
    return annotatedClasses.mapNotNull { parentSymbol ->
      val annotationScopeClassId =
        extractScopeClassId(
          parentSymbol,
          ContributesMultibindingScopedIds.CONTRIBUTES_MULTIBINDING_SCOPED_CLASS_ID,
          session,
        ) ?: return@mapNotNull null

      if (annotationScopeClassId != scopeClassId) return@mapNotNull null

      val contributionInterfaceClassId =
        parentSymbol.classId.createNestedClassId(
          ContributesMultibindingScopedIds.NESTED_INTERFACE_NAME
        )

      val contributionSymbol =
        session.symbolProvider.getClassLikeSymbolByClassId(contributionInterfaceClassId)
          as? FirRegularClassSymbol ?: return@mapNotNull null

      // Trigger Metro's FIR generator to create the MetroContribution nested class.
      val scope = contributionSymbol.declaredMemberScope(session, memberRequiredPhase = null)
      val metroContributionName =
        scope.getClassifierNames().firstOrNull { it.identifier.startsWith("MetroContributionTo") }
          ?: return@mapNotNull null

      val metroContributionSymbol =
        scope.getSingleClassifier(metroContributionName) as? FirRegularClassSymbol
          ?: return@mapNotNull null

      MetroContributionExtension.Contribution(
        supertype = metroContributionSymbol.defaultType(),
        replaces = emptyList(),
        originClassId = parentSymbol.classId,
      )
    }
  }

  public class Factory : MetroContributionExtension.Factory {
    override fun create(
      session: FirSession,
      options: MetroOptions,
      compatContext: CompatContext,
    ): MetroContributionExtension {
      return ContributesMultibindingScopedMetroExtension(session)
    }
  }
}