FIR IDE: generate lazy initializers for properties

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaField.kt

@@ -112,6 +112,10 @@ class FirJavaField @FirImplementationDetail constructor(
         return this
     }
 
+    override fun replaceInitializer(newInitializer: FirExpression?) {
+        initializer = newInitializer
+    }
+
     override fun <D> transformTypeParameters(transformer: FirTransformer<D>, data: D): FirField {
         typeParameters.transformInplace(transformer, data)
         return this

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

@@ -1194,9 +1194,15 @@ class RawFirBuilder(
             val propertyType = typeReference.toFirOrImplicitType()
             val propertyName = nameAsSafeName
             val isVar = isVar
-            val propertyInitializer = if (hasInitializer()) {
-                { initializer }.toFirExpression("Should have initializer")
-            } else null
+            val propertyInitializer = when {
+                !hasInitializer() -> null
+                lazyBodiesMode -> buildLazyExpression {
+                    source = initializer?.toFirSourceElement()
+                }
+                else -> {
+                    { initializer }.toFirExpression("Should have initializer")
+                }
+            }
             val delegateExpression by lazy { delegate?.expression }
             val propertySource = toFirSourceElement()
 

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

@@ -51,5 +51,5 @@ FILE: annotation.kt
 
     }
     @base() public? final? fun foo(@base() y: @base() Int): Int { LAZY_BLOCK }
-    @base() public? final? val z: <implicit> = IntegerLiteral(0)
+    @base() public? final? val z: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>

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

@@ -44,11 +44,5 @@ FILE: constructorInObject.kt
         }
 
     }
-    public? final? val anonObject: <implicit> = object : R|kotlin/Any| {
-        public? constructor(): R|<anonymous>| {
-            super<R|kotlin/Any|>()
-        }
-
-    }
-
+    public? final? val anonObject: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/newSyntax/propertyAccessorsContractDescription.lazyBodies.txt

@@ -4,7 +4,7 @@ FILE: propertyAccessorsContractDescription.kt
             super<R|kotlin/Any|>()
         }
 
-        public? final? var myInt: Int = IntegerLiteral(0)
+        public? final? var myInt: Int = LAZY_EXPRESSION
             public? get(): Int
                 [Contract description] <
                     returnsNotNull#()
@@ -18,7 +18,7 @@ FILE: propertyAccessorsContractDescription.kt
             super<R|kotlin/Any|>()
         }
 
-        public? final? var anotherInt: Int = IntegerLiteral(0)
+        public? final? var anotherInt: Int = LAZY_EXPRESSION
             public? get(): Int
                 [Contract description] <
                     returnsNotNull#()
@@ -36,7 +36,7 @@ FILE: propertyAccessorsContractDescription.kt
             super<R|kotlin/Any|>()
         }
 
-        public? final? var someInt: Int = IntegerLiteral(0)
+        public? final? var someInt: Int = LAZY_EXPRESSION
             public? get(): Int
                 [Contract description] <
                     returnsNotNull#()

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/oldSyntax/contractDescription.lazyBodies.txt

@@ -1,7 +1,7 @@
 FILE: contractDescription.kt
     public? final? fun test_1(): R|kotlin/Unit| { LAZY_BLOCK }
     public? final? fun test_2(): R|kotlin/Unit| { LAZY_BLOCK }
-    public? final? var test_3: Int = IntegerLiteral(1)
+    public? final? var test_3: Int = LAZY_EXPRESSION
         public? get(): Int { LAZY_BLOCK }
         public? set(value: Int): R|kotlin/Unit| { LAZY_BLOCK }
     public? final? fun test_4(): R|kotlin/Unit| { LAZY_BLOCK }

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

@@ -52,7 +52,7 @@ FILE: enums.kt
 
         }
 
-        public? final? val g: Double = G#.times#(m#).div#(r#.times#(r#))
+        public? final? val g: Double = LAZY_EXPRESSION
             public? get(): Double
 
         public? abstract fun sayHello(): R|kotlin/Unit|
