[FIR IDE] Fix invalid lazy expressions for property initializers

compiler/fir/raw-fir/light-tree2fir/tests-gen/org/jetbrains/kotlin/fir/lightTree/LightTree2FirConverterTestCaseGenerated.java

@@ -353,6 +353,11 @@ public class LightTree2FirConverterTestCaseGenerated extends AbstractLightTree2F
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
         }
 
+        @TestMetadata("localDeclarationWithExpression.kt")
+        public void testLocalDeclarationWithExpression() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt");
+        }
+
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/locals.kt");

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

@@ -210,7 +210,7 @@ open class RawFirBuilder(
 
         private fun KtElement?.toFirExpression(
             errorReason: String,
-            kind: DiagnosticKind = DiagnosticKind.ExpressionExpected,
+            kind: DiagnosticKind = DiagnosticKind.ExpressionExpected
         ): FirExpression {
             if (stubMode) {
                 return buildExpressionStub()
@@ -1401,7 +1401,7 @@ open class RawFirBuilder(
                 else -> initializer.toFirExpression("Should have initializer")
 
             }
-            val delegateExpression by lazy { delegate?.expression }
+
             val propertySource = toFirSourceElement()
 
             return buildProperty {
@@ -1417,27 +1417,31 @@ open class RawFirBuilder(
                 if (this@toFirProperty.isLocal) {
                     isLocal = true
                     symbol = FirPropertySymbol(propertyName)
-                    val delegateBuilder = delegateExpression?.let {
-                        FirWrappedDelegateExpressionBuilder().apply {
-                            source = it.toFirSourceElement(FirFakeSourceElementKind.WrappedDelegate)
-                            expression = it.toFirExpression("Incorrect delegate expression")
-                        }
-                    }
-
                     status = FirDeclarationStatusImpl(Visibilities.Local, Modality.FINAL).apply {
                         isLateInit = hasModifier(LATEINIT_KEYWORD)
                     }
 
-                    val receiver = delegateExpression?.toFirExpression("Incorrect delegate expression")
+                    if (hasDelegate()) {
-                    generateAccessorsByDelegate(
+                        fun extractDelegateExpression() =
-                        delegateBuilder,
+                            this@toFirProperty.delegate?.expression.toFirExpression("Incorrect delegate expression")
-                        baseModuleData,
+
-                        ownerRegularOrAnonymousObjectSymbol = null,
+                        val delegateBuilder = FirWrappedDelegateExpressionBuilder().apply {
-                        ownerRegularClassTypeParametersCount = null,
+                            val delegateFirExpression = extractDelegateExpression()
-                        isExtension = false,
+                            source = delegateFirExpression.source?.fakeElement(FirFakeSourceElementKind.WrappedDelegate)
-                        stubMode = stubMode,
+                            expression = delegateFirExpression
-                        receiver = receiver
+                        }
-                    )
+
+                        generateAccessorsByDelegate(
+                            delegateBuilder,
+                            baseModuleData,
+                            ownerRegularOrAnonymousObjectSymbol = null,
+                            ownerRegularClassTypeParametersCount = null,
+                            isExtension = false,
+                            stubMode = stubMode,
+                            //TODO This expression should be the same for wrapper and receiver
+                            receiver = extractDelegateExpression()
+                        )
+                    }
                 } else {
                     isLocal = false
                     receiverTypeRef = receiverTypeReference.convertSafe()
@@ -1445,21 +1449,6 @@ open class RawFirBuilder(
                     dispatchReceiverType = currentDispatchReceiverType()
                     extractTypeParametersTo(this, symbol)
                     withCapturedTypeParameters(true, this.typeParameters) {
-                        val delegateBuilder = if (hasDelegate()) {
-                            FirWrappedDelegateExpressionBuilder().apply {
-                                source =
-                                    if (stubMode) null else delegateExpression?.toFirSourceElement(FirFakeSourceElementKind.WrappedDelegate)
-                                expression = when (mode) {
-                                    BodyBuildingMode.NORMAL -> delegateExpression.toFirExpression("Should have delegate")
-                                    BodyBuildingMode.STUBS -> buildExpressionStub()
-                                    BodyBuildingMode.LAZY_BODIES -> buildLazyExpression {
-                                        source = delegateExpression!!.toFirSourceElement()
-                                    }
-                                }
-
-                            }
-                        } else null
-
                         getter = this@toFirProperty.getter.toFirPropertyAccessor(this@toFirProperty, propertyType, isGetter = true)
                         setter = this@toFirProperty.setter.toFirPropertyAccessor(this@toFirProperty, propertyType, isGetter = false)
 
@@ -1472,16 +1461,34 @@ open class RawFirBuilder(
                             isExternal = hasModifier(EXTERNAL_KEYWORD)
                         }
 
-                        val receiver = delegateExpression?.toFirExpression("Should have delegate")
+                        if (hasDelegate()) {
-                        generateAccessorsByDelegate(
+                            fun extractDelegateExpression() = when (mode) {
-                            delegateBuilder,
+                                BodyBuildingMode.NORMAL -> this@toFirProperty.delegate?.expression.toFirExpression("Should have delegate")
-                            baseModuleData,
+                                BodyBuildingMode.STUBS -> buildExpressionStub()
-                            ownerRegularOrAnonymousObjectSymbol,
+                                BodyBuildingMode.LAZY_BODIES -> this@toFirProperty.delegate?.expression?.let {
-                            ownerRegularClassTypeParametersCount,
+                                    buildLazyExpression { source = it.toFirSourceElement() }
-                            isExtension = receiverTypeReference != null,
+                                } ?: buildErrorExpression {
-                            stubMode = stubMode,
+                                    ConeSimpleDiagnostic("Should have delegate", DiagnosticKind.ExpressionExpected)
-                            receiver = receiver
+                                }
-                        )
+                            }
+
+                            val delegateBuilder = FirWrappedDelegateExpressionBuilder().apply {
+                                val delegateExpression = extractDelegateExpression()
+                                source = delegateExpression.source?.fakeElement(FirFakeSourceElementKind.WrappedDelegate)
+                                expression = extractDelegateExpression()
+                            }
+
+                            generateAccessorsByDelegate(
+                                delegateBuilder,
+                                baseModuleData,
+                                ownerRegularOrAnonymousObjectSymbol,
+                                ownerRegularClassTypeParametersCount,
+                                isExtension = receiverTypeReference != null,
+                                stubMode = stubMode,
+                                //TODO This expression should be the same for wrapper and receiver
+                                receiver = extractDelegateExpression()
+                            )
+                        }
                     }
                 }
                 extractAnnotationsTo(this)

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt

@@ -0,0 +1,4 @@
+private val nonLocalProperty: List<XXX> by lazy {
+    val localProperty = mutableListOf<KtLightField>()
+    localProperty
+}

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

@@ -0,0 +1,5 @@
+FILE: localDeclarationWithExpression.kt
+    private final? val nonLocalProperty: List<XXX>by LAZY_EXPRESSION
+        public? get(): <implicit> {
+            ^ D|/nonLocalProperty|.getValue#(Null(null), ::R|/nonLocalProperty|)
+        }

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.txt

@@ -0,0 +1,9 @@
+FILE: localDeclarationWithExpression.kt
+    private final? val nonLocalProperty: List<XXX>by lazy#(<L> = lazy@fun <implicit>.<anonymous>(): <implicit> <inline=Unknown>  {
+        lval localProperty: <implicit> = mutableListOf#<KtLightField>()
+        localProperty#
+    }
+    )
+        public? get(): <implicit> {
+            ^ D|/nonLocalProperty|.getValue#(Null(null), ::R|/nonLocalProperty|)
+        }

compiler/fir/raw-fir/psi2fir/tests-gen/org/jetbrains/kotlin/fir/builder/RawFirBuilderLazyBodiesTestCaseGenerated.java

@@ -353,6 +353,11 @@ public class RawFirBuilderLazyBodiesTestCaseGenerated extends AbstractRawFirBuil
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
         }
 
+        @TestMetadata("localDeclarationWithExpression.kt")
+        public void testLocalDeclarationWithExpression() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt");
+        }
+
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/locals.kt");

compiler/fir/raw-fir/psi2fir/tests-gen/org/jetbrains/kotlin/fir/builder/RawFirBuilderTestCaseGenerated.java

@@ -353,6 +353,11 @@ public class RawFirBuilderTestCaseGenerated extends AbstractRawFirBuilderTestCas
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
         }
 
+        @TestMetadata("localDeclarationWithExpression.kt")
+        public void testLocalDeclarationWithExpression() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt");
+        }
+
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/locals.kt");

