[FIR] Add resolve phase for resolving arguments of plugin's annotations

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/declarations/GeneratedDeclarationValidation.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.declarations
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationResolveStatus
 import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.references.FirNamedReference
@@ -30,7 +31,7 @@ object FirGeneratedElementsValidator : FirDefaultVisitorVoid() {
     }
 
     override fun visitAnnotationCall(annotationCall: FirAnnotationCall) {
-        require(annotationCall.resolved)
+        require(annotationCall.resolveStatus == FirAnnotationResolveStatus.Resolved)
         annotationCall.acceptChildren(this)
     }
 

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

@@ -32,12 +32,17 @@ class FirExtensionService(val session: FirSession) : ComponentArrayOwner<FirExte
     var registeredExtensionsSize: Int = 0
         private set
 
+    var registeredPredicateBasedExtensionsSize: Int = 0
+        private set
+
     @PluginServicesInitialization
     fun registerExtensions(extensionClass: KClass<out FirExtension>, extensionFactories: List<FirExtension.Factory<*>>) {
         registeredExtensionsSize += extensionFactories.size
+        val extensions = extensionFactories.map { it.create(session) }
+        registeredPredicateBasedExtensionsSize += extensions.count { it is FirPredicateBasedExtension }
         registerComponent(
             extensionClass,
-            extensionFactories.map { it.create(session) }
+            extensions
         )
     }
 
@@ -51,3 +56,6 @@ val FirSession.extensionService: FirExtensionService by FirSession.sessionCompon
 
 val FirExtensionService.hasExtensions: Boolean
     get() = registeredExtensionsSize > 0
+
+val FirExtensionService.hasPredicateBasedExtensions: Boolean
+    get() = registeredPredicateBasedExtensionsSize > 0

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

@@ -30,6 +30,7 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
         TYPES -> FirTypeResolveProcessor(session, scopeSession)
+        ARGUMENTS_OF_PLUGIN_ANNOTATIONS -> FirAnnotationArgumentsResolveProcessor(session, scopeSession)
         EXTENSION_STATUS_UPDATE -> FirGlobalExtensionStatusProcessor(session, scopeSession)
         STATUS -> FirStatusResolveProcessor(session, scopeSession)
         CONTRACTS -> FirContractResolveProcessor(session, scopeSession)
