[FIR] Move call of status transform extensions to FirStatusResolver

It makes more sense to run status transformer extension
  right in status resolve stage, because it reduces
  number of compiler stages

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransformerExtension.kt

@@ -8,9 +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.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import kotlin.reflect.KClass
 
@@ -24,11 +22,94 @@ abstract class FirStatusTransformerExtension(session: FirSession) : FirPredicate
 
     final override val extensionType: KClass<out FirExtension> = FirStatusTransformerExtension::class
 
-    abstract fun transformStatus(
-        declaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>,
-        status: FirDeclarationStatus
-    ): FirDeclarationStatus
+    protected open fun transformStatus(
+        status: FirDeclarationStatus,
+        declaration: FirAnnotatedDeclaration
+    ): FirDeclarationStatus {
+        return status
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        property: FirProperty,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, property)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        function: FirSimpleFunction,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, function)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        regularClass: FirRegularClass,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, regularClass)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        typeAlias: FirTypeAlias,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, typeAlias)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        propertyAccessor: FirPropertyAccessor,
+        containingClass: FirClass?,
+        containingProperty: FirProperty?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, propertyAccessor)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        constructor: FirConstructor,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, constructor)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        field: FirField,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, field)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        backingField: FirBackingField,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, backingField)
+    }
+
+    open fun transformStatus(
+        status: FirDeclarationStatus,
+        enumEntry: FirEnumEntry,
+        containingClass: FirClass?,
+        isLocal: Boolean
+    ): FirDeclarationStatus {
+        return transformStatus(status, enumEntry)
+    }
 
     fun interface Factory : FirExtension.Factory<FirStatusTransformerExtension>
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolver.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.effectiveVisibility
 import org.jetbrains.kotlin.fir.declarations.utils.isExpect
 import org.jetbrains.kotlin.fir.declarations.utils.isOverride
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
+import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
@@ -44,6 +45,22 @@ class FirStatusResolver(
             FirDeclarationStatusImpl.Modifier.values().toList() - NOT_INHERITED_MODIFIERS
     }
 
