FIR IDE: implement lazy delegate initializers

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

@@ -1246,7 +1246,14 @@ class RawFirBuilder(
                             FirWrappedDelegateExpressionBuilder().apply {
                                 source =
                                     if (stubMode) null else delegateExpression?.toFirSourceElement(FirFakeSourceElementKind.WrappedDelegate)
-                                expression = { delegateExpression }.toFirExpression("Should have delegate")
+                                expression = when (mode) {
+                                    RawFirBuilderMode.NORMAL -> delegateExpression.toFirExpression("Should have delegate")
+                                    RawFirBuilderMode.STUBS -> buildExpressionStub()
+                                    RawFirBuilderMode.LAZY_BODIES -> buildLazyExpression {
+                                        source = delegateExpression!!.toFirSourceElement()
+                                    }
+                                }
+
                             }
                         } else null
 

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

@@ -25,5 +25,7 @@ abstract class FirLambdaArgumentExpression : FirWrappedArgumentExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
+    abstract override fun replaceExpression(newExpression: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirLambdaArgumentExpression
 }

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

@@ -27,5 +27,7 @@ abstract class FirNamedArgumentExpression : FirWrappedArgumentExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
+    abstract override fun replaceExpression(newExpression: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirNamedArgumentExpression
 }

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

@@ -25,5 +25,7 @@ abstract class FirSpreadArgumentExpression : FirWrappedArgumentExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
+    abstract override fun replaceExpression(newExpression: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirSpreadArgumentExpression
 }

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

@@ -25,5 +25,7 @@ abstract class FirWrappedArgumentExpression : FirWrappedExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
+    abstract override fun replaceExpression(newExpression: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirWrappedArgumentExpression
 }

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

@@ -25,5 +25,7 @@ abstract class FirWrappedDelegateExpression : FirWrappedExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
+    abstract override fun replaceExpression(newExpression: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirWrappedDelegateExpression
 }

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

@@ -24,5 +24,7 @@ abstract class FirWrappedExpression : FirExpression() {
 
     abstract override fun replaceTypeRef(newTypeRef: FirTypeRef)
 
+    abstract fun replaceExpression(newExpression: FirExpression)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirWrappedExpression
 }

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

@@ -42,4 +42,8 @@ internal class FirLambdaArgumentExpressionImpl(
     }
 
     override fun replaceTypeRef(newTypeRef: FirTypeRef) {}
+
+    override fun replaceExpression(newExpression: FirExpression) {
+        expression = newExpression
+    }
 }

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

@@ -44,4 +44,8 @@ internal class FirNamedArgumentExpressionImpl(
     }
 
     override fun replaceTypeRef(newTypeRef: FirTypeRef) {}
+
+    override fun replaceExpression(newExpression: FirExpression) {
+        expression = newExpression
+    }
 }

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

@@ -42,4 +42,8 @@ internal class FirSpreadArgumentExpressionImpl(
     }
 
     override fun replaceTypeRef(newTypeRef: FirTypeRef) {}
+
+    override fun replaceExpression(newExpression: FirExpression) {
+        expression = newExpression
+    }
 }

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

@@ -44,4 +44,8 @@ internal class FirWrappedDelegateExpressionImpl(
     }
 
     override fun replaceTypeRef(newTypeRef: FirTypeRef) {}
+
+    override fun replaceExpression(newExpression: FirExpression) {
+        expression = newExpression
+    }
 }

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

@@ -516,7 +516,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         wrappedExpression.configure {
-            +field(expression)
+            +field(expression).withReplace()
         }
 
         wrappedDelegateExpression.configure {

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 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.FirWrappedDelegateExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
 import org.jetbrains.kotlin.fir.psi
@@ -67,6 +68,12 @@ internal object FirLazyBodiesCalculator {
             if (firProperty.initializer is FirLazyExpression) {
                 firProperty.replaceInitializer(newProperty.initializer)
             }
+
+            val delegate = firProperty.delegate
+            if (delegate is FirWrappedDelegateExpression && delegate.expression is FirLazyExpression) {
+                val newDelegate = newProperty.delegate as FirWrappedDelegateExpression
+                delegate.replaceExpression(newDelegate.expression)
+            }
         }
     }
 
