From 4b43e72440dea24b75803bccb09db7183cdbd6b2 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Thu, 13 Jul 2023 18:20:59 +0900
Subject: [PATCH] [FIR] Create 'FirCodeFragment' with a lazy block if requested

---
 .../jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt   | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt
index da822651793..1267dc4771d 100644
--- a/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt
+++ b/compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt
@@ -1197,10 +1197,12 @@ open class PsiRawFirBuilder(
                 moduleData = baseModuleData
                 origin = FirDeclarationOrigin.Source
                 symbol = FirCodeFragmentSymbol()
-                block = when (file) {
-                    is KtExpressionCodeFragment -> file.getContentElement()?.toFirBlock() ?: buildEmptyExpressionBlock()
-                    is KtBlockCodeFragment -> configureBlockWithoutBuilding(file.getContentElement()).build()
-                    else -> error("Unexpected code fragment type: ${file::class}")
+                block = buildOrLazyBlock {
+                    when (file) {
+                        is KtExpressionCodeFragment -> file.getContentElement()?.toFirBlock() ?: buildEmptyExpressionBlock()
+                        is KtBlockCodeFragment -> configureBlockWithoutBuilding(file.getContentElement()).build()
+                        else -> error("Unexpected code fragment type: ${file::class}")
+                    }
                 }
             }
         }