FIR: add forgotten annotation calls resolve at some points

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/CopyUtils.kt

@@ -96,19 +96,22 @@ fun FirTypeParameter.copy(
 
 fun FirWhenExpression.copy(
     resultType: FirTypeRef = this.typeRef,
-    calleeReference: FirReference = this.calleeReference
+    calleeReference: FirReference = this.calleeReference,
+    annotations: List<FirAnnotationCall> = this.annotations
 ): FirWhenExpressionImpl = FirWhenExpressionImpl(source, subject, subjectVariable).apply {
     this.calleeReference = calleeReference
     this@apply.branches.addAll(this@copy.branches)
     this.typeRef = resultType
-    this.calleeReference = calleeReference
+    this.annotations += annotations
 }
 
 fun FirTryExpression.copy(
     resultType: FirTypeRef = this.typeRef,
-    calleeReference: FirReference = this.calleeReference
+    calleeReference: FirReference = this.calleeReference,
+    annotations: List<FirAnnotationCall> = this.annotations
 ): FirTryExpressionImpl = FirTryExpressionImpl(source, tryBlock, finallyBlock).apply {
     this.calleeReference = calleeReference
     this@apply.catches.addAll(this@copy.catches)
     this.typeRef = resultType
+    this.annotations += annotations
 }

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

@@ -52,6 +52,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         if (whenExpression.calleeReference is FirResolvedNamedReference && whenExpression.resultType !is FirImplicitTypeRef) {
             return whenExpression.compose()
         }
+        whenExpression.annotations.forEach { it.accept(this, data) }
         dataFlowAnalyzer.enterWhenExpression(whenExpression)
         return withScopeCleanup(localScopes) with@{
             if (whenExpression.subjectVariable != null) {
@@ -128,6 +129,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             return tryExpression.compose()
         }
 
+        tryExpression.annotations.forEach { it.accept(this, data) }
         dataFlowAnalyzer.enterTryExpression(tryExpression)
         tryExpression.transformTryBlock(transformer, ResolutionMode.ContextDependent)
         dataFlowAnalyzer.exitTryMainBlock(tryExpression)

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

@@ -53,7 +53,10 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
 
     override fun transformExpression(expression: FirExpression, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         if (expression.resultType is FirImplicitTypeRef && expression !is FirWrappedExpression) {
-            val type = FirErrorTypeRefImpl(expression.source, FirSimpleDiagnostic("Type calculating for ${expression::class} is not supported", DiagnosticKind.InferenceError))
+            val type = FirErrorTypeRefImpl(
+                expression.source,
+                FirSimpleDiagnostic("Type calculating for ${expression::class} is not supported", DiagnosticKind.InferenceError)
+            )
             expression.resultType = type
         }
         return (expression.transformChildren(transformer, data) as FirStatement).compose()
@@ -63,6 +66,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         qualifiedAccessExpression: FirQualifiedAccessExpression,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
+        qualifiedAccessExpression.annotations.forEach { it.accept(this, data) }
         var result = when (val callee = qualifiedAccessExpression.calleeReference) {
             is FirExplicitThisReference -> {
                 val labelName = callee.labelName
@@ -131,6 +135,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
             storeTypeFromCallee(functionCall)
         }
         if (functionCall.calleeReference !is FirSimpleNamedReference) return functionCall.compose()
+        functionCall.annotations.forEach { it.accept(this, data) }
         functionCall.transform<FirFunctionCall, Nothing?>(InvocationKindTransformer, null)
         functionCall.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
         val expectedTypeRef = data.expectedType
@@ -194,6 +199,8 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
 
         if (operatorCall.operation in FirOperation.ASSIGNMENTS) {
             require(operatorCall.operation != FirOperation.ASSIGN)
+
+            operatorCall.annotations.forEach { it.accept(this, data) }
             @Suppress("NAME_SHADOWING")
             val operatorCall = operatorCall.transformArguments(this, ResolutionMode.ContextIndependent)
             val (leftArgument, rightArgument) = operatorCall.arguments
@@ -245,7 +252,10 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         throw IllegalArgumentException(operatorCall.render())
     }
 
-    override fun transformTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformTypeOperatorCall(
+        typeOperatorCall: FirTypeOperatorCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         val symbolProvider = session.firSymbolProvider
         val resolved = transformExpression(typeOperatorCall, data).single
         when ((resolved as FirTypeOperatorCall).operation) {
@@ -295,6 +305,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
         // val resolvedAssignment = transformCallee(variableAssignment)
+        variableAssignment.annotations.forEach { it.accept(this, data) }
         val resolvedAssignment = callResolver.resolveVariableAccessAndSelectCandidate(variableAssignment, file)
         val result = if (resolvedAssignment is FirVariableAssignment) {
             val completeAssignment = callCompleter.completeCall(resolvedAssignment, noExpectedType) // TODO: check
@@ -317,6 +328,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
             return callableReferenceAccess.compose()
         }
 
+        callableReferenceAccess.annotations.forEach { it.accept(this, data) }
         val transformedLHS =
             callableReferenceAccess.explicitReceiver?.transformSingle(this, ResolutionMode.ContextIndependent)
 
@@ -384,6 +396,7 @@ class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransformer) :
     override fun <T> transformConstExpression(constExpression: FirConstExpression<T>, data: ResolutionMode): CompositeTransformResult<FirStatement> {
         // TODO: add support of IntegerLiteralType
 
+        constExpression.annotations.forEach { it.accept(this, data) }
         val kind = constExpression.kind
         val symbol = when (kind) {
             IrConstKind.Null -> StandardClassIds.Nothing(symbolProvider)

compiler/fir/resolve/testData/resolve/expresssions/annotations.kt

@@ -0,0 +1,19 @@
+annotation class MyAnn
+
+fun bar(x: Int) {}
+
+fun foo() {
+    @MyAnn
+    val x: Int
+    @MyAnn
+    x = @MyAnn 42
+    @MyAnn
+    bar(@MyAnn x)
+
+    val y = @MyAnn if (false) x else x
+    val z = @MyAnn try {
+        x
+    } catch (t: Throwable) {
+        0
+    }
+}

compiler/fir/resolve/testData/resolve/expresssions/annotations.txt

@@ -0,0 +1,30 @@
+FILE: annotations.kt
+    public final annotation class MyAnn : R|kotlin/Annotation| {
+        public constructor(): R|MyAnn| {
+            super<R|kotlin/Annotation|>()
+        }
+
+    }
+    public final fun bar(x: R|kotlin/Int|): R|kotlin/Unit| {
+    }
+    public final fun foo(): R|kotlin/Unit| {
+        @R|MyAnn|() lval x: R|kotlin/Int|
+        @R|MyAnn|() R|<local>/x| = @R|MyAnn|() Int(42)
+        @R|MyAnn|() R|/bar|(@R|MyAnn|() R|<local>/x|)
+        lval y: R|kotlin/Int| = @R|MyAnn|() when () {
+            Boolean(false) ->  {
+                R|<local>/x|
+            }
+            else ->  {
+                R|<local>/x|
+            }
+        }
+
+        lval z: R|kotlin/Int| = @R|MyAnn|() try {
+            R|<local>/x|
+        }
+        catch (t: R|kotlin/Throwable|) {
+            Int(0)
+        }
+
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -364,6 +364,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/fir/resolve/testData/resolve/expresssions"), Pattern.compile("^([^.]+)\\.kt$"), true);
         }
 
+        @TestMetadata("annotations.kt")
+        public void testAnnotations() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/expresssions/annotations.kt");
+        }
+
         @TestMetadata("checkArguments.kt")
         public void testCheckArguments() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/expresssions/checkArguments.kt");