@@ -62,7 +62,7 @@ FILE: enums.kt
                 super<R|kotlin/Any|>()
             }
 
-            public? final? const val G: <implicit> = Double(6.67E-11)
+            public? final? const val G: <implicit> = LAZY_EXPRESSION
                 public? get(): <implicit>
 
         }

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

@@ -15,5 +15,5 @@ FILE: expectActual.kt
 
     }
     public? final? actual fun foo(): <implicit> { LAZY_BLOCK }
-    public? final? actual val x: <implicit> = IntegerLiteral(42)
+    public? final? actual val x: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>

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

@@ -11,7 +11,7 @@ FILE: simpleClass.kt
             super<<implicit>>()
         }
 
-        private final? val baz: <implicit> = IntegerLiteral(42)
+        private final? val baz: <implicit> = LAZY_EXPRESSION
             private get(): <implicit>
 
         public? open? override fun foo(x: Int, y: String): String { LAZY_BLOCK }

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

@@ -1,5 +1,5 @@
 FILE: arrayAccess.kt
-    public? final? val p: <implicit> = IntegerLiteral(0)
+    public? final? val p: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
     public? final? fun foo(): <implicit> { LAZY_BLOCK }
     public? final? class Wrapper : R|kotlin/Any| {

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

@@ -6,25 +6,25 @@ FILE: callableReferences.kt
 
         public? final? fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
 
-        public? final? val bar: <implicit> = IntegerLiteral(0)
+        public? final? val bar: <implicit> = LAZY_EXPRESSION
             public? get(): <implicit>
 
     }
     public? final? fun A.qux(): R|kotlin/Unit| { LAZY_BLOCK }
     public? final? fun baz(): R|kotlin/Unit| { LAZY_BLOCK }
-    public? final? val test1: <implicit> = A#()::foo#
+    public? final? val test1: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test2: <implicit> = A#()::bar#
+    public? final? val test2: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test3: <implicit> = A#()::qux#
+    public? final? val test3: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test4: <implicit> = A#::foo#
+    public? final? val test4: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test5: <implicit> = A#::bar#
+    public? final? val test5: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test6: <implicit> = A#::qux#
+    public? final? val test6: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test7: <implicit> = ::baz#
+    public? final? val test7: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val test8: <implicit> = A#?::foo#
+    public? final? val test8: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>

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

@@ -21,13 +21,7 @@ FILE: lambda.kt
     public? final? fun foo(): Int { LAZY_BLOCK }
     public? final? fun bar(): Int { LAZY_BLOCK }
     public? final? fun test(list: List<Int>): R|kotlin/Unit| { LAZY_BLOCK }
-    public? final? val simple: <implicit> = fun <implicit>.<anonymous>(): <implicit> {
-        ^ Unit
-    }
-
+    public? final? val simple: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>
-    public? final? val another: <implicit> = fun <implicit>.<anonymous>(): <implicit> {
-        IntegerLiteral(42)
-    }
-
+    public? final? val another: <implicit> = LAZY_EXPRESSION
         public? get(): <implicit>

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

@@ -51,6 +51,8 @@ abstract class FirEnumEntry : FirVariable<FirEnumEntry>(), FirCallableMemberDecl
 
     abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
+    abstract override fun replaceInitializer(newInitializer: FirExpression?)
+
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirEnumEntry
 
     abstract override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirEnumEntry

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

@@ -51,6 +51,8 @@ abstract class FirErrorProperty : FirVariable<FirErrorProperty>(), FirDiagnostic
 
     abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
+    abstract override fun replaceInitializer(newInitializer: FirExpression?)
+
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirErrorProperty
 
     abstract override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirErrorProperty

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

@@ -51,6 +51,8 @@ abstract class FirField : FirVariable<FirField>(), FirTypeParametersOwner, FirCa
 
     abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
+    abstract override fun replaceInitializer(newInitializer: FirExpression?)
+
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirField
 
     abstract override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirField

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

@@ -56,6 +56,8 @@ abstract class FirProperty : FirVariable<FirProperty>(), FirTypeParametersOwner,
 
     abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
