[FIR] Make enum entry initializer lazy in RawFirBuilder

Third step for KT-52615

Merge-request: KT-MR-7769
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,6 +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.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor
@@ -118,6 +119,20 @@ internal object FirLazyBodiesCalculator {
         }
     }
 
+    fun calculateLazyInitializerForEnumEntry(designation: FirDeclarationDesignation) {
+        val enumEntry = designation.declaration as FirEnumEntry
+        require(enumEntry.initializer is FirLazyExpression)
+        val newEntry = RawFirNonLocalDeclarationBuilder.buildWithFunctionSymbolRebind(
+            session = enumEntry.moduleData.session,
+            scopeProvider = enumEntry.moduleData.session.kotlinScopeProvider,
+            designation = designation,
+            rootNonLocalDeclaration = enumEntry.psi as KtEnumEntry,
+        ) as FirEnumEntry
+        enumEntry.apply {
+            replaceInitializer(newEntry.initializer)
+        }
+    }
+
     fun needCalculatingLazyBodyForProperty(firProperty: FirProperty): Boolean =
         firProperty.getter?.body is FirLazyBlock
                 || firProperty.setter?.body is FirLazyBlock
@@ -179,4 +194,12 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<PersistentLis
     override fun transformPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: PersistentList<FirDeclaration>): FirStatement {
         return propertyAccessor.also { transformProperty(it.propertySymbol.fir, data) }
     }
+
+    override fun transformEnumEntry(enumEntry: FirEnumEntry, data: PersistentList<FirDeclaration>): FirStatement {
+        if (enumEntry.initializer is FirLazyExpression) {
+            val designation = FirDeclarationDesignation(data, enumEntry)
+            FirLazyBodiesCalculator.calculateLazyInitializerForEnumEntry(designation)
+        }
+        return enumEntry
+    }
 }

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

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDeclarationDesignation
+import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
+import org.jetbrains.kotlin.analysis.utils.errors.withPsiEntry
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
@@ -31,7 +33,6 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
     private val functionsToRebind: Set<FirFunction>? = null,
     private val replacementApplier: RawFirReplacement.Applier? = null
 ) : RawFirBuilder(session, baseScopeProvider, bodyBuildingMode = BodyBuildingMode.NORMAL) {
-
     companion object {
         fun buildWithReplacement(
             session: FirSession,
@@ -95,7 +96,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
         }
     }
 
-    private inner class VisitorWithReplacement : Visitor() {
+    private inner class VisitorWithReplacement(private val containingClass: FirRegularClass?) : Visitor() {
         override fun convertElement(element: KtElement): FirElement? =
             super.convertElement(replacementApplier?.tryReplace(element) ?: element)
 
@@ -130,11 +131,27 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
                 additionalAnnotations = additionalAnnotations
             )
         }
+
+        override fun visitEnumEntry(enumEntry: KtEnumEntry, data: Unit?): FirElement {
+            val owner = containingClass ?: buildErrorWithAttachment("Enum entry outside of class") {
+                withPsiEntry("enumEntry", enumEntry)
+            }
+            val classOrObject = owner.psi as KtClassOrObject
+            val primaryConstructor = classOrObject.primaryConstructor
+            val ownerClassHasDefaultConstructor =
+                primaryConstructor?.valueParameters?.isEmpty() ?: classOrObject.secondaryConstructors.let { constructors ->
+                    constructors.isEmpty() || constructors.any { it.valueParameters.isEmpty() }
+                }
+            val typeParameters = mutableListOf<FirTypeParameterRef>()
+            context.appendOuterTypeParameters(ignoreLastLevel = false, typeParameters)
+            val selfType = classOrObject.toDelegatedSelfType(typeParameters, owner.symbol)
+            return enumEntry.toFirEnumEntry(selfType, ownerClassHasDefaultConstructor)
+        }
     }
 
     private fun moveNext(iterator: Iterator<FirDeclaration>, containingClass: FirRegularClass?): FirDeclaration {
         if (!iterator.hasNext()) {
-            val visitor = VisitorWithReplacement()
+            val visitor = VisitorWithReplacement(containingClass)
             return when (declarationToBuild) {
                 is KtProperty -> {
                     val ownerSymbol = containingClass?.symbol

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/FirLazyBodiesCalculatorTestGenerated.java

@@ -105,6 +105,12 @@ public class FirLazyBodiesCalculatorTestGenerated extends AbstractFirLazyBodiesC
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/enums2.kt");
         }
 
+        @Test
+        @TestMetadata("enums3.kt")
+        public void testEnums3() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/enums3.kt");
+        }
+
         @Test
         @TestMetadata("expectActual.kt")
         public void testExpectActual() throws Exception {