[FIR IDE] Fix invalid lazy expressions for property initializers

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
-        if (delegate is FirWrappedDelegateExpression && delegate.expression is FirLazyExpression) {
-            val newDelegate = newProperty.delegate as FirWrappedDelegateExpression
+        val delegate = firProperty.delegate as? FirWrappedDelegateExpression
+        val delegateExpression = delegate?.expression
+        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 {