FIR IDE: implement lazy functions, constructors & accessors bodies building for raw fir

test

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

@@ -130,6 +130,10 @@ class FirJavaConstructor @FirImplementationDetail constructor(
     override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?) {}
 
     override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {}
+
+    override fun replaceBody(newBody: FirBlock?) {
+        error("Body cannot be replaced for FirJavaConstructor")
+    }
 }
 
 @FirBuilderDsl

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

@@ -26,7 +26,6 @@ import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
-import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.builder.*
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
@@ -50,7 +49,7 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class RawFirBuilder(
-    session: FirSession, val baseScopeProvider: FirScopeProvider, val stubMode: Boolean
+    session: FirSession, val baseScopeProvider: FirScopeProvider, val stubMode: Boolean, val lazyBodiesMode: Boolean = false
 ) : BaseFirBuilder<PsiElement>(session) {
 
     fun buildFirFile(file: KtFile): FirFile {
@@ -62,16 +61,22 @@ class RawFirBuilder(
     }
 
     fun buildFunctionWithBody(function: KtNamedFunction): FirFunction<*> {
-        assert(!stubMode) { "Building FIR function with body isn't supported in stub mode" }
-        setupContextForPosition(function)
-        return function.accept(Visitor(), Unit) as FirFunction<*>
+        return buildDeclaration(function) as FirFunction<*>
+    }
+
+    fun buildSecondaryConstructor(secondaryConstructor: KtSecondaryConstructor): FirConstructor {
+        return buildDeclaration(secondaryConstructor) as FirConstructor
     }
 
     fun buildPropertyWithBody(property: KtProperty): FirProperty {
         require(!property.isLocal) { "Should not be used to build local properties (variables)" }
-        assert(!stubMode) { "Building FIR function with body isn't supported in stub mode" }
-        setupContextForPosition(property)
-        return property.accept(Visitor(), Unit) as FirProperty
+        return buildDeclaration(property) as FirProperty
+    }
+
+    private fun buildDeclaration(declaration: KtDeclaration): FirDeclaration {
+        assert(!stubMode) { "Building FIR declarations isn't supported in stub mode" }
+        setupContextForPosition(declaration)
+        return declaration.accept(Visitor(), Unit) as FirDeclaration
     }
 
     private fun setupContextForPosition(position: KtElement) {
@@ -253,6 +258,13 @@ class RawFirBuilder(
             when {
                 !hasBody() ->
                     null to null
+                lazyBodiesMode -> {
+                    val block = buildLazyBlock {
+                        source = bodyExpression?.toFirSourceElement()
+                            ?: error("hasBody() == true but body is null")
+                    }
+                    block to null
+                }
                 hasBlockBody() -> if (!stubMode) {
                     val block = bodyBlockExpression?.accept(this@Visitor, Unit) as? FirBlock
                     if (hasContractEffectList()) {
@@ -1144,6 +1156,8 @@ class RawFirBuilder(
                 extractAnnotationsTo(this)
                 typeParameters += constructorTypeParametersFromConstructedClass(ownerTypeParameters)
                 extractValueParametersTo(this)
+
+
                 val (body, _) = buildFirBody()
                 this.body = body
                 this@RawFirBuilder.context.firFunctionTargets.removeLast()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt

@@ -100,10 +100,21 @@ fun <F : FirClass<F>> F.runContractAndBodiesResolutionForLocalClass(
     }
 }
 
-fun createReturnTypeCalculatorForIDE(session: FirSession, scopeSession: ScopeSession): ReturnTypeCalculator =
-    ReturnTypeCalculatorWithJump(session, scopeSession, ImplicitBodyResolveComputationSession())
+fun createReturnTypeCalculatorForIDE(
+    session: FirSession,
+    scopeSession: ScopeSession,
+    createTransformer: (
+        designation: Iterator<FirElement>,
+        FirSession,
+        ScopeSession,
+        ImplicitBodyResolveComputationSession,
+        ReturnTypeCalculator,
+        BodyResolveContext?
+    ) -> FirDesignatedBodyResolveTransformerForReturnTypeCalculator
+): ReturnTypeCalculator =
+    ReturnTypeCalculatorWithJump(session, scopeSession, ImplicitBodyResolveComputationSession(), createTransformer = createTransformer)
 
-private open class FirImplicitAwareBodyResolveTransformer(
+open class FirImplicitAwareBodyResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
     private val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
@@ -167,7 +178,15 @@ private class ReturnTypeCalculatorWithJump(
     private val session: FirSession,
     private val scopeSession: ScopeSession,
     val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
-    val designationMapForLocalClasses: Map<FirCallableMemberDeclaration<*>, List<FirClass<*>>> = mapOf()
+    val designationMapForLocalClasses: Map<FirCallableMemberDeclaration<*>, List<FirClass<*>>> = mapOf(),
+    private val createTransformer: (
+        designation: Iterator<FirElement>,
+        session: FirSession,
+        scopeSession: ScopeSession,
+        implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+        returnTypeCalculator: ReturnTypeCalculator,
+        outerBodyResolveContext: BodyResolveContext?
+    ) -> FirDesignatedBodyResolveTransformerForReturnTypeCalculator = ::FirDesignatedBodyResolveTransformerForReturnTypeCalculator,
 ) : ReturnTypeCalculator {
 
     var outerBodyResolveContext: BodyResolveContext? = null
@@ -232,7 +251,7 @@ private class ReturnTypeCalculatorWithJump(
             (listOf(file) + outerClasses.filterNotNull().asReversed()) to null
         }
 
-        val transformer = FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
+        val transformer = createTransformer(
             (designation.drop(1) + declaration).iterator(),
             session,
             scopeSession,
@@ -252,7 +271,7 @@ private class ReturnTypeCalculatorWithJump(
     }
 }
 
-private class FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
+open class FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
     private val designation: Iterator<FirElement>,
     session: FirSession,
     scopeSession: ScopeSession,
@@ -283,10 +302,10 @@ private class FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
     }
 }
 
-private class ImplicitBodyResolveComputationSession {
+class ImplicitBodyResolveComputationSession {
     private val implicitBodyResolveStatusMap = hashMapOf<FirCallableSymbol<*>, ImplicitBodyResolveComputationStatus>()
 
-    fun getStatus(symbol: FirCallableSymbol<*>): ImplicitBodyResolveComputationStatus {
+    internal fun getStatus(symbol: FirCallableSymbol<*>): ImplicitBodyResolveComputationStatus {
         if (symbol is FirAccessorSymbol) {
             val fir = symbol.fir
             if (fir is FirSyntheticProperty) {
@@ -321,7 +340,7 @@ private class ImplicitBodyResolveComputationSession {
     }
 }
 
-private sealed class ImplicitBodyResolveComputationStatus {
+internal sealed class ImplicitBodyResolveComputationStatus {
     object NotComputed : ImplicitBodyResolveComputationStatus()
     object Computing : ImplicitBodyResolveComputationStatus()
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt

@@ -49,7 +49,7 @@ internal class FirModuleResolveStateImpl(
         elementBuilder.getOrBuildFirFor(element, rootModuleSession.cache, fileStructureCache)
 
     override fun getFirFile(ktFile: KtFile): FirFile =
-        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache)
+        firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, lazyBodiesMode = false)
 
     override fun getDiagnostics(element: KtElement): List<Diagnostic> =
         diagnosticsCollector.getDiagnosticsFor(element)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirPhaseRunner.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
 import java.util.concurrent.locks.ReentrantLock
 import kotlin.concurrent.withLock
@@ -29,6 +30,9 @@ internal class FirPhaseRunner(val transformerProvider: FirTransformerProvider) {
 
     private fun runPhaseWithoutLock(firFile: FirFile, phase: FirResolvePhase) {
         val phaseProcessor = transformerProvider.getTransformerForPhase(firFile.session, phase)
-        executeWithoutPCE { phaseProcessor.processFile(firFile) }
+        executeWithoutPCE {
+            FirLazyBodiesCalculator.calculateLazyBodiesIfPhaseRequires(firFile, phase)
+            phaseProcessor.processFile(firFile)
+        }
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostics/AbstractFirIdeDiagnosticsCollector.kt

@@ -9,15 +9,13 @@ import com.intellij.openapi.diagnostic.Logger
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
-import org.jetbrains.kotlin.fir.analysis.collectors.components.*
 import org.jetbrains.kotlin.fir.analysis.collectors.registerAllComponents
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.createReturnTypeCalculatorForIDE
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.addValueFor
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.psi.KtElement
 
@@ -25,7 +23,11 @@ internal abstract class AbstractFirIdeDiagnosticsCollector(
     session: FirSession,
 ) : AbstractDiagnosticCollector(
     session,
-    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, ScopeSession())
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(
+        session,
+        ScopeSession(),
+        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+    )
 ) {
     init {
         registerAllComponents()

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirDesignatedBodyResolveTransformerForIDE.kt

@@ -28,7 +28,11 @@ internal class FirDesignatedBodyResolveTransformerForIDE(
     phase = phase,
     implicitTypeOnly = phase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
     scopeSession = scopeSession,
-    returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession)
+    returnTypeCalculator = createReturnTypeCalculatorForIDE(
+        session,
+        scopeSession,
+        ::FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator
+    )
 ) {
 
     override fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
+
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.ResolutionMode
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.BodyResolveContext
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDesignatedBodyResolveTransformerForReturnTypeCalculator
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator
+
+class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator(
+    designation: Iterator<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+    returnTypeCalculator: ReturnTypeCalculator,
+    outerBodyResolveContext: BodyResolveContext?,
+) : FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
+    designation,
+    session,
+    scopeSession,
+    implicitBodyResolveComputationSession,
+    returnTypeCalculator,
+    outerBodyResolveContext
+) {
+    override fun transformSimpleFunction(
+        simpleFunction: FirSimpleFunction,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirSimpleFunction> {
+        FirLazyBodiesCalculator.calculateLazyBodiesForFunction(simpleFunction)
+        return super.transformSimpleFunction(simpleFunction, data)
+    }
+
+    override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
+        FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(constructor)
+        return super.transformConstructor(constructor, data)
+    }
+
+    override fun transformProperty(property: FirProperty, data: ResolutionMode): CompositeTransformResult<FirProperty> {
+        FirLazyBodiesCalculator.calculateLazyBodyForProperty(property)
+        return super.transformProperty(property, data)
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/builder/FirFileBuilder.kt

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.psi.KtFile
@@ -31,9 +30,10 @@ internal class FirFileBuilder(
      */
     fun buildRawFirFileWithCaching(
         ktFile: KtFile,
-        cache: ModuleFileCache
+        cache: ModuleFileCache,
+        lazyBodiesMode: Boolean
     ): FirFile = cache.fileCached(ktFile) {
-        RawFirBuilder(cache.session, scopeProvider, stubMode = false).buildFirFile(ktFile)
+        RawFirBuilder(cache.session, scopeProvider, stubMode = false, lazyBodiesMode = lazyBodiesMode).buildFirFile(ktFile)
     }
 
     fun isFirFileBuilt(
@@ -47,8 +47,9 @@ internal class FirFileBuilder(
         @Suppress("SameParameterValue") toPhase: FirResolvePhase,
         checkPCE: Boolean
     ): FirFile {
-        val firFile = buildRawFirFileWithCaching(ktFile, cache)
-        if (toPhase > FirResolvePhase.RAW_FIR) {
+        val needResolve = toPhase > FirResolvePhase.RAW_FIR
+        val firFile = buildRawFirFileWithCaching(ktFile, cache, lazyBodiesMode = !needResolve)
+        if (needResolve) {
             cache.firFileLockProvider.withLock(firFile) {
                 if (firFile.resolvePhase >= toPhase) return@withLock
                 runResolveWithoutLock(firFile, fromPhase = firFile.resolvePhase, toPhase = toPhase, checkPCE = checkPCE)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructureCache.kt

@@ -21,7 +21,7 @@ internal class FileStructureCache(
     private val cache = ConcurrentHashMap<KtFile, FileStructure>()
 
     fun getFileStructure(ktFile: KtFile, moduleFileCache: ModuleFileCache): FileStructure = cache.computeIfAbsent(ktFile) {
-        val firFile = fileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
+        val firFile = fileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = false)
         FileStructure(ktFile, firFile, firLazyDeclarationResolver, fileBuilder, moduleFileCache)
     }
 }

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

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+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.impl.FirLazyBlock
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.compose
+import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
+import org.jetbrains.kotlin.psi.*
+
+internal object FirLazyBodiesCalculator {
+    fun calculateLazyBodiesInside(element: FirElement) {
+        element.transform<FirElement, Nothing?>(FirLazyBodiesCalculatorTransformer, null)
+    }
+
+    fun calculateLazyBodiesIfPhaseRequires(firFile: FirFile, phase: FirResolvePhase) {
+        if (phase == FIRST_PHASE_WHICH_NEEDS_BODIES) {
+            calculateLazyBodiesInside(firFile)
+        }
+    }
+
+    fun calculateLazyBodiesForFunction(simpleFunction: FirSimpleFunction) {
+        if (simpleFunction.body !is FirLazyBlock) return
+        val rawFirBuilder = createRawFirBuilder(simpleFunction)
+        val newFunction = rawFirBuilder.buildFunctionWithBody(simpleFunction.psi as KtNamedFunction) as FirSimpleFunction
+        simpleFunction.apply {
+            replaceBody(newFunction.body)
+            replaceContractDescription(newFunction.contractDescription)
+        }
+    }
+
+    fun calculateLazyBodyForSecondaryConstructor(secondaryConstructor: FirConstructor) {
+        require(!secondaryConstructor.isPrimary)
+        if (secondaryConstructor.body !is FirLazyBlock) return
+        val rawFirBuilder = createRawFirBuilder(secondaryConstructor)
+        val newFunction = rawFirBuilder.buildSecondaryConstructor(secondaryConstructor.psi as KtSecondaryConstructor)
+        secondaryConstructor.apply {
+            replaceBody(newFunction.body)
+        }
+    }
+
+    fun calculateLazyBodyForProperty(firProperty: FirProperty) {
+        if (firProperty.getter?.body !is FirLazyBlock && firProperty.setter?.body !is FirLazyBlock) return
+
+        val rawFirBuilder = createRawFirBuilder(firProperty)
+        val newProperty = rawFirBuilder.buildPropertyWithBody(firProperty.psi as KtProperty)
+        newProperty.apply {
+            getter?.takeIf { it.body is FirLazyBlock }?.let { getter ->
+                val newGetter = newProperty.getter!!
+                getter.replaceBody(newGetter.body)
+                getter.replaceContractDescription(newGetter.contractDescription)
+            }
+            setter?.takeIf { it.body is FirLazyBlock }?.let { setter ->
+                val newSetter = newProperty.setter!!
+                setter.replaceBody(newSetter.body)
+                setter.replaceContractDescription(newSetter.contractDescription)
+            }
+        }
+    }
+
+
+    private fun createRawFirBuilder(firDeclaration: FirDeclaration): RawFirBuilder {
+        val scopeProvider = firDeclaration.session.firIdeProvider.kotlinScopeProvider
+        return RawFirBuilder(firDeclaration.session, scopeProvider, stubMode = false)
+    }
+
+    private val FIRST_PHASE_WHICH_NEEDS_BODIES = FirResolvePhase.CONTRACTS
+}
+
+private object FirLazyBodiesCalculatorTransformer : FirTransformer<Nothing?>() {
+    override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
+        @Suppress("UNCHECKED_CAST")
+        return (element.transformChildren(this, data) as E).compose()
+    }
+
+    override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        if (simpleFunction.body is FirLazyBlock) {
+            FirLazyBodiesCalculator.calculateLazyBodiesForFunction(simpleFunction)
+            return simpleFunction.compose()
+        }
+        return (simpleFunction.transformChildren(this, data) as FirDeclaration).compose()
+    }
+
+    override fun transformConstructor(constructor: FirConstructor, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        if (constructor.body is FirLazyBlock) {
+            FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(constructor)
+            return constructor.compose()
+        }
+        return (constructor.transformChildren(this, data) as FirDeclaration).compose()
+    }
+
+    override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+        if (property.getter?.body is FirLazyBlock || property.setter?.body is FirLazyBlock) {
+            FirLazyBodiesCalculator.calculateLazyBodyForProperty(property)
+        }
+        return super.transformProperty(property, data)
+    }
+}

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

@@ -38,8 +38,12 @@ internal class FirLazyDeclarationResolver(
             ?: error("FirFile was not found for\n${declaration.render()}")
         val provider = firFile.session.firIdeProvider
         val fromPhase = if (reresolveFile) declaration.resolvePhase else minOf(firFile.resolvePhase, declaration.resolvePhase)
+
         if (checkPCE) {
             firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) {
+                executeWithoutPCE {
+                    calculateLazyBodies(declaration)
+                }
                 runLazyResolveWithoutLock(
                     declaration,
                     moduleFileCache,
@@ -54,6 +58,7 @@ internal class FirLazyDeclarationResolver(
         } else {
             firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) {
                 executeWithoutPCE {
+                    calculateLazyBodies(declaration)
                     runLazyResolveWithoutLock(
                         declaration,
                         moduleFileCache,
@@ -69,6 +74,10 @@ internal class FirLazyDeclarationResolver(
         }
     }
 
+    private fun calculateLazyBodies(firDeclaration: FirDeclaration) {
+        FirLazyBodiesCalculator.calculateLazyBodiesInside(firDeclaration)
+    }
+
     fun runLazyResolveWithoutLock(
         firDeclarationToResolve: FirDeclaration,
         moduleFileCache: ModuleFileCache,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/providers/FirIdeProvider.kt

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.fir.ThreadSafeMutableState
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
-import org.jetbrains.kotlin.fir.dependenciesWithoutSelf
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
@@ -24,11 +23,9 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
-import org.jetbrains.kotlin.idea.fir.low.level.api.PackageExistenceCheckerForMultipleModules
 import org.jetbrains.kotlin.idea.fir.low.level.api.PackageExistenceCheckerForSingleModule
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.collectTransitiveDependenciesWithSelf
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -40,7 +37,7 @@ internal class FirIdeProvider(
     project: Project,
     val session: FirSession,
     moduleInfo: ModuleSourceInfo,
-    private val kotlinScopeProvider: KotlinScopeProvider,
+    val kotlinScopeProvider: KotlinScopeProvider,
     firFileBuilder: FirFileBuilder,
     val cache: ModuleFileCache,
     searchScope: GlobalSearchScope,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/providers/FirProviderHelper.kt

@@ -37,7 +37,7 @@ internal class FirProviderHelper(
                     ?: indexHelper.typeAliasFromIndexByClassId(classId)
                     ?: return@computeIfAbsent Optional.empty()
                 if (ktClass is KtEnumEntry) return@computeIfAbsent Optional.empty()
-                val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile, cache)
+                val firFile = firFileBuilder.buildRawFirFileWithCaching(ktClass.containingKtFile, cache, lazyBodiesMode = true)
                 val classifier = FirElementFinder.findElementIn<FirClassLikeDeclaration<*>>(firFile) { classifier ->
                     classifier.symbol.classId == classId
                 }
@@ -59,7 +59,7 @@ internal class FirProviderHelper(
                 @OptIn(ExperimentalStdlibApi::class)
                 buildList {
                     files.forEach { ktFile ->
-                        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, cache)
+                        val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, cache, lazyBodiesMode = true)
                         firFile.collectCallableDeclarationsTo(this, name)
                     }
                 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt

@@ -32,7 +32,7 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
                 containerClassFir.declarations
             } else {
                 val ktFile = containingKtFile
-                val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache)
+                val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = true)
                 firFile.declarations
             }
             val original = originalDeclaration