@@ -74,6 +81,7 @@ internal object FirLazyBodiesCalculator {
         firProperty.getter?.body is FirLazyBlock
                 || firProperty.setter?.body is FirLazyBlock
                 || firProperty.initializer is FirLazyExpression
+                || (firProperty.delegate as? FirWrappedDelegateExpression)?.expression is FirLazyExpression
 
 
     private fun createRawFirBuilder(firDeclaration: FirDeclaration): RawFirBuilder {

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

@@ -0,0 +1,23 @@
+import kotlin.properties.ReadWriteProperty
+import kotlin.reflect.KProperty
+
+fun resolveMe() {
+    receive(valueWithExplicitType)
+    receive(valueWithImplicitType)
+
+    variableWithExplicitType = 10
+    variableWithImplicitType = 10
+}
+
+fun receive(value: Int){}
+
+val delegate = object: ReadWriteProperty<Any?, Int> {
+    override fun getValue(thisRef: Any?, property: KProperty<*>): Int = 1
+    override fun setValue(thisRef: Any?, property: KProperty<*>, value: Int) {}
+}
+
+val valueWithExplicitType: Int by delegate
+val valueWithImplicitType by delegate
+
+var variableWithExplicitType: Int by delegate
+var variableWithImplicitType by delegate

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

@@ -0,0 +1,45 @@
+FILE: delegates.kt
+    public final [BODY_RESOLVE] fun resolveMe(): R|kotlin/Unit| {
+        R|/receive|(R|/valueWithExplicitType|)
+        R|/receive|(R|/valueWithImplicitType|)
+        R|/variableWithExplicitType| = Int(10)
+        R|/variableWithImplicitType| = Int(10)
+    }
+    public final [STATUS] fun receive(value: R|kotlin/Int|): R|kotlin/Unit| { LAZY_BLOCK }
+    public final [IMPLICIT_TYPES_BODY_RESOLVE] val delegate: R|kotlin/properties/ReadWriteProperty<kotlin/Any?, kotlin/Int>| = object : R|kotlin/properties/ReadWriteProperty<kotlin/Any?, kotlin/Int>| {
+        private constructor(): R|<anonymous>| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final override [BODY_RESOLVE] fun getValue(thisRef: R|kotlin/Any?|, property: R|kotlin/reflect/KProperty<*>|): R|kotlin/Int| {
+            ^getValue Int(1)
+        }
+
+        public final override [BODY_RESOLVE] fun setValue(thisRef: R|kotlin/Any?|, property: R|kotlin/reflect/KProperty<*>|, value: R|kotlin/Int|): R|kotlin/Unit| {
+        }
+
+    }
+
+        public get(): R|kotlin/properties/ReadWriteProperty<kotlin/Any?, kotlin/Int>|
+    public final [STATUS] val valueWithExplicitType: R|kotlin/Int|by LAZY_EXPRESSION
+        public get(): <implicit> {
+            ^ D|/valueWithExplicitType|.getValue#(Null(null), ::R|/valueWithExplicitType|)
+        }
+    public final [IMPLICIT_TYPES_BODY_RESOLVE] val valueWithImplicitType: R|kotlin/Int|by R|/delegate|
+        public get(): R|kotlin/Int| {
+            ^ D|/valueWithImplicitType|.R|FakeOverride<kotlin/properties/ReadWriteProperty.getValue: R|kotlin/Int|>|(Null(null), ::R|/valueWithImplicitType|)
+        }
+    public final [STATUS] var variableWithExplicitType: R|kotlin/Int|by LAZY_EXPRESSION
+        public get(): <implicit> {
+            ^ D|/variableWithExplicitType|.getValue#(Null(null), ::R|/variableWithExplicitType|)
+        }
+        public set(<set-?>: <implicit>): R|kotlin/Unit| {
+            D|/variableWithExplicitType|.setValue#(Null(null), ::R|/variableWithExplicitType|, R|<local>/variableWithExplicitType|)
+        }
+    public final [IMPLICIT_TYPES_BODY_RESOLVE] var variableWithImplicitType: R|kotlin/Int|by R|/delegate|
+        public get(): R|kotlin/Int| {
+            ^ D|/variableWithImplicitType|.R|FakeOverride<kotlin/properties/ReadWriteProperty.getValue: R|kotlin/Int|>|(Null(null), ::R|/variableWithImplicitType|)
+        }
+        public set(<set-?>: R|kotlin/Int|): R|kotlin/Unit| {
+            D|/variableWithImplicitType|.R|FakeOverride<kotlin/properties/ReadWriteProperty.setValue: R|kotlin/Unit|>|(Null(null), ::R|/variableWithImplicitType|, R|<local>/variableWithImplicitType|)
+        }

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirLazyDeclarationResolveTest.kt

@@ -6,11 +6,13 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.util.io.FileUtil
+import com.intellij.testFramework.LightProjectDescriptor
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacade
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtPsiUtil
@@ -44,4 +46,6 @@ abstract class AbstractFirLazyDeclarationResolveTest : KotlinLightCodeInsightFix
             KotlinTestUtils.assertEqualsToFile(testDataFile.parentFile.resolve(expectedFileName), rendered)
         }
     }
+
+    override fun getProjectDescriptor(): LightProjectDescriptor = KotlinWithJdkAndRuntimeLightProjectDescriptor.INSTANCE
 }

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

@@ -33,6 +33,11 @@ public class FirLazyDeclarationResolveTestGenerated extends AbstractFirLazyDecla
         runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/classMembers.kt");
     }
 
+    @TestMetadata("delegates.kt")
+    public void testDelegates() throws Exception {
+        runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/delegates.kt");
+    }
+
     @TestMetadata("propertyWithGetter.kt")
     public void testPropertyWithGetter() throws Exception {
         runTest("idea/idea-frontend-fir/idea-fir-low-level-api/testdata/lazyResolve/propertyWithGetter.kt");