diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt index 41d450c5efe..19a6f431c40 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt @@ -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>> + abstract fun getOwnersOfDeclaration(declaration: FirAnnotatedDeclaration): List? + abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList) 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? { + return cache.ownersForDeclaration[declaration] + } + override fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration) { val owners = cache.ownersForDeclaration.getValue(owner).add(owner) registerAnnotatedDeclaration(declaration, owners) diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt index 55b72a32c61..716a1564c5e 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt @@ -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 = FirStatusTransformerExtension::class - abstract fun transformStatus(declaration: FirDeclaration, status: FirDeclarationStatus): FirDeclarationStatus + abstract fun transformStatus(declaration: FirDeclaration, owners: List, status: FirDeclarationStatus): FirDeclarationStatus fun interface Factory : FirExtension.Factory } diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalExtensionStatusProcessor.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalExtensionStatusProcessor.kt index d7cd03df993..fa055d71bde 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalExtensionStatusProcessor.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalExtensionStatusProcessor.kt @@ -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() { 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) }