MerlinTHS
12/02/2023, 10:42 PMFirCompanionGenerationProcessorrunResolutionJavier
12/03/2023, 8:42 AMMerlinTHS
12/03/2023, 10:35 AMimport org.jetbrains.kotlin.GeneratedDeclarationKey
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
import org.jetbrains.kotlin.fir.extensions.*
import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
import org.jetbrains.kotlin.fir.plugin.createCompanionObject
import org.jetbrains.kotlin.fir.plugin.createDefaultPrivateConstructor
import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlin.name.SpecialNames
import org.jetbrains.kotlin.name.SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT
import org.jetbrains.kotlin.utils.addToStdlib.runIf
class CompanionGenerator(session: FirSession, annotations: List<String>) : FirDeclarationGenerationExtension(session) {
    object Key : GeneratedDeclarationKey() {
       override fun toString() = "CompanionGeneratorKey"
    }
    private val companionPredicate = LookupPredicate.create { 
       annotated(annotations.map(::FqName)) 
    }
    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
       register(companionPredicate)
    }
    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name, context: NestedClassGenerationContext): FirClassLikeSymbol<*>? =
        runIf(name == DEFAULT_NAME_FOR_COMPANION_OBJECT) { createCompanionObject(owner, Key).symbol }
    override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
        val constructor = createDefaultPrivateConstructor(context.owner, Key)
        return listOf(constructor.symbol)
    }
    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
       if (!classSymbol.isCompanion) return emptySet()
        val origin = classSymbol.origin as? FirDeclarationOrigin.Plugin
        return runIf(origin?.key == Key) { setOf(SpecialNames.INIT) }.orEmpty()
    }
    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> =
        runIf(classSymbol matches companionPredicate and classSymbol.needsCompanion) {
            setOf(DEFAULT_NAME_FOR_COMPANION_OBJECT)
        }.orEmpty()
    private infix fun FirClassSymbol<*>.matches(predicate: LookupPredicate): Boolean =
       session.predicateBasedProvider.matches(predicate, this)
}
private val FirClassSymbol<*>.isCompanion get() =
    isSingleton && with(classId) {
       isNestedClass && shortClassName == DEFAULT_NAME_FOR_COMPANION_OBJECT
    }
private val FirClassSymbol<*>.needsCompanion get() =
    !isSingleton && declarationSymbols.none { (it as? FirClassSymbol<*>)?.isCompanion ?: false }
private val FirClassSymbol<*>.isSingleton get() =
    classKind == ClassKind.OBJECTneedsCompanionJavier
12/03/2023, 10:38 AMMerlinTHS
12/03/2023, 11:20 AMJavier
12/03/2023, 11:27 AMJavier
12/03/2023, 11:28 AMMerlinTHS
12/03/2023, 11:33 AMMerlinTHS
12/03/2023, 11:43 AMJavier
12/03/2023, 2:32 PM