+    abstract override fun replaceInitializer(newInitializer: FirExpression?)
+
     abstract override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
 
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirProperty

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

@@ -53,6 +53,8 @@ abstract class FirValueParameter : FirVariable<FirValueParameter>(), FirControlF
 
     abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
+    abstract override fun replaceInitializer(newInitializer: FirExpression?)
+
     abstract override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
 
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirValueParameter

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

@@ -49,6 +49,8 @@ abstract class FirVariable<F : FirVariable<F>> : FirPureAbstractElement(), FirCa
 
     abstract override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?)
 
+    abstract fun replaceInitializer(newInitializer: FirExpression?)
+
     abstract override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirVariable<F>
 
     abstract override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirVariable<F>

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

@@ -133,6 +133,10 @@ internal class FirDefaultSetterValueParameter(
         receiverTypeRef = newReceiverTypeRef
     }
 
+    override fun replaceInitializer(newInitializer: FirExpression?) {
+        initializer = newInitializer
+    }
+
     override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
         controlFlowGraphReference = newControlFlowGraphReference
     }

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

@@ -128,4 +128,8 @@ internal class FirEnumEntryImpl(
     }
 
     override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?) {}
+
+    override fun replaceInitializer(newInitializer: FirExpression?) {
+        initializer = newInitializer
+    }
 }

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

@@ -117,4 +117,6 @@ internal class FirErrorPropertyImpl(
     }
 
     override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?) {}
+
+    override fun replaceInitializer(newInitializer: FirExpression?) {}
 }

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

@@ -125,4 +125,6 @@ internal class FirFieldImpl(
     }
 
     override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?) {}
+
+    override fun replaceInitializer(newInitializer: FirExpression?) {}
 }

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

@@ -151,6 +151,10 @@ internal class FirPropertyImpl(
         receiverTypeRef = newReceiverTypeRef
     }
 
+    override fun replaceInitializer(newInitializer: FirExpression?) {
+        initializer = newInitializer
+    }
+
     override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
         controlFlowGraphReference = newControlFlowGraphReference
     }

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

@@ -117,6 +117,8 @@ open class FirValueParameterImpl @FirImplementationDetail constructor(
 
     override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?) {}
 
+    override fun replaceInitializer(newInitializer: FirExpression?) {}
+
     override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
         controlFlowGraphReference = newControlFlowGraphReference
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/synthetic/FirSyntheticProperty.kt

@@ -141,4 +141,8 @@ class FirSyntheticProperty(
     override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
         throw AssertionError("Mutation of synthetic property isn't supported")
     }
+
+    override fun replaceInitializer(newInitializer: FirExpression?) {
+        throw AssertionError("Mutation of synthetic property isn't supported")
+    }
 }

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

@@ -368,7 +368,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             parentArg(callableDeclaration, "F", "F")
             +name
             +symbol("FirVariableSymbol", "F")
-            +initializer.withTransform()
+            +initializer.withTransform().withReplace()
             +field("delegate", expression, nullable = true).withTransform()
             +field("delegateFieldSymbol", delegateFieldSymbolType, "F", nullable = true)
             generateBooleanFields("var", "val")

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

@@ -92,4 +92,8 @@ fun Field.withTransform(needTransformInOtherChildren: Boolean = false): Field =
     this.needTransformInOtherChildren = needTransformInOtherChildren
 }
 
+fun Field.withReplace(): Field = copy().apply {
+    withReplace = true
+}
+
 fun FieldSet.withTransform(): FieldSet = this.map { it.withTransform() }

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -48,7 +49,7 @@ internal object FirLazyBodiesCalculator {
     }
 
     fun calculateLazyBodyForProperty(firProperty: FirProperty) {
-        if (firProperty.getter?.body !is FirLazyBlock && firProperty.setter?.body !is FirLazyBlock) return
+        if (!needCalculatingLazyBodyForProperty(firProperty)) return
 
         val rawFirBuilder = createRawFirBuilder(firProperty)
         val newProperty = rawFirBuilder.buildPropertyWithBody(firProperty.psi as KtProperty)
@@ -63,9 +64,17 @@ internal object FirLazyBodiesCalculator {
                 setter.replaceBody(newSetter.body)
                 setter.replaceContractDescription(newSetter.contractDescription)
             }
+            if (firProperty.initializer is FirLazyExpression) {
+                firProperty.replaceInitializer(newProperty.initializer)
+            }
         }
     }
 
