[FIR] Add owners to FirStatusTransformerExtension.transformStatus

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<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)

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<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>
 }

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)
+            val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
-            for (declaration in declarations) {
+            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)
         }