@@ -51,6 +52,7 @@ fun FirResolvePhase.createTransformerBasedProcessorByPhase(
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
         TYPES -> FirTypeResolveProcessor(session, scopeSession)
+        ARGUMENTS_OF_PLUGIN_ANNOTATIONS -> FirAnnotationArgumentsResolveProcessor(session, scopeSession)
         EXTENSION_STATUS_UPDATE -> FirTransformerBasedExtensionStatusProcessor(session, scopeSession)
         STATUS -> FirStatusResolveProcessor(session, scopeSession)
         CONTRACTS -> FirContractResolveProcessor(session, scopeSession)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -39,7 +39,7 @@ open class FirBodyResolveTransformer(
     final override val components: BodyResolveTransformerComponents =
         BodyResolveTransformerComponents(session, scopeSession, this, context)
 
-    internal val expressionsTransformer = FirExpressionsResolveTransformer(this)
+    internal open val expressionsTransformer = FirExpressionsResolveTransformer(this)
     protected open val declarationsTransformer = FirDeclarationsResolveTransformer(this)
     private val controlFlowStatementsTransformer = FirControlFlowStatementsResolveTransformer(this)
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -40,7 +40,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance
 
-class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
+open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) : FirPartialBodyResolveTransformer(transformer) {
     private inline val builtinTypes: BuiltinTypes get() = session.builtinTypes
     private val arrayOfCallTransformer = FirArrayOfCallTransformer()
     var enableArrayOfCallTransformation = false
@@ -74,7 +74,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         qualifiedAccessExpression: FirQualifiedAccessExpression,
         data: ResolutionMode,
     ): CompositeTransformResult<FirStatement> {
-        qualifiedAccessExpression.annotations.forEach { it.accept(this, data) }
+        qualifiedAccessExpression.transformAnnotations(this, data)
         qualifiedAccessExpression.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
 
         var result = when (val callee = qualifiedAccessExpression.calleeReference) {
@@ -148,6 +148,9 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
                 // NB: here we can get raw expression because of dropped qualifiers (see transform callee),
                 // so candidate existence must be checked before calling completion
                 if (transformedCallee is FirQualifiedAccessExpression && transformedCallee.candidate() != null) {
+                    if (!transformedCallee.isAcceptableResolvedQualifiedAccess()) {
+                        return qualifiedAccessExpression.compose()
+                    }
                     callCompleter.completeCall(transformedCallee, data.expectedType).result
                 } else {
                     transformedCallee
@@ -167,6 +170,10 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         return result.compose()
     }
 
+    protected open fun FirQualifiedAccessExpression.isAcceptableResolvedQualifiedAccess(): Boolean {
+        return true
+    }
+
     override fun transformSafeCallExpression(
         safeCallExpression: FirSafeCallExpression,
         data: ResolutionMode
@@ -623,15 +630,24 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
 
     override fun transformAnnotationCall(annotationCall: FirAnnotationCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         if (annotationCall.resolveStatus == FirAnnotationResolveStatus.Resolved) return annotationCall.compose()
+        return resolveAnnotationCall(annotationCall, data, FirAnnotationResolveStatus.Resolved)
+    }
+
+    protected fun resolveAnnotationCall(
+        annotationCall: FirAnnotationCall,
+        data: ResolutionMode,
+        status: FirAnnotationResolveStatus
+    ): CompositeTransformResult<FirAnnotationCall> {
         dataFlowAnalyzer.enterAnnotationCall(annotationCall)
         return withFirArrayOfCallTransformer {
             (annotationCall.transformChildren(transformer, data) as FirAnnotationCall).also {
                 // TODO: it's temporary incorrect solution until we design resolve and completion for annotation calls
                 it.argumentList.transformArguments(integerLiteralTypeApproximator, null)
-                it.replaceResolveStatus(FirAnnotationResolveStatus.Resolved)
+                it.replaceResolveStatus(status)
                 dataFlowAnalyzer.exitAnnotationCall(it)
             }.compose()
         }
+
     }
 
     private fun ConeTypeProjection.toFirTypeProjection(): FirTypeProjection = when (this) {

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

@@ -0,0 +1,42 @@
+/*
+ * 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.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.AdapterForResolveProcessor
+import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+
+@OptIn(AdapterForResolveProcessor::class)
+class FirAnnotationArgumentsResolveProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirTransformerBasedResolveProcessor(session, scopeSession) {
+    override val transformer: FirTransformer<Nothing?> = FirAnnotationArgumentsResolveTransformerAdapter(session, scopeSession)
+}
+
+@AdapterForResolveProcessor
+class FirAnnotationArgumentsResolveTransformerAdapter(session: FirSession, scopeSession: ScopeSession) : FirTransformer<Nothing?>() {
+    private val transformer = FirAnnotationArgumentsResolveTransformer(session, scopeSession)
+    private val hasAnnotations = session.registeredPluginAnnotations.annotations.isNotEmpty()
+
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+        return element.compose()
+    }
+
+    override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        if (!hasAnnotations) return file.compose()
+        return file.transform(transformer, ResolutionMode.ContextIndependent)
+    }
+}

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

@@ -0,0 +1,200 @@
+/*
+ * 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.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
+import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
+import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirExpressionsResolveTransformer
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.compose
+
+class FirAnnotationArgumentsResolveTransformer(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    outerBodyResolveContext: BodyResolveContext? = null
+) : FirBodyResolveTransformer(
+    session,
+    FirResolvePhase.ARGUMENTS_OF_PLUGIN_ANNOTATIONS,
+    implicitTypeOnly = false,
+    scopeSession,
+    outerBodyResolveContext = outerBodyResolveContext
+) {
+    override val expressionsTransformer: FirExpressionsResolveTransformer = FirExpressionsResolveTransformerForSpecificAnnotations(
+        this,
+        session.registeredPluginAnnotations.annotations
+    )
+
+    override val declarationsTransformer: FirDeclarationsResolveTransformer = FirDeclarationsResolveTransformerForArgumentAnnotations(this)
+}
+
+private class FirDeclarationsResolveTransformerForArgumentAnnotations(
+    transformer: FirBodyResolveTransformer
+) : FirDeclarationsResolveTransformer(transformer) {
+    override fun transformWrappedDelegateExpression(
+        wrappedDelegateExpression: FirWrappedDelegateExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return wrappedDelegateExpression.compose()
+    }
+
+    override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        return regularClass.transformAnnotations(this, data).transformDeclarations(this, data).compose()
+    }
+
+    override fun transformAnonymousInitializer(
+        anonymousInitializer: FirAnonymousInitializer,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirDeclaration> {
+        return anonymousInitializer.compose()
+    }
+
+    override fun transformSimpleFunction(
+        simpleFunction: FirSimpleFunction,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirSimpleFunction> {
+        return simpleFunction.transformAnnotations(this, data).compose()
+    }
+
+    override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+        return constructor.transformAnnotations(this, data).compose()
+    }
+
+    override fun transformValueParameter(valueParameter: FirValueParameter, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        return valueParameter.transformAnnotations(this, data).compose()
+    }
+
+    override fun transformProperty(property: FirProperty, data: ResolutionMode): CompositeTransformResult<FirProperty> {
+        property.transformAnnotations(this, data)
+        property.transformGetter(this, data)
+        property.transformSetter(this, data)
+        return property.compose()
+    }
+
+    override fun transformPropertyAccessor(
+        propertyAccessor: FirPropertyAccessor,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        propertyAccessor.transformAnnotations(this, data)
+        return propertyAccessor.compose()
+    }
+}
+
+private class FirExpressionsResolveTransformerForSpecificAnnotations(
+    transformer: FirBodyResolveTransformer,
+    private val annotations: Set<AnnotationFqn>
+) : FirExpressionsResolveTransformer(transformer) {
+    private var annotationArgumentsMode: Boolean = false
+
+    override fun transformAnnotationCall(annotationCall: FirAnnotationCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        if (annotationArgumentsMode) {
+            return resolveAnnotationCall(annotationCall, data, FirAnnotationResolveStatus.PartiallyResolved)
+        }
+
+        annotationCall.transformAnnotationTypeRef(transformer, data)
+        val classId = annotationCall.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.lookupTag?.classId
+            ?: return annotationCall.compose()
+        if (classId.asSingleFqName() !in annotations) {
+            return annotationCall.compose()
+        }
+        annotationArgumentsMode = true
+        return resolveAnnotationCall(annotationCall, data, FirAnnotationResolveStatus.PartiallyResolved).also {
+            annotationArgumentsMode = false
+        }
+    }
+
+    override fun transformExpression(expression: FirExpression, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        return expression.compose()
+    }
+
+    override fun FirQualifiedAccessExpression.isAcceptableResolvedQualifiedAccess(): Boolean {
+        return calleeReference !is FirErrorNamedReference
+    }
+
+    override fun transformFunctionCall(functionCall: FirFunctionCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        return functionCall.compose()
+    }
+
+    override fun transformBlock(block: FirBlock, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        return block.compose()
+    }
+
+    override fun transformThisReceiverExpression(
+        thisReceiverExpression: FirThisReceiverExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return thisReceiverExpression.compose()
+    }
+
+    override fun transformComparisonExpression(
+        comparisonExpression: FirComparisonExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return comparisonExpression.compose()
+    }
+
+    override fun transformOperatorCall(operatorCall: FirOperatorCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+        return operatorCall.compose()
+    }
+
+    override fun transformTypeOperatorCall(
+        typeOperatorCall: FirTypeOperatorCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return typeOperatorCall.compose()
+    }
+
+    override fun transformCheckNotNullCall(
+        checkNotNullCall: FirCheckNotNullCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return checkNotNullCall.compose()
+    }
+
+    override fun transformBinaryLogicExpression(
+        binaryLogicExpression: FirBinaryLogicExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return binaryLogicExpression.compose()
+    }
+
+    override fun transformVariableAssignment(
+        variableAssignment: FirVariableAssignment,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return variableAssignment.compose()
+    }
+
+    override fun transformCallableReferenceAccess(
+        callableReferenceAccess: FirCallableReferenceAccess,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return callableReferenceAccess.compose()
+    }
+
+    override fun transformDelegatedConstructorCall(
+        delegatedConstructorCall: FirDelegatedConstructorCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return delegatedConstructorCall.compose()
+    }
+
+    override fun transformAugmentedArraySetCall(
+        augmentedArraySetCall: FirAugmentedArraySetCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
+        return augmentedArraySetCall.compose()
+    }
+}

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

@@ -44,7 +44,7 @@ class FirPluginAnnotationsResolveTransformer(
 
     override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirFile> {
         checkSessionConsistency(file)
-        if (!extensionService.hasExtensions) return file.compose()
+        if (!extensionService.hasPredicateBasedExtensions) return file.compose()
         val registeredPluginAnnotations = file.session.registeredPluginAnnotations
         file.replaceResolvePhase(FirResolvePhase.ANNOTATIONS_FOR_PLUGINS)
         val newAnnotations = file.resolveAnnotations(registeredPluginAnnotations.annotations, registeredPluginAnnotations.metaAnnotations)

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

@@ -7,16 +7,17 @@ package org.jetbrains.kotlin.fir.declarations
 
 enum class FirResolvePhase {
     RAW_FIR,
-    ANNOTATIONS_FOR_PLUGINS, // run only if some extensions are registered
+    ANNOTATIONS_FOR_PLUGINS, // plugin phase
     CLASS_GENERATION, // plugin phase
     IMPORTS,
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,
     TYPES,
-    EXTENSION_STATUS_UPDATE,
+    ARGUMENTS_OF_PLUGIN_ANNOTATIONS, // plugin phase
+    EXTENSION_STATUS_UPDATE, // plugin phase
     STATUS,
     CONTRACTS,
-    NEW_MEMBERS_GENERATION, // plugin
+    NEW_MEMBERS_GENERATION, // plugin phase
     IMPLICIT_TYPES_BODY_RESOLVE,
     BODY_RESOLVE;