From 1ac4d684337da9148923129062e2365e5f39da8e Mon Sep 17 00:00:00 2001
From: Ivan Kochurkin <ivan.kochurkin@jetbrains.com>
Date: Thu, 2 Jun 2022 15:21:58 +0300
Subject: [PATCH] [FIR] Initialize annotations for block for LightTree

---
 .../fir/lightTree/converter/ExpressionsConverter.kt | 10 +++++++++-
 .../jetbrains/kotlin/fir/builder/RawFirBuilder.kt   | 13 +++++++++++--
 .../diagnostics/tests/annotations/AnnotatedBlock.kt |  6 +++++-
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt
index 0cabbe17c57..f5d0f02289e 100644
--- a/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt
+++ b/compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt
@@ -435,6 +435,7 @@ class ExpressionsConverter(
             when (it.tokenType) {
                 ANNOTATION -> firAnnotationList += declarationsConverter.convertAnnotation(it)
                 ANNOTATION_ENTRY -> firAnnotationList += declarationsConverter.convertAnnotationEntry(it)
+                BLOCK -> firExpression = declarationsConverter.convertBlockExpression(it)
                 else -> if (it.isExpression()) {
                     context.forwardLabelUsagePermission(annotatedExpression, it)
                     firExpression = getAsFirExpression(it)
@@ -442,7 +443,7 @@ class ExpressionsConverter(
             }
         }
 
-        val result = firExpression ?: buildErrorExpression(null, ConeNotAnnotationContainer(firExpression?.render() ?: "???"))
+        val result = firExpression ?: buildErrorExpression(null, ConeNotAnnotationContainer("???"))
         require(result is FirAnnotationContainer)
         (result.annotations as MutableList<FirAnnotation>) += firAnnotationList
         return result
@@ -1149,6 +1150,13 @@ class ExpressionsConverter(
         body?.forEachChildren {
             when (it.tokenType) {
                 BLOCK -> firBlock = declarationsConverter.convertBlockExpression(it)
+                ANNOTATED_EXPRESSION -> {
+                    if (it.getChildNodeByType(BLOCK) != null) {
+                        firBlock = getAsFirExpression(it)
+                    } else {
+                        firStatement = getAsFirExpression(it)
+                    }
+                }
                 else -> if (it.isExpression()) firStatement = getAsFirExpression(it)
             }
         }
diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
index 943d468f2e5..c086bff86d4 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt
@@ -295,8 +295,17 @@ open class RawFirBuilder(
                     accept(this@Visitor, Unit) as FirBlock
                 null ->
                     buildEmptyExpressionBlock()
-                else ->
-                    FirSingleExpressionBlock(convert())
+                else -> {
+                    var firBlock: FirBlock? = null
+                    if (this is KtAnnotatedExpression) {
+                        val lastChild = children.lastOrNull()
+                        if (lastChild is KtBlockExpression) {
+                            firBlock = lastChild.toFirBlock()
+                            extractAnnotationsTo(firBlock.annotations as MutableList<FirAnnotation>)
+                        }
+                    }
+                    firBlock ?: FirSingleExpressionBlock(convert())
+                }
             }
 
         private fun KtDeclarationWithBody.buildFirBody(): Pair<FirBlock?, FirContractDescription?> =
diff --git a/compiler/testData/diagnostics/tests/annotations/AnnotatedBlock.kt b/compiler/testData/diagnostics/tests/annotations/AnnotatedBlock.kt
index 7b30f955d2e..b7b5c9ee9bb 100644
--- a/compiler/testData/diagnostics/tests/annotations/AnnotatedBlock.kt
+++ b/compiler/testData/diagnostics/tests/annotations/AnnotatedBlock.kt
@@ -3,8 +3,12 @@
 // ISSUE: KT-52175
 
 annotation class Ann
+annotation class Ann2
 
 fun test(x: String?) {
     if (x != null)
-        @Ann() { Unit }
+        @Ann() @Ann2() { Unit } // It should be Block with annotations
+
+    if (x != null)
+        @Ann() @Ann2() Unit // It should be SingleExpressionBlock  with annotations
 }
\ No newline at end of file