[FIR] Make anonymous initializer lazy in RawFirBuilder

Fourth step for KT-52615

Merge-request: KT-MR-7798
Merged-by: Egor Kulikov <Egor.Kulikov@jetbrains.com>

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

@@ -21,10 +21,7 @@ import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
-import org.jetbrains.kotlin.psi.KtEnumEntry
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 
 internal object FirLazyBodiesCalculator {
     fun calculateLazyBodiesInside(designation: FirDeclarationDesignation) {
@@ -133,6 +130,20 @@ internal object FirLazyBodiesCalculator {
         }
     }
 
+    fun calculateLazyBodyForAnonymousInitializer(designation: FirDeclarationDesignation) {
+        val initializer = designation.declaration as FirAnonymousInitializer
+        require(initializer.body is FirLazyBlock)
+        val newInitializer = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
+            session = initializer.moduleData.session,
+            scopeProvider = initializer.moduleData.session.kotlinScopeProvider,
+            designation = designation,
+            rootNonLocalDeclaration = initializer.psi as KtAnonymousInitializer,
+        ) as FirAnonymousInitializer
+        initializer.apply {
+            replaceBody(newInitializer.body)
+        }
+    }
+
     fun needCalculatingLazyBodyForProperty(firProperty: FirProperty): Boolean =
         firProperty.getter?.body is FirLazyBlock
                 || firProperty.setter?.body is FirLazyBlock
@@ -202,4 +213,14 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
         }
         return enumEntry
     }
+
+    override fun transformAnonymousInitializer(
+        anonymousInitializer: FirAnonymousInitializer, data: PersistentList<FirDeclaration>
+    ): FirAnonymousInitializer {
+        if (anonymousInitializer.body is FirLazyBlock) {
+            val designation = FirDeclarationDesignation(data, anonymousInitializer)
+            FirLazyBodiesCalculator.calculateLazyBodyForAnonymousInitializer(designation)
+        }
+        return anonymousInitializer
+    }
 }

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

@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.psi.psiUtil.hasExpectModifier
 internal class RawFirNonLocalDeclarationBuilder private constructor(
     session: FirSession,
     baseScopeProvider: FirScopeProvider,
-    private val originalDeclaration: FirTypeParameterRefsOwner,
+    private val originalDeclarationIfParamertized: FirTypeParameterRefsOwner?,
     private val declarationToBuild: KtDeclaration,
     private val functionsToRebind: Set<FirFunction>? = null,
     private val replacementApplier: RawFirReplacement.Applier? = null
@@ -45,7 +45,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             val builder = RawFirNonLocalDeclarationBuilder(
                 session = session,
                 baseScopeProvider = scopeProvider,
-                originalDeclaration = designation.declaration as FirTypeParameterRefsOwner,
+                originalDeclarationIfParamertized = designation.declaration as? FirTypeParameterRefsOwner,
                 declarationToBuild = rootNonLocalDeclaration,
                 replacementApplier = replacementApplier
             )
@@ -70,7 +70,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             val builder = RawFirNonLocalDeclarationBuilder(
                 session = session,
                 baseScopeProvider = scopeProvider,
-                originalDeclaration = designation.declaration as FirTypeParameterRefsOwner,
+                originalDeclarationIfParamertized = designation.declaration as? FirTypeParameterRefsOwner,
                 declarationToBuild = rootNonLocalDeclaration,
                 functionsToRebind = functionsToRebind,
             )
@@ -89,8 +89,8 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
         declarationSource: KtSourceElement?,
         currentFirTypeParameters: List<FirTypeParameterRef>
     ) {
-        if (declarationSource?.psi == originalDeclaration.psi) {
+        if (originalDeclarationIfParamertized != null && declarationSource?.psi == originalDeclarationIfParamertized.psi) {
-            super.addCapturedTypeParameters(status, declarationSource, originalDeclaration.typeParameters)
+            super.addCapturedTypeParameters(status, declarationSource, originalDeclarationIfParamertized.typeParameters)
         } else {
             super.addCapturedTypeParameters(status, declarationSource, currentFirTypeParameters)
         }

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1083,12 +1083,6 @@ open class RawFirBuilder(
             }
         }
 
-        override fun visitClassInitializer(initializer: KtClassInitializer, data: Unit?): FirElement {
-            return disabledLazyMode {
-                super.visitClassInitializer(initializer, data)
-            }
-        }
-
         private fun convertContextReceivers(receivers: List<KtContextReceiver>): List<FirContextReceiver> {
             return receivers.map { contextReceiverElement ->
                 buildContextReceiver {
@@ -1746,7 +1740,7 @@ open class RawFirBuilder(
                 source = initializer.toFirSourceElement()
                 moduleData = baseModuleData
                 origin = FirDeclarationOrigin.Source
-                body = initializer.body.toFirBlock()
+                body = buildOrLazyBlock { initializer.body.toFirBlock() }
             }
         }
 

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorInObject.lazyBodies.txt

@@ -4,8 +4,7 @@ FILE: constructorInObject.kt
             super<R|kotlin/Any|>()
         }
 
-        init {
+        init { LAZY_BLOCK }
-        }
 
     }
     public? final? enum class B : R|kotlin/Enum<B>| {

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/enums3.lazyBodies.txt

@@ -4,77 +4,6 @@ FILE: enums3.kt
             super<R|kotlin/Any|>()
         }
 
-        init {
+        init { LAZY_BLOCK }
-            local final? enum class Planet : R|kotlin/Enum<C.Planet>| {
-                private constructor(m: Double, r: Double): R|C.Planet| {
-                    super<R|kotlin/Enum<C.Planet>|>()
-                }
-
-                public? final? val m: Double = R|<local>/m|
-                    public? get(): Double
-
-                internal final? val r: Double = R|<local>/r|
-                    internal get(): Double
-
-                public final static enum entry MERCURY: R|C.Planet| = object : R|C.Planet| {
-                    private constructor(): R|<anonymous>| {
-                        super<R|C.Planet|>(Double(1.0), Double(2.0))
-                    }
-
-                    public? open? override fun sayHello(): R|kotlin/Unit| {
-                        println#(String(Hello!!!))
-                    }
-
-                }
-
-                public final static enum entry VENERA: R|C.Planet| = object : R|C.Planet| {
-                    private constructor(): R|<anonymous>| {
-                        super<R|C.Planet|>(Double(3.0), Double(4.0))
-                    }
-
-                    public? open? override fun sayHello(): R|kotlin/Unit| {
-                        println#(String(Ola!!!))
-                    }
-
-                }
-
-                public final static enum entry EARTH: R|C.Planet| = object : R|C.Planet| {
-                    private constructor(): R|<anonymous>| {
-                        super<R|C.Planet|>(Double(5.0), Double(6.0))
-                    }
-
-                    public? open? override fun sayHello(): R|kotlin/Unit| {
-                        println#(String(Privet!!!))
-                    }
-
-                }
-
-                public? final? val g: Double = G#.times#(m#).div#(r#.times#(r#))
-                    public? get(): Double
-
-                public? abstract fun sayHello(): R|kotlin/Unit|
-
-                local final? companion object Companion : R|kotlin/Any| {
-                    private constructor(): R|C.Planet.Companion| {
-                        super<R|kotlin/Any|>()
-                    }
-
-                    public? final? const val G: <implicit> = Double(6.67E-11)
-                        public? get(): <implicit>
-
-                }
-
-                public final static fun values(): R|kotlin/Array<C.Planet>| {
-                }
-
-                public final static fun valueOf(value: R|kotlin/String|): R|C.Planet| {
-                }
-
-                public final static val entries: R|kotlin/enums/EnumEntries<C.Planet>|
-                    public get(): R|kotlin/enums/EnumEntries<C.Planet>|
-
-            }
-
-        }
 
     }

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.lazyBodies.txt