compiler/visualizer/tests-gen/org/jetbrains/kotlin/visualizer/fir/FirVisualizerForRawFirDataGenerated.java

@@ -385,6 +385,12 @@ public class FirVisualizerForRawFirDataGenerated extends AbstractFirVisualizerTe
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
         }
 
+        @Test
+        @TestMetadata("localDeclarationWithExpression.kt")
+        public void testLocalDeclarationWithExpression() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt");
+        }
+
         @Test
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {

compiler/visualizer/tests-gen/org/jetbrains/kotlin/visualizer/psi/PsiVisualizerForRawFirDataGenerated.java

@@ -385,6 +385,12 @@ public class PsiVisualizerForRawFirDataGenerated extends AbstractPsiVisualizerTe
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
         }
 
+        @Test
+        @TestMetadata("localDeclarationWithExpression.kt")
+        public void testLocalDeclarationWithExpression() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt");
+        }
+
         @Test
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {

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

@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirWrappedDelegateExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
@@ -90,10 +92,20 @@ internal object FirLazyBodiesCalculator {
             firProperty.replaceInitializer(newProperty.initializer)
         }
 
-        val delegate = firProperty.delegate
+        val delegate = firProperty.delegate as? FirWrappedDelegateExpression
-        if (delegate is FirWrappedDelegateExpression && delegate.expression is FirLazyExpression) {
+        val delegateExpression = delegate?.expression
-            val newDelegate = newProperty.delegate as FirWrappedDelegateExpression
+        if (delegateExpression is FirLazyExpression) {
+            val newDelegate = newProperty.delegate as? FirWrappedDelegateExpression
+            check(newDelegate != null) { "Invalid replacement delegate" }
             delegate.replaceExpression(newDelegate.expression)
+
+            val delegateProviderCall = delegate.delegateProvider as? FirFunctionCall
+            val delegateProviderExplicitReceiver = delegateProviderCall?.explicitReceiver
+            if (delegateProviderExplicitReceiver is FirLazyExpression) {
+                val newDelegateProviderExplicitReceiver = (newDelegate.delegateProvider as? FirFunctionCall)?.explicitReceiver
+                check(newDelegateProviderExplicitReceiver != null) { "Invalid replacement expression" }
+                delegateProviderCall.replaceExplicitReceiver(newDelegateProviderExplicitReceiver)
+            }
         }
     }
 

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

@@ -6,10 +6,14 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import junit.framework.TestCase
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
+import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.AbstractLowLevelApiSingleFileTest
@@ -18,6 +22,15 @@ import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 
 abstract class AbstractFirLazyBodiesCalculatorTest : AbstractLowLevelApiSingleFileTest() {
+
+    private val lazyChecker = object : FirVisitorVoid() {
+        override fun visitElement(element: FirElement) {
+            TestCase.assertFalse("${FirLazyBlock::class.qualifiedName} should not present in the tree", element is FirLazyBlock)
+            TestCase.assertFalse("${FirLazyExpression::class.qualifiedName} should not present in the tree", element is FirLazyExpression)
+            element.acceptChildren(this)
+        }
+    }
+
     override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
         resolveWithClearCaches(ktFile) { resolveState ->
             val session = resolveState.rootModuleSession
@@ -31,6 +44,7 @@ abstract class AbstractFirLazyBodiesCalculatorTest : AbstractLowLevelApiSingleFi
             ).buildFirFile(ktFile)
 
             FirLazyBodiesCalculator.calculateLazyBodies(laziedFirFile)
+            laziedFirFile.accept(lazyChecker)
 
             val fullFirFile = RawFirBuilder(
                 session,

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

@@ -385,6 +385,12 @@ public class FirLazyBodiesCalculatorTestGenerated extends AbstractFirLazyBodiesC
             runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
         }
 
+        @Test
+        @TestMetadata("localDeclarationWithExpression.kt")
+        public void testLocalDeclarationWithExpression() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/localDeclarationWithExpression.kt");
+        }
+
         @Test
         @TestMetadata("locals.kt")
         public void testLocals() throws Exception {