[FIR] Add owners to FirStatusTransformerExtension.transformStatus

This commit is contained in:
Dmitriy Novozhilov
2020-06-03 10:55:47 +03:00
parent b79d6aced4
commit 514cbc77fc
3 changed files with 15 additions and 8 deletions
@@ -5,17 +5,14 @@
package org.jetbrains.kotlin.fir.extensions
import com.google.common.collect.ArrayListMultimap
import com.google.common.collect.LinkedHashMultimap
import com.google.common.collect.Multimap
import kotlinx.collections.immutable.PersistentList
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.FirSessionComponent
import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
import org.jetbrains.kotlin.fir.declarations.FirRegularClass
import org.jetbrains.kotlin.fir.extensions.predicate.*
import org.jetbrains.kotlin.fir.resolve.fqName
import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
abstract class FirPredicateBasedProvider : FirSessionComponent {
companion object {
@@ -40,6 +37,8 @@ abstract class FirPredicateBasedProvider : FirSessionComponent {
predicate: DeclarationPredicate
): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>>
abstract fun getOwnersOfDeclaration(declaration: FirAnnotatedDeclaration): List<FirAnnotatedDeclaration>?
abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>)
abstract fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration)
@@ -91,6 +90,10 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
cache.annotationsOfDeclaration.putAll(declaration, matchingAnnotations)
}
override fun getOwnersOfDeclaration(declaration: FirAnnotatedDeclaration): List<FirAnnotatedDeclaration>? {
return cache.ownersForDeclaration[declaration]
}
override fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration) {
val owners = cache.ownersForDeclaration.getValue(owner).add(owner)
registerAnnotatedDeclaration(declaration, owners)
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.extensions
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.descriptors.Visibility
import org.jetbrains.kotlin.fir.FirSession
import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
import org.jetbrains.kotlin.fir.declarations.FirDeclaration
import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
@@ -23,7 +24,7 @@ abstract class FirStatusTransformerExtension(session: FirSession) : FirPredicate
final override val extensionType: KClass<out FirExtension> = FirStatusTransformerExtension::class
abstract fun transformStatus(declaration: FirDeclaration, status: FirDeclarationStatus): FirDeclarationStatus
abstract fun transformStatus(declaration: FirDeclaration, owners: List<FirAnnotatedDeclaration>, status: FirDeclarationStatus): FirDeclarationStatus
fun interface Factory : FirExtension.Factory<FirStatusTransformerExtension>
}
@@ -29,11 +29,11 @@ class FirGlobalExtensionStatusProcessor(
if (extensions.isEmpty()) return
val provider = session.predicateBasedProvider
for (extension in extensions) {
val declarations = provider.getSymbolsByPredicate(extension.predicate)
for (declaration in declarations) {
val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
for ((declaration, owners) in declarations) {
// TODO: maybe replace with visitor?
if (declaration is FirMemberDeclaration) {
val newStatus = extension.transformStatus(declaration, declaration.status)
val newStatus = extension.transformStatus(declaration, owners, declaration.status)
declaration.transformStatus(ReplaceStatus, newStatus)
}
}
@@ -49,12 +49,15 @@ class FirTransformerBasedExtensionStatusProcessor(
private inner class StatusUpdater : FirDefaultTransformer<Nothing?>() {
private val extensions = session.extensionService.statusTransformerExtensions
private val predicateBasedProvider = session.predicateBasedProvider
private fun FirMemberDeclaration.updateStatus() {
if (extensions.isEmpty()) return
val owners = predicateBasedProvider.getOwnersOfDeclaration(this)
requireNotNull(owners)
var status = this.status
for (extension in extensions) {
status = extension.transformStatus(this, status)
status = extension.transformStatus(this, owners, status)
}
transformStatus(ReplaceStatus, status)
}