+    fun needCalculatingLazyBodyForProperty(firProperty: FirProperty): Boolean =
+        firProperty.getter?.body is FirLazyBlock
+                || firProperty.setter?.body is FirLazyBlock
+                || firProperty.initializer is FirLazyExpression
+
 
     private fun createRawFirBuilder(firDeclaration: FirDeclaration): RawFirBuilder {
         val scopeProvider = firDeclaration.session.firIdeProvider.kotlinScopeProvider
@@ -98,8 +107,9 @@ private object FirLazyBodiesCalculatorTransformer : FirTransformer<Nothing?>() {
     }
 
     override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
-        if (property.getter?.body is FirLazyBlock || property.setter?.body is FirLazyBlock) {
+        if (FirLazyBodiesCalculator.needCalculatingLazyBodyForProperty(property)) {
             FirLazyBodiesCalculator.calculateLazyBodyForProperty(property)
+            return property.compose()
         }
         return super.transformProperty(property, data)
     }

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/classMembers.txt

@@ -8,6 +8,9 @@ FILE: classMembers.kt
             this@R|/A|.R|/A.receive|(this@R|/A|.R|/A.functionWithLazyBody|())
         }
 
+        public final [STATUS] val x: R|kotlin/Int| = LAZY_EXPRESSION
+            public get(): R|kotlin/Int| { LAZY_BLOCK }
+
         public final [STATUS] fun receive(value: R|kotlin/String|): R|kotlin/Unit| { LAZY_BLOCK }
 
         public final [STATUS] fun functionWithLazyBody(): R|kotlin/String| { LAZY_BLOCK }

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/propertyWithGetterAndSetter.txt

@@ -4,6 +4,6 @@ FILE: propertyWithGetterAndSetter.kt
         R|/withGetterAndSetter| = Int(123)
     }
     public final [STATUS] fun receive(value: R|kotlin/Int|): R|kotlin/Unit| { LAZY_BLOCK }
-    public final [STATUS] var withGetterAndSetter: R|kotlin/Int| = IntegerLiteral(42)
+    public final [STATUS] var withGetterAndSetter: R|kotlin/Int| = LAZY_EXPRESSION
         public get(): R|kotlin/Int| { LAZY_BLOCK }
         public set(value: R|kotlin/Int|): R|kotlin/Unit| { LAZY_BLOCK }

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/propertyWithInitializer.kt

@@ -0,0 +1,5 @@
+fun resolveMe() {
+    receive(property)
+}
+
+val property: Int = 10

idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/propertyWithInitializer.txt

@@ -0,0 +1,6 @@
+FILE: propertyWithInitializer.kt
+    public final [BODY_RESOLVE] fun resolveMe(): R|kotlin/Unit| {
+        <Unresolved name: receive>#(R|/property|)
+    }
+    public final [STATUS] val property: R|kotlin/Int| = LAZY_EXPRESSION
+        public get(): R|kotlin/Int|

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/FirLazyDeclarationResolveTestGenerated.java

@@ -43,6 +43,11 @@ public class FirLazyDeclarationResolveTestGenerated extends AbstractFirLazyDecla
         runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/propertyWithGetterAndSetter.kt");
     }
 
+    @TestMetadata("propertyWithInitializer.kt")
+    public void testPropertyWithInitializer() throws Exception {
+        runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/propertyWithInitializer.kt");
+    }
+
     @TestMetadata("secondaryConstructor.kt")
     public void testSecondaryConstructor() throws Exception {
         runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/secondaryConstructor.kt");