[FIR] Introduce utility service for plugins for matching common predicate pattern

This commit is contained in:
Dmitriy Novozhilov
2022-05-20 11:56:53 +03:00
committed by teamcity
parent 7048cac770
commit 8acb8aa043
3 changed files with 70 additions and 36 deletions
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
class FirAllOpenExtensionRegistrar(val allOpenAnnotationFqNames: List<String>) : FirExtensionRegistrar() {
override fun ExtensionRegistrarContext.configurePlugin() {
+FirAllOpenStatusTransformer.getFactory(allOpenAnnotationFqNames)
+::FirAllOpenStatusTransformer
+FirAllOpenPredicateMatcher.getFactory(allOpenAnnotationFqNames)
}
}
@@ -8,66 +8,36 @@ package org.jetbrains.kotlin.allopen.fir
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.descriptors.Modality
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.caches.getValue
import org.jetbrains.kotlin.fir.copy
import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
import org.jetbrains.kotlin.fir.declarations.FirDeclaration
import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
import org.jetbrains.kotlin.fir.declarations.FirRegularClass
import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent.Factory
import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension
import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension.Factory
import org.jetbrains.kotlin.fir.extensions.predicate.has
import org.jetbrains.kotlin.fir.extensions.predicate.metaHas
import org.jetbrains.kotlin.fir.extensions.predicate.or
import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
import org.jetbrains.kotlin.fir.extensions.utils.AbstractSimpleClassPredicateMatchingService
import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
import org.jetbrains.kotlin.name.FqName
class FirAllOpenStatusTransformer(allOpenAnnotationFqNames: List<String>, session: FirSession) : FirStatusTransformerExtension(session) {
companion object {
fun getFactory(allOpenAnnotationFqNames: List<String>): Factory {
return Factory { session -> FirAllOpenStatusTransformer(allOpenAnnotationFqNames, session) }
}
}
private val annotationFqNames = allOpenAnnotationFqNames.map { FqName(it) }
private val hasPredicate = has(annotationFqNames) or metaHas(annotationFqNames)
private val cache: FirCache<FirRegularClassSymbol, Boolean, Nothing?> = session.firCachesFactory.createCache { symbol, _ ->
symbol.annotated()
}
override fun FirDeclarationPredicateRegistrar.registerPredicates() {
register(hasPredicate)
}
class FirAllOpenStatusTransformer(session: FirSession) : FirStatusTransformerExtension(session) {
override fun needTransformStatus(declaration: FirDeclaration): Boolean {
return when (declaration) {
is FirRegularClass -> declaration.classKind == ClassKind.CLASS && cache.getValue(declaration.symbol)
is FirRegularClass -> declaration.classKind == ClassKind.CLASS && session.allOpenPredicateMatcher.isAnnotated(declaration.symbol)
is FirCallableDeclaration -> {
val parentClassId = declaration.symbol.callableId.classId ?: return false
if (parentClassId.isLocal) return false
val parentClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(parentClassId) as? FirRegularClassSymbol
?: return false
cache.getValue(parentClassSymbol)
session.allOpenPredicateMatcher.isAnnotated(parentClassSymbol)
}
else -> false
}
}
private fun FirRegularClassSymbol.annotated(): Boolean {
if (session.predicateBasedProvider.matches(hasPredicate, this)) return true
return resolvedSuperTypes.any {
val superSymbol = it.toRegularClassSymbol(session) ?: return@any false
cache.getValue(superSymbol)
}
}
override fun transformStatus(status: FirDeclarationStatus, declaration: FirDeclaration): FirDeclarationStatus {
return if (status.modality == null) {
status.copy(modality = Modality.OPEN)
@@ -76,3 +46,21 @@ class FirAllOpenStatusTransformer(allOpenAnnotationFqNames: List<String>, sessio
}
}
}
class FirAllOpenPredicateMatcher(
session: FirSession,
allOpenAnnotationFqNames: List<String>
) : AbstractSimpleClassPredicateMatchingService(session) {
companion object {
fun getFactory(allOpenAnnotationFqNames: List<String>): Factory {
return Factory { session -> FirAllOpenPredicateMatcher(session, allOpenAnnotationFqNames) }
}
}
override val predicate = run {
val annotationFqNames = allOpenAnnotationFqNames.map { FqName(it) }
has(annotationFqNames) or metaHas(annotationFqNames)
}
}
val FirSession.allOpenPredicateMatcher: FirAllOpenPredicateMatcher by FirSession.sessionComponentAccessor()