+    private val extensionStatusTransformers = session.extensionService.statusTransformerExtensions
+
+    private inline fun <T> T.applyExtensionTransformers(
+        operation: FirStatusTransformerExtension.(FirDeclarationStatus) -> FirDeclarationStatus
+    ): FirDeclarationStatus where T : FirMemberDeclaration, T : FirAnnotatedDeclaration {
+        if (extensionStatusTransformers.isEmpty()) return status
+        val declaration = this
+        return extensionStatusTransformers.fold(status) { acc, it ->
+            if (session.predicateBasedProvider.matches(it.predicate, declaration)) {
+                it.operation(acc)
+            } else {
+                acc
+            }
+        }
+    }
+
     fun resolveStatus(
         declaration: FirDeclaration,
         containingClass: FirClass?,
@@ -65,7 +82,8 @@ class FirStatusResolver(
 
     @OptIn(ExperimentalStdlibApi::class)
     fun resolveStatus(property: FirProperty, containingClass: FirClass?, isLocal: Boolean): FirResolvedDeclarationStatus {
-        return resolveStatus(property, property.status, containingClass, null, isLocal) l@{
+        val status = property.applyExtensionTransformers { transformStatus(it, property, containingClass, isLocal) }
+        return resolveStatus(property, status, containingClass, null, isLocal) l@{
             if (containingClass == null) return@l emptyList()
             @Suppress("RemoveExplicitTypeArguments") // Workaround for KT-42175
             buildList<FirProperty> {
@@ -84,7 +102,8 @@ class FirStatusResolver(
 
     @OptIn(ExperimentalStdlibApi::class)
     fun resolveStatus(function: FirSimpleFunction, containingClass: FirClass?, isLocal: Boolean): FirResolvedDeclarationStatus {
-        return resolveStatus(function, function.status, containingClass, null, isLocal) l@{
+        val status = function.applyExtensionTransformers { transformStatus(it, function, containingClass, isLocal) }
+        return resolveStatus(function, status, containingClass, null, isLocal) l@{
             if (containingClass == null) return@l emptyList()
             @Suppress("RemoveExplicitTypeArguments") // Workaround for KT-42175
             buildList<FirCallableDeclaration> {
@@ -106,7 +125,8 @@ class FirStatusResolver(
         containingClass: FirClass?,
         isLocal: Boolean
     ): FirResolvedDeclarationStatus {
-        return resolveStatus(regularClass, regularClass.status, containingClass, null, isLocal) { emptyList() }
+        val status = regularClass.applyExtensionTransformers { transformStatus(it, regularClass, containingClass, isLocal) }
+        return resolveStatus(regularClass, status, containingClass, null, isLocal) { emptyList() }
     }
 
     fun resolveStatus(
@@ -114,7 +134,8 @@ class FirStatusResolver(
         containingClass: FirClass?,
         isLocal: Boolean
     ): FirResolvedDeclarationStatus {
-        return resolveStatus(typeAlias, typeAlias.status, containingClass, null, isLocal) { emptyList() }
+        val status = typeAlias.applyExtensionTransformers { transformStatus(it, typeAlias, containingClass, isLocal) }
+        return resolveStatus(typeAlias, status, containingClass, null, isLocal) { emptyList() }
     }
 
     fun resolveStatus(
@@ -123,11 +144,15 @@ class FirStatusResolver(
         containingProperty: FirProperty?,
         isLocal: Boolean
     ): FirResolvedDeclarationStatus {
-        return resolveStatus(propertyAccessor, propertyAccessor.status, containingClass, containingProperty, isLocal) { emptyList() }
+        val status = propertyAccessor.applyExtensionTransformers {
+            transformStatus(it, propertyAccessor, containingClass, containingProperty, isLocal)
+        }
+        return resolveStatus(propertyAccessor, status, containingClass, containingProperty, isLocal) { emptyList() }
     }
 
     fun resolveStatus(constructor: FirConstructor, containingClass: FirClass?, isLocal: Boolean): FirResolvedDeclarationStatus {
-        return resolveStatus(constructor, constructor.status, containingClass, null, isLocal) { emptyList() }
+        val status = constructor.applyExtensionTransformers { transformStatus(it, constructor, containingClass, isLocal) }
+        return resolveStatus(constructor, status, containingClass, null, isLocal) { emptyList() }
     }
 
     fun resolveStatus(field: FirField, containingClass: FirClass?, isLocal: Boolean): FirResolvedDeclarationStatus {
@@ -139,11 +164,13 @@ class FirStatusResolver(
         containingClass: FirClass?,
         isLocal: Boolean
     ): FirResolvedDeclarationStatus {
-        return resolveStatus(backingField, backingField.status, containingClass, null, isLocal) { emptyList() }
+        val status = backingField.applyExtensionTransformers { transformStatus(it, backingField, containingClass, isLocal) }
+        return resolveStatus(backingField, status, containingClass, null, isLocal) { emptyList() }
     }
 
     fun resolveStatus(enumEntry: FirEnumEntry, containingClass: FirClass?, isLocal: Boolean): FirResolvedDeclarationStatus {
-        return resolveStatus(enumEntry, enumEntry.status, containingClass, null, isLocal) { emptyList() }
+        val status = enumEntry.applyExtensionTransformers { transformStatus(it, enumEntry, containingClass, isLocal) }
+        return resolveStatus(enumEntry, status, containingClass, null, isLocal) { emptyList() }
     }
 
     private inline fun resolveStatus(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirTotalResolveProcessor.kt

@@ -71,7 +71,6 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
         TYPES -> FirTypeResolveProcessor(session, scopeSession)
-        EXTENSION_STATUS_UPDATE -> FirGlobalExtensionStatusProcessor(session, scopeSession)
         STATUS -> FirStatusResolveProcessor(session, scopeSession)
         ARGUMENTS_OF_ANNOTATIONS -> FirAnnotationArgumentsResolveProcessor(session, scopeSession)
         CONTRACTS -> FirContractResolveProcessor(session, scopeSession)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalExtensionStatusProcessor.kt

@@ -1,118 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.fir.resolve.transformers.plugin
-
-import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
-import org.jetbrains.kotlin.fir.extensions.statusTransformerExtensions
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirGlobalResolveProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
-import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
-import org.jetbrains.kotlin.fir.visitors.FirTransformer
-
-class FirGlobalExtensionStatusProcessor(
-    session: FirSession,
-    scopeSession: ScopeSession
-) : FirGlobalResolveProcessor(session, scopeSession) {
-    override fun process(files: Collection<FirFile>) {
-        val extensions = session.extensionService.statusTransformerExtensions
-        if (extensions.isEmpty()) return
-        val provider = session.predicateBasedProvider
-        for (extension in extensions) {
-            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, owners, declaration.status)
-                    declaration.transformStatus(ReplaceStatus, newStatus)
-                }
-            }
-        }
-    }
-}
-
-class FirTransformerBasedExtensionStatusProcessor(
-    session: FirSession,
-    scopeSession: ScopeSession
-) : FirTransformerBasedResolveProcessor(session, scopeSession) {
-    override val transformer: FirTransformer<Nothing?> = StatusUpdater()
-
-    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 as FirAnnotatedDeclaration)
-            requireNotNull(owners)
-            var status = this.status
-            for (extension in extensions) {
-                status = extension.transformStatus(this, owners, status)
-            }
-            transformStatus(ReplaceStatus, status)
-        }
-
-        override fun <E : FirElement> transformElement(element: E, data: Nothing?): E {
-            return element
-        }
-
-        override fun transformTypeAlias(typeAlias: FirTypeAlias, data: Nothing?): FirTypeAlias {
-            typeAlias.updateStatus()
-            return typeAlias
-        }
-
-        override fun transformRegularClass(regularClass: FirRegularClass, data: Nothing?): FirStatement {
-            regularClass.updateStatus()
-            regularClass.transformDeclarations(this, data)
-            return regularClass
-        }
-
-        override fun transformConstructor(constructor: FirConstructor, data: Nothing?): FirConstructor {
-            constructor.updateStatus()
-            return constructor
-        }
-
-        override fun transformProperty(property: FirProperty, data: Nothing?): FirProperty {
-            property.updateStatus()
-            property.transformGetter(this, data)
-            property.transformSetter(this, data)
-            return property
-        }
-
-        override fun transformField(field: FirField, data: Nothing?): FirField {
-            field.updateStatus()
-            return field
-        }
-
-        override fun transformEnumEntry(enumEntry: FirEnumEntry, data: Nothing?): FirEnumEntry {
-            enumEntry.updateStatus()
-            return enumEntry
-        }
-
-        override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?): FirSimpleFunction {
-            simpleFunction.updateStatus()
-            return simpleFunction
-        }
-    }
-}
-
-private object ReplaceStatus : FirTransformer<FirDeclarationStatus>() {
-    override fun <E : FirElement> transformElement(element: E, data: FirDeclarationStatus): E {
-        return element
-    }
-
-    override fun transformDeclarationStatus(
-        declarationStatus: FirDeclarationStatus,
-        data: FirDeclarationStatus
-    ): FirDeclarationStatus {
-        return data
-    }
-}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirResolvePhase.kt

@@ -13,7 +13,6 @@ enum class FirResolvePhase(val pluginPhase: Boolean = false, val noProcessor: Bo
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,
     TYPES,
-    EXTENSION_STATUS_UPDATE(pluginPhase = true),
     STATUS,
     ARGUMENTS_OF_ANNOTATIONS,
     CONTRACTS,

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/AllOpenStatusTransformer.kt

@@ -24,7 +24,7 @@ class AllOpenStatusTransformer(session: FirSession) : FirStatusTransformerExtens
         private val ALL_OPEN = FqName("org.jetbrains.kotlin.fir.plugin.AllOpen")
     }
 
-    override fun transformStatus(declaration: FirDeclaration, owners: List<FirAnnotatedDeclaration>, status: FirDeclarationStatus): FirDeclarationStatus {
+    override fun transformStatus(status: FirDeclarationStatus, declaration: FirAnnotatedDeclaration): FirDeclarationStatus {
         if (status.modality != null) return status
         return status.transform(modality = Modality.OPEN)
     }

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/AllOpenVisibilityTransformer.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.expressions.arguments
 import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.hasOrUnder
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.transform
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
@@ -39,11 +40,8 @@ class AllOpenVisibilityTransformer(session: FirSession) : FirStatusTransformerEx
         private val ProtectedName = Name.identifier("Protected")
     }
 
-    override fun transformStatus(
-        declaration: FirDeclaration,
-        owners: List<FirAnnotatedDeclaration>,
-        status: FirDeclarationStatus
-    ): FirDeclarationStatus {
+    override fun transformStatus(status: FirDeclarationStatus, declaration: FirAnnotatedDeclaration): FirDeclarationStatus {
+        val owners = session.predicateBasedProvider.getOwnersOfDeclaration(declaration) ?: emptyList()
         val visibility = findVisibility(declaration, owners) ?: return status
         if (visibility == status.visibility) return status
         return status.transform(visibility = visibility)