[FIR][tree] introduce lazy expression to simplify code

^KT-56543

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -14,9 +14,7 @@ import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyDelegatedConstructorCall
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
 import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.CompilerRequiredAnnotationsHelper
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyBodiesCalculatorTest.kt

@@ -13,8 +13,8 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.Analys
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.renderer.FirRenderer
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.references.FirResolvedErrorReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.*

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -7,6 +7,9 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.PrivateForInline
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 
@@ -35,4 +38,12 @@ abstract class FirPartialBodyResolveTransformer(
     override fun <E : FirElement> transformElement(element: E, data: ResolutionMode): E {
         return element.transform(transformer, data)
     }
+
+    override fun transformLazyExpression(lazyExpression: FirLazyExpression, data: ResolutionMode): FirStatement {
+        error("FirLazyExpression should be calculated before accessing")
+    }
+
+    override fun transformLazyBlock(lazyBlock: FirLazyBlock, data: ResolutionMode): FirStatement {
+        error("FirLazyBlock should be calculated before accessing")
+    }
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirLazyBlock.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2023 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.expressions
+
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.visitors.*
+
+/*
+ * This file was generated automatically
+ * DO NOT MODIFY IT MANUALLY
+ */
+
+abstract class FirLazyBlock : FirBlock() {
+    abstract override val source: KtSourceElement?
+    abstract override val annotations: List<FirAnnotation>
+    abstract override val statements: List<FirStatement>
+    abstract override val typeRef: FirTypeRef
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitLazyBlock(this, data)
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <E : FirElement, D> transform(transformer: FirTransformer<D>, data: D): E =
+        transformer.transformLazyBlock(this, data) as E
+
+    abstract override fun replaceAnnotations(newAnnotations: List<FirAnnotation>)
+
+    abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
+
+    abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLazyBlock
+
+    abstract override fun <D> transformStatements(transformer: FirTransformer<D>, data: D): FirLazyBlock
+
+    abstract override fun <D> transformOtherChildren(transformer: FirTransformer<D>, data: D): FirLazyBlock
+}

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirLazyExpression.kt

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2023 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.expressions
+
+import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.visitors.*
+
+/*
+ * This file was generated automatically
+ * DO NOT MODIFY IT MANUALLY
+ */
+
+abstract class FirLazyExpression : FirExpression() {
+    abstract override val source: KtSourceElement?
+    abstract override val typeRef: FirTypeRef
+    abstract override val annotations: List<FirAnnotation>
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitLazyExpression(this, data)
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <E : FirElement, D> transform(transformer: FirTransformer<D>, data: D): E =
+        transformer.transformLazyExpression(this, data) as E
+
+    abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
+
+    abstract override fun replaceAnnotations(newAnnotations: List<FirAnnotation>)
+
+    abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLazyExpression
+}

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirLazyBlockBuilder.kt

@@ -9,15 +9,14 @@ package org.jetbrains.kotlin.fir.expressions.builder
 
 import kotlin.contracts.*
 import org.jetbrains.kotlin.KtSourceElement
-import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlockImpl
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 
@@ -26,7 +25,6 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-@OptIn(FirImplementationDetail::class)
-fun buildLazyBlock(): FirBlock {
-    return FirLazyBlock()
+fun buildLazyBlock(): FirLazyBlock {
+    return FirLazyBlockImpl()
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirLazyExpressionBuilder.kt

@@ -9,13 +9,13 @@ package org.jetbrains.kotlin.fir.expressions.builder
 
 import kotlin.contracts.*
 import org.jetbrains.kotlin.KtSourceElement
-import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
 import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirExpression
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
+import org.jetbrains.kotlin.fir.expressions.builder.FirExpressionBuilder
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpressionImpl
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 
@@ -25,23 +25,29 @@ import org.jetbrains.kotlin.fir.visitors.*
  */
 
 @FirBuilderDsl
-class FirLazyExpressionBuilder : FirAnnotationContainerBuilder {
+class FirLazyExpressionBuilder : FirAnnotationContainerBuilder, FirExpressionBuilder {
     override var source: KtSourceElement? = null
 
-    @OptIn(FirImplementationDetail::class)
-    override fun build(): FirExpression {
-        return FirLazyExpression(
+    override fun build(): FirLazyExpression {
+        return FirLazyExpressionImpl(
             source,
         )
     }
 
 
+    @Deprecated("Modification of 'typeRef' has no impact for FirLazyExpressionBuilder", level = DeprecationLevel.HIDDEN)
+    override var typeRef: FirTypeRef
+        get() = throw IllegalStateException()
+        set(_) {
+            throw IllegalStateException()
+        }
+
     @Deprecated("Modification of 'annotations' has no impact for FirLazyExpressionBuilder", level = DeprecationLevel.HIDDEN)
     override val annotations: MutableList<FirAnnotation> = mutableListOf()
 }
 
 @OptIn(ExperimentalContracts::class)
-inline fun buildLazyExpression(init: FirLazyExpressionBuilder.() -> Unit = {}): FirExpression {
+inline fun buildLazyExpression(init: FirLazyExpressionBuilder.() -> Unit = {}): FirLazyExpression {
     contract {
         callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirLazyBlockImpl.kt

@@ -8,9 +8,8 @@
 package org.jetbrains.kotlin.fir.expressions.impl
 
 import org.jetbrains.kotlin.KtSourceElement
-import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
@@ -22,7 +21,7 @@ import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
  * DO NOT MODIFY IT MANUALLY
  */
 
-class FirLazyBlock : FirBlock() {
+internal class FirLazyBlockImpl : FirLazyBlock() {
     override val source: KtSourceElement? get() = error("FirLazyBlock should be calculated before accessing")
     override val annotations: List<FirAnnotation> get() = error("FirLazyBlock should be calculated before accessing")
     override val statements: List<FirStatement> get() = error("FirLazyBlock should be calculated before accessing")
@@ -31,20 +30,20 @@ class FirLazyBlock : FirBlock() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
     }
 
-    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirLazyBlock {
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirLazyBlockImpl {
         transformOtherChildren(transformer, data)
         return this
     }
 
-    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLazyBlock {
+    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLazyBlockImpl {
         return this
     }
 
-    override fun <D> transformStatements(transformer: FirTransformer<D>, data: D): FirLazyBlock {
+    override fun <D> transformStatements(transformer: FirTransformer<D>, data: D): FirLazyBlockImpl {
         return this
     }
 
-    override fun <D> transformOtherChildren(transformer: FirTransformer<D>, data: D): FirLazyBlock {
+    override fun <D> transformOtherChildren(transformer: FirTransformer<D>, data: D): FirLazyBlockImpl {
         return this
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirLazyExpressionImpl.kt

@@ -8,9 +8,8 @@
 package org.jetbrains.kotlin.fir.expressions.impl
 
 import org.jetbrains.kotlin.KtSourceElement
-import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.fir.MutableOrEmptyList
@@ -21,20 +20,20 @@ import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
  * DO NOT MODIFY IT MANUALLY
  */
 
-class FirLazyExpression @FirImplementationDetail constructor(
+internal class FirLazyExpressionImpl(
     override val source: KtSourceElement?,
-) : FirExpression() {
+) : FirLazyExpression() {
     override val typeRef: FirTypeRef get() = error("FirLazyExpression should be calculated before accessing")
     override val annotations: List<FirAnnotation> get() = error("FirLazyExpression should be calculated before accessing")
 
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
     }
 
-    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirLazyExpression {
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirLazyExpressionImpl {
         return this
     }
 
-    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLazyExpression {
+    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLazyExpressionImpl {
         return this
     }
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirDefaultVisitor.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirControlFlowGraphOwner
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.declarations.FirContextReceiver
 import org.jetbrains.kotlin.fir.FirElementWithResolvePhase
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
@@ -61,6 +62,7 @@ import org.jetbrains.kotlin.fir.expressions.FirErrorLoop
 import org.jetbrains.kotlin.fir.expressions.FirDoWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirBinaryLogicExpression
 import org.jetbrains.kotlin.fir.expressions.FirJump
 import org.jetbrains.kotlin.fir.expressions.FirLoopJump
@@ -166,6 +168,8 @@ abstract class FirDefaultVisitor<out R, in D> : FirVisitor<R, D>() {
 
     override fun visitExpression(expression: FirExpression, data: D): R  = visitStatement(expression, data)
 
+    override fun visitLazyExpression(lazyExpression: FirLazyExpression, data: D): R  = visitExpression(lazyExpression, data)
+
     override fun visitTypeParametersOwner(typeParametersOwner: FirTypeParametersOwner, data: D): R  = visitTypeParameterRefsOwner(typeParametersOwner, data)
 
     override fun visitCallableDeclaration(callableDeclaration: FirCallableDeclaration, data: D): R  = visitMemberDeclaration(callableDeclaration, data)
@@ -190,6 +194,8 @@ abstract class FirDefaultVisitor<out R, in D> : FirVisitor<R, D>() {
 
     override fun visitBlock(block: FirBlock, data: D): R  = visitExpression(block, data)
 
+    override fun visitLazyBlock(lazyBlock: FirLazyBlock, data: D): R  = visitBlock(lazyBlock, data)
+
     override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: D): R  = visitExpression(binaryLogicExpression, data)
 
     override fun <E : FirTargetElement> visitJump(jump: FirJump<E>, data: D): R  = visitExpression(jump, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirDefaultVisitorVoid.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirControlFlowGraphOwner
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.declarations.FirContextReceiver
 import org.jetbrains.kotlin.fir.FirElementWithResolvePhase
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
@@ -61,6 +62,7 @@ import org.jetbrains.kotlin.fir.expressions.FirErrorLoop
 import org.jetbrains.kotlin.fir.expressions.FirDoWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirBinaryLogicExpression
 import org.jetbrains.kotlin.fir.expressions.FirJump
 import org.jetbrains.kotlin.fir.expressions.FirLoopJump
@@ -166,6 +168,8 @@ abstract class FirDefaultVisitorVoid : FirVisitorVoid() {
 
     override fun visitExpression(expression: FirExpression)  = visitStatement(expression)
 
+    override fun visitLazyExpression(lazyExpression: FirLazyExpression)  = visitExpression(lazyExpression)
+
     override fun visitTypeParametersOwner(typeParametersOwner: FirTypeParametersOwner)  = visitTypeParameterRefsOwner(typeParametersOwner)
 
     override fun visitCallableDeclaration(callableDeclaration: FirCallableDeclaration)  = visitMemberDeclaration(callableDeclaration)
@@ -190,6 +194,8 @@ abstract class FirDefaultVisitorVoid : FirVisitorVoid() {
 
     override fun visitBlock(block: FirBlock)  = visitExpression(block)
 
+    override fun visitLazyBlock(lazyBlock: FirLazyBlock)  = visitBlock(lazyBlock)
+
     override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression)  = visitExpression(binaryLogicExpression)
 
     override fun <E : FirTargetElement> visitJump(jump: FirJump<E>)  = visitExpression(jump)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirTransformer.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirControlFlowGraphOwner
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.declarations.FirContextReceiver
 import org.jetbrains.kotlin.fir.FirElementWithResolvePhase
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
@@ -61,6 +62,7 @@ import org.jetbrains.kotlin.fir.expressions.FirErrorLoop
 import org.jetbrains.kotlin.fir.expressions.FirDoWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirBinaryLogicExpression
 import org.jetbrains.kotlin.fir.expressions.FirJump
 import org.jetbrains.kotlin.fir.expressions.FirLoopJump
@@ -205,6 +207,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformElement(expression, data)
     }
 
+    open fun transformLazyExpression(lazyExpression: FirLazyExpression, data: D): FirStatement {
+        return transformElement(lazyExpression, data)
+    }
+
     open fun transformContextReceiver(contextReceiver: FirContextReceiver, data: D): FirContextReceiver {
         return transformElement(contextReceiver, data)
     }
@@ -381,6 +387,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformElement(block, data)
     }
 
+    open fun transformLazyBlock(lazyBlock: FirLazyBlock, data: D): FirStatement {
+        return transformElement(lazyBlock, data)
+    }
+
     open fun transformBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: D): FirStatement {
         return transformElement(binaryLogicExpression, data)
     }
@@ -789,6 +799,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformExpression(expression, data)
     }
 
+    final override fun visitLazyExpression(lazyExpression: FirLazyExpression, data: D): FirStatement {
+        return transformLazyExpression(lazyExpression, data)
+    }
+
     final override fun visitContextReceiver(contextReceiver: FirContextReceiver, data: D): FirContextReceiver {
         return transformContextReceiver(contextReceiver, data)
     }
@@ -965,6 +979,10 @@ abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {
         return transformBlock(block, data)
     }
 
+    final override fun visitLazyBlock(lazyBlock: FirLazyBlock, data: D): FirStatement {
+        return transformLazyBlock(lazyBlock, data)
+    }
+
     final override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: D): FirStatement {
         return transformBinaryLogicExpression(binaryLogicExpression, data)
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirVisitor.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirControlFlowGraphOwner
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.declarations.FirContextReceiver
 import org.jetbrains.kotlin.fir.FirElementWithResolvePhase
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
@@ -61,6 +62,7 @@ import org.jetbrains.kotlin.fir.expressions.FirErrorLoop
 import org.jetbrains.kotlin.fir.expressions.FirDoWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirBinaryLogicExpression
 import org.jetbrains.kotlin.fir.expressions.FirJump
 import org.jetbrains.kotlin.fir.expressions.FirLoopJump
@@ -182,6 +184,8 @@ abstract class FirVisitor<out R, in D> {
 
     open fun visitExpression(expression: FirExpression, data: D): R  = visitElement(expression, data)
 
+    open fun visitLazyExpression(lazyExpression: FirLazyExpression, data: D): R  = visitElement(lazyExpression, data)
+
     open fun visitContextReceiver(contextReceiver: FirContextReceiver, data: D): R  = visitElement(contextReceiver, data)
 
     open fun visitElementWithResolvePhase(elementWithResolvePhase: FirElementWithResolvePhase, data: D): R  = visitElement(elementWithResolvePhase, data)
@@ -270,6 +274,8 @@ abstract class FirVisitor<out R, in D> {
 
     open fun visitBlock(block: FirBlock, data: D): R  = visitElement(block, data)
 
+    open fun visitLazyBlock(lazyBlock: FirLazyBlock, data: D): R  = visitElement(lazyBlock, data)
+
     open fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: D): R  = visitElement(binaryLogicExpression, data)
 
     open fun <E : FirTargetElement> visitJump(jump: FirJump<E>, data: D): R  = visitElement(jump, data)

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/visitors/FirVisitorVoid.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvedDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirControlFlowGraphOwner
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.declarations.FirContextReceiver
 import org.jetbrains.kotlin.fir.FirElementWithResolvePhase
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
@@ -61,6 +62,7 @@ import org.jetbrains.kotlin.fir.expressions.FirErrorLoop
 import org.jetbrains.kotlin.fir.expressions.FirDoWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirWhileLoop
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirBinaryLogicExpression
 import org.jetbrains.kotlin.fir.expressions.FirJump
 import org.jetbrains.kotlin.fir.expressions.FirLoopJump
@@ -204,6 +206,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitElement(expression)
     }
 
+    open fun visitLazyExpression(lazyExpression: FirLazyExpression) {
+        visitElement(lazyExpression)
+    }
+
     open fun visitContextReceiver(contextReceiver: FirContextReceiver) {
         visitElement(contextReceiver)
     }
@@ -380,6 +386,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitElement(block)
     }
 
+    open fun visitLazyBlock(lazyBlock: FirLazyBlock) {
+        visitElement(lazyBlock)
+    }
+
     open fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression) {
         visitElement(binaryLogicExpression)
     }
@@ -788,6 +798,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitExpression(expression)
     }
 
+    final override fun visitLazyExpression(lazyExpression: FirLazyExpression, data: Nothing?) {
+        visitLazyExpression(lazyExpression)
+    }
+
     final override fun visitContextReceiver(contextReceiver: FirContextReceiver, data: Nothing?) {
         visitContextReceiver(contextReceiver)
     }
@@ -964,6 +978,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitBlock(block)
     }
 
+    final override fun visitLazyBlock(lazyBlock: FirLazyBlock, data: Nothing?) {
+        visitLazyBlock(lazyBlock)
+    }
+
     final override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: Nothing?) {
         visitBinaryLogicExpression(binaryLogicExpression)
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/renderer/FirBodyRenderer.kt

@@ -8,8 +8,8 @@ package org.jetbrains.kotlin.fir.renderer
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirDelegatedConstructorCall
+import org.jetbrains.kotlin.fir.expressions.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 
 class FirBodyRenderer {
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/renderer/FirRenderer.kt

@@ -595,13 +595,10 @@ class FirRenderer(
         }
 
         override fun visitExpression(expression: FirExpression) {
-            if (expression !is FirLazyExpression) {
-                annotationRenderer?.render(expression)
-            }
+            annotationRenderer?.render(expression)
             print(
                 when (expression) {
                     is FirExpressionStub -> "STUB"
-                    is FirLazyExpression -> "LAZY_EXPRESSION"
                     is FirUnitExpression -> "Unit"
                     is FirElseIfTrueCondition -> "else"
                     is FirNoReceiverExpression -> ""
@@ -610,6 +607,10 @@ class FirRenderer(
             )
         }
 
+        override fun visitLazyExpression(lazyExpression: FirLazyExpression) {
+            print("LAZY_EXPRESSION")
+        }
+
         override fun <T> visitConstExpression(constExpression: FirConstExpression<T>) {
             annotationRenderer?.render(constExpression)
             val kind = constExpression.kind

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/FirTreeBuilder.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -26,6 +26,7 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
 
     val statement by element(Expression, annotationContainer)
     val expression by element(Expression, statement)
+    val lazyExpression by element(Expression, expression)
 
     val contextReceiver by element(Declaration)
 
@@ -83,6 +84,8 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
     val whileLoop by element(Expression, loop)
 
     val block by element(Expression, expression)
+    val lazyBlock by element(Expression, block)
+
     val binaryLogicExpression by element(Expression, expression)
     val jump by sealedElement(Expression, expression)
     val loopJump by element(Expression, jump)

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/ImplementationConfigurator.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -173,7 +173,7 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
             publicImplementation()
         }
 
-        impl(block, "FirLazyBlock") {
+        impl(lazyBlock) {
             val error = """error("FirLazyBlock should be calculated before accessing")"""
             default("source") {
                 value = error
@@ -191,7 +191,6 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
                 value = error
                 withGetter = true
             }
-            publicImplementation()
         }
 
         impl(errorLoop) {
@@ -204,7 +203,7 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
             publicImplementation()
         }
 
-        impl(expression, "FirLazyExpression") {
+        impl(lazyExpression) {
             val error = """error("FirLazyExpression should be calculated before accessing")"""
             default("typeRef") {
                 value = error
@@ -214,7 +213,6 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
                 value = error
                 withGetter = true
             }
-            publicImplementation()
         }
 
         impl(functionCall) {