@@ -4,21 +4,7 @@ FILE: initBlockWithDeclarations.kt
             super<R|kotlin/Any|>()
         }
 
-        init {
+        init { LAZY_BLOCK }
-            local final? class classInInit : R|kotlin/Any| {
-                public? constructor(): R|X.classInInit| {
-                    super<R|kotlin/Any|>()
-                }
-
-                public? final? fun funInClassInInit(): R|kotlin/Unit| {
-                }
-
-            }
-
-            local final? fun funInInit(): R|kotlin/Unit| {
-            }
-
-        }
 
     }
     public? final? object <no name provided> : R|kotlin/Any| {
@@ -26,20 +12,6 @@ FILE: initBlockWithDeclarations.kt
             super<R|kotlin/Any|>()
         }
 
-        init {
+        init { LAZY_BLOCK }
-            local final? class classInInit : R|kotlin/Any| {
-                public? constructor(): R|<no name provided>.classInInit| {
-                    super<R|kotlin/Any|>()
-                }
-
-                public? final? fun funInClassInInit(): R|kotlin/Unit| {
-                }
-
-            }
-
-            local final? fun funInInit(): R|kotlin/Unit| {
-            }
-
-        }
 
     }

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/init.lazyBodies.txt

@@ -7,8 +7,6 @@ FILE: init.kt
         public? final? val x: Int
             public? get(): Int
 
-        init {
+        init { LAZY_BLOCK }
-            this#.x# = x#
-        }
 
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirAnonymousInitializer.kt

@@ -40,5 +40,7 @@ abstract class FirAnonymousInitializer : FirDeclaration(), FirControlFlowGraphOw
 
     abstract override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
 
+    abstract fun replaceBody(newBody: FirBlock?)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirAnonymousInitializer
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirAnonymousInitializerImpl.kt

@@ -63,4 +63,8 @@ internal class FirAnonymousInitializerImpl(
     override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
         controlFlowGraphReference = newControlFlowGraphReference
     }
+
+    override fun replaceBody(newBody: FirBlock?) {
+        body = newBody
+    }
 }

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

@@ -437,7 +437,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         anonymousInitializer.configure {
-            +body(nullable = true)
+            +body(nullable = true, withReplace = true)
             +symbol("FirAnonymousInitializerSymbol")
         }