[FIR] Add processor for transforming statuses with plugins

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

@@ -12,7 +12,9 @@ import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirImplicitTypeBodyResolveProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractResolveProcessor
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirGlobalExtensionStatusProcessor
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirPluginAnnotationsResolveProcessor
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirTransformerBasedExtensionStatusProcessor
 
 fun FirResolvePhase.createProcessorByPhase(
     session: FirSession,
@@ -27,6 +29,10 @@ fun FirResolvePhase.createProcessorByPhase(
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
         TYPES -> FirTypeResolveProcessor(session, scopeSession)
+        EXTENSION_STATUS_UPDATE -> when (compilerMode) {
+            CompilerMode.CLI -> FirGlobalExtensionStatusProcessor(session, scopeSession)
+            CompilerMode.IDE -> FirTransformerBasedExtensionStatusProcessor(session, scopeSession)
+        }
         STATUS -> FirStatusResolveProcessor(session, scopeSession)
         CONTRACTS -> FirContractResolveProcessor(session, scopeSession)
         IMPLICIT_TYPES_BODY_RESOLVE -> FirImplicitTypeBodyResolveProcessor(session, scopeSession)

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

@@ -0,0 +1,118 @@
+/*
+ * 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.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+
+class FirGlobalExtensionStatusProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirGlobalResolveProcessor(session, scopeSession) {
+    override fun process() {
+        val extensions = session.extensionService.statusTransformerExtensions
+        if (extensions.isEmpty()) return
+        val provider = session.predicateBasedProvider
+        for (extension in extensions) {
+            val declarations = provider.getSymbolsByPredicate(extension.predicate)
+            for (declaration in declarations) {
+                // TODO: maybe replace with visitor?
+                if (declaration is FirMemberDeclaration) {
+                    val newStatus = extension.transformStatus(declaration, 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 fun FirMemberDeclaration.updateStatus() {
+            if (extensions.isEmpty()) return
+            var status = this.status
+            for (extension in extensions) {
+                status = extension.transformStatus(this, status)
+            }
+            transformStatus(ReplaceStatus, status)
+        }
+
+        override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+            return element.compose()
+        }
+
+        override fun transformTypeAlias(typeAlias: FirTypeAlias, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+            typeAlias.updateStatus()
+            return typeAlias.compose()
+        }
+
+        override fun transformRegularClass(regularClass: FirRegularClass, data: Nothing?): CompositeTransformResult<FirStatement> {
+            regularClass.updateStatus()
+            regularClass.transformDeclarations(this, data)
+            return regularClass.compose()
+        }
+
+        override fun transformConstructor(constructor: FirConstructor, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+            constructor.updateStatus()
+            return constructor.compose()
+        }
+
+        override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+            property.updateStatus()
+            property.transformGetter(this, data)
+            property.transformSetter(this, data)
+            return property.compose()
+        }
+
+        override fun transformField(field: FirField, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+            field.updateStatus()
+            return field.compose()
+        }
+
+        override fun transformEnumEntry(enumEntry: FirEnumEntry, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+            enumEntry.updateStatus()
+            return enumEntry.compose()
+        }
+
+        override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+            simpleFunction.updateStatus()
+            return simpleFunction.compose()
+        }
+    }
+}
+
+
+private object ReplaceStatus : FirTransformer<FirDeclarationStatus>() {
+    override fun <E : FirElement> transformElement(element: E, data: FirDeclarationStatus): CompositeTransformResult<E> {
+        return element.compose()
+    }
+
+    override fun transformDeclarationStatus(
+        declarationStatus: FirDeclarationStatus,
+        data: FirDeclarationStatus
+    ): CompositeTransformResult<FirDeclarationStatus> {
+        return data.compose()
+    }
+}

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

@@ -13,6 +13,7 @@ enum class FirResolvePhase {
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,
     TYPES,
+    EXTENSION_STATUS_UPDATE,
     STATUS,
     CONTRACTS,
     IMPLICIT_TYPES_BODY_RESOLVE,
@@ -22,6 +23,7 @@ enum class FirResolvePhase {
         get() = when (this) {
             RAW_FIR -> RAW_FIR
             IMPORTS -> RAW_FIR
+            STATUS -> TYPES
             IMPLICIT_TYPES_BODY_RESOLVE, BODY_RESOLVE -> STATUS
             else -> values()[ordinal - 1]
         }