[FIR IDE] RawFirBuilder for lazy bodies

2021-03-10 18:18:06 +01:00
parent 1558bd7373
commit 25929b50d9
30 changed files with 956 additions and 172 deletions
@@ -116,6 +116,11 @@ public class LightTree2FirConverterTestCaseGenerated extends AbstractLightTree2F
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericProperty.kt");
        }

+        @TestMetadata("initBlockWithDeclarations.kt")
+        public void testInitBlockWithDeclarations() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.kt");
+        }
+
        @TestMetadata("nestedClass.kt")
        public void testNestedClass() throws Exception {
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/nestedClass.kt");
@@ -0,0 +1,71 @@
+/*
+ * Copyright 2010-2021 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.fir.builder
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.containingClass
+import org.jetbrains.kotlin.fir.containingClassAttr
+import org.jetbrains.kotlin.fir.declarations.FirCallableMemberDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.isInner
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.psi.*
+
+class RawFirFragmentForLazyBodiesBuilder private constructor(
+    session: FirSession,
+    baseScopeProvider: FirScopeProvider,
+    private val declaration: KtDeclaration,
+) : RawFirBuilder(session, baseScopeProvider, RawFirBuilderMode.NORMAL) {
+
+    companion object {
+        fun build(
+            session: FirSession,
+            baseScopeProvider: FirScopeProvider,
+            designation: List<FirDeclaration>,
+            declaration: KtDeclaration,
+        ): FirDeclaration {
+            require(declaration is KtNamedFunction || declaration is KtProperty) { "Not implemented for ${declaration::class.qualifiedName}" }
+            val builder = RawFirFragmentForLazyBodiesBuilder(session, baseScopeProvider, declaration)
+            builder.context.packageFqName = declaration.containingKtFile.packageFqName
+            return builder.moveNext(designation.iterator())
+        }
+    }
+
+    private fun moveNext(iterator: Iterator<FirDeclaration>): FirDeclaration {
+        if (!iterator.hasNext()) {
+            return if (declaration is KtProperty) {
+                with(Visitor()) {
+                    declaration.toFirProperty(null)
+                }
+            } else {
+                declaration.accept(Visitor(), Unit) as FirDeclaration
+            }
+        }
+
+        val parent = iterator.next()
+        if (parent !is FirRegularClass) return moveNext(iterator)
+
+        val classOrObject = parent.psi
+        check(classOrObject is KtClassOrObject)
+
+        withChildClassName(classOrObject.nameAsSafeName, false) {
+            withCapturedTypeParameters {
+                if (!parent.isInner) context.capturedTypeParameters = context.capturedTypeParameters.clear()
+                addCapturedTypeParameters(parent.typeParameters.take(classOrObject.typeParameters.size))
+                registerSelfType(classOrObject.toDelegatedSelfType(parent))
+                return moveNext(iterator)
+            }
+        }
+    }
+
+    private fun PsiElement?.toDelegatedSelfType(firClass: FirRegularClass): FirResolvedTypeRef =
+        toDelegatedSelfType(firClass.typeParameters, firClass.symbol)
+}
+
@@ -49,12 +49,25 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.runIf

-class RawFirBuilder(
-    session: FirSession, val baseScopeProvider: FirScopeProvider, val mode: RawFirBuilderMode = RawFirBuilderMode.NORMAL
+open class RawFirBuilder(
+    session: FirSession, val baseScopeProvider: FirScopeProvider, builderMode: RawFirBuilderMode = RawFirBuilderMode.NORMAL
 ) : BaseFirBuilder<PsiElement>(session) {

    private val stubMode get() = mode == RawFirBuilderMode.STUBS

+    var mode: RawFirBuilderMode = builderMode
+        private set
+
+    private inline fun <T> disabledLazyMode(body: () -> T): T {
+        if (mode != RawFirBuilderMode.LAZY_BODIES) return body()
+        return try {
+            mode = RawFirBuilderMode.NORMAL
+            body()
+        } finally {
+            mode = RawFirBuilderMode.LAZY_BODIES
+        }
+    }
+
    fun buildFirFile(file: KtFile): FirFile {
        return file.accept(Visitor(), Unit) as FirFile
    }
@@ -63,49 +76,6 @@ class RawFirBuilder(
        return reference.accept(Visitor(), Unit) as FirTypeRef
    }

-    fun buildFunctionWithBody(function: KtNamedFunction, original: FirFunction<*>?): FirFunction<*> {
-        return buildDeclaration(function, original) as FirFunction<*>
-    }
-
-    fun buildSecondaryConstructor(secondaryConstructor: KtSecondaryConstructor, original: FirConstructor?): FirConstructor {
-        return buildDeclaration(secondaryConstructor, original) as FirConstructor
-    }
-
-    fun buildPropertyWithBody(property: KtProperty, original: FirProperty?): FirProperty {
-        require(!property.isLocal) { "Should not be used to build local properties (variables)" }
-        return buildDeclaration(property, original) as FirProperty
-    }
-
-    private fun buildDeclaration(declaration: KtDeclaration, original: FirDeclaration?): FirDeclaration {
-        assert(mode ==  RawFirBuilderMode.NORMAL) { "Building FIR declarations isn't supported in stub or lazy mode mode" }
-        setupContextForPosition(declaration,)
-        val firDeclaration = declaration.accept(Visitor(), Unit) as FirDeclaration
-        original?.let { firDeclaration.copyContainingClassAttrFrom(it) }
-        return firDeclaration
-    }
-
-    // TODO this is a (temporary) hack, instead we should properly initialize [context]
-    private fun FirDeclaration.copyContainingClassAttrFrom(from: FirDeclaration) {
-        (this as? FirCallableMemberDeclaration<*>)?.let {
-            it.containingClassAttr = (from as? FirCallableMemberDeclaration<*>)?.containingClass()
-        }
-    }
-
-    private fun setupContextForPosition(position: KtElement) {
-        val parentsUpToFile = position.parents
-        for (parent in parentsUpToFile.toList().asReversed()) {
-            when (parent) {
-                is KtFile -> {
-                    context.packageFqName = parent.packageFqName
-                }
-                is KtClassOrObject -> {
-                    context.className = context.className.child(parent.nameAsSafeName)
-                    context.localBits.add(parent.isLocal || parent.getStrictParentOfType<KtEnumEntry>() != null)
-                }
-            }
-        }
-    }
-
    override fun PsiElement.toFirSourceElement(kind: FirFakeSourceElementKind?): FirPsiSourceElement<*> {
        val actualKind = kind ?: this@RawFirBuilder.context.forcedElementSourceKind ?: FirRealSourceElementKind
        return this.toFirPsiSourceElement(actualKind)
@@ -178,7 +148,7 @@ class RawFirBuilder(
            }
        }

-    private inner class Visitor : KtVisitor<FirElement, Unit>() {
+    protected inner class Visitor : KtVisitor<FirElement, Unit>() {
        private inline fun <reified R : FirElement> KtElement?.convertSafe(): R? =
            this?.accept(this@Visitor, Unit) as? R

@@ -237,12 +207,14 @@ class RawFirBuilder(
        ): FirDeclaration {
            return when (this) {
                is KtSecondaryConstructor -> {
-                    toFirConstructor(
-                        delegatedSuperType,
-                        delegatedSelfType,
-                        owner,
-                        ownerTypeParameters
-                    )
+                    disabledLazyMode {
+                        toFirConstructor(
+                            delegatedSuperType,
+                            delegatedSelfType,
+                            owner,
+                            ownerTypeParameters
+                        )
+                    }
                }
                is KtEnumEntry -> {
                    val primaryConstructor = owner.primaryConstructor
@@ -250,7 +222,9 @@ class RawFirBuilder(
                        primaryConstructor?.valueParameters?.isEmpty() ?: owner.secondaryConstructors.let { constructors ->
                            constructors.isEmpty() || constructors.any { it.valueParameters.isEmpty() }
                        }
-                    toFirEnumEntry(delegatedSelfType, ownerClassHasDefaultConstructor)
+                    disabledLazyMode {
+                        toFirEnumEntry(delegatedSelfType, ownerClassHasDefaultConstructor)
+                    }
                }
                is KtProperty -> {
                    toFirProperty(ownerClassBuilder)
@@ -548,7 +522,7 @@ class RawFirBuilder(
            container.argumentList = argumentList
        }

-        private fun KtClassOrObject.extractSuperTypeListEntriesTo(
+        fun KtClassOrObject.extractSuperTypeListEntriesTo(
            container: FirClassBuilder,
            delegatedSelfTypeRef: FirTypeRef?,
            delegatedEnumSuperTypeRef: FirTypeRef?,
@@ -790,6 +764,12 @@ class RawFirBuilder(
            }
        }

+        override fun visitClassInitializer(initializer: KtClassInitializer, data: Unit?): FirElement {
+            return disabledLazyMode {
+                super.visitClassInitializer(initializer, data)
+            }
+        }
+
        override fun visitClassOrObject(classOrObject: KtClassOrObject, data: Unit): FirElement {
            return withChildClassName(
                classOrObject.nameAsSafeName,
@@ -874,7 +854,8 @@ class RawFirBuilder(
                        }

                        if (classOrObject.hasModifier(DATA_KEYWORD) && firPrimaryConstructor != null) {
-                            val zippedParameters = classOrObject.primaryConstructorParameters.filter { it.hasValOrVar() } zip declarations.filterIsInstance<FirProperty>()
+                            val zippedParameters =
+                                classOrObject.primaryConstructorParameters.filter { it.hasValOrVar() } zip declarations.filterIsInstance<FirProperty>()
                            DataClassMembersGenerator(
                                baseSession,
                                classOrObject,
@@ -1215,7 +1196,7 @@ class RawFirBuilder(
            }
        }

-        private fun KtProperty.toFirProperty(ownerClassBuilder: FirClassBuilder?): FirProperty {
+        fun KtProperty.toFirProperty(ownerClassBuilder: FirClassBuilder?): FirProperty {
            val propertyType = typeReference.toFirOrImplicitType()
            val propertyName = nameAsSafeName
            val isVar = isVar
@@ -0,0 +1,10 @@
+// FUNCTION: foo
+
+package test.classes
+
+class Outer<X> {
+    fun foo() {
+        val z = object { }
+    }
+}
+
@@ -0,0 +1,9 @@
+public? final? fun test/classes/Outer.foo(): R|kotlin/Unit| {
+    lval <local>/z: <implicit> = object : R|kotlin/Any| {
+        private constructor(): R|<anonymous><X>| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+
+}
@@ -30,7 +30,9 @@ FILE: enums.kt
                super<R|Planet|>(Double(1.0), Double(2.0))
            }

-            public? open? override fun sayHello(): R|kotlin/Unit| { LAZY_BLOCK }
+            public? open? override fun sayHello(): R|kotlin/Unit| {
+                println#(String(Hello!!!))
+            }

        }

@@ -39,7 +41,9 @@ FILE: enums.kt
                super<R|Planet|>(Double(3.0), Double(4.0))
            }

-            public? open? override fun sayHello(): R|kotlin/Unit| { LAZY_BLOCK }
+            public? open? override fun sayHello(): R|kotlin/Unit| {
+                println#(String(Ola!!!))
+            }

        }

@@ -48,7 +52,9 @@ FILE: enums.kt
                super<R|Planet|>(Double(5.0), Double(6.0))
            }

-            public? open? override fun sayHello(): R|kotlin/Unit| { LAZY_BLOCK }
+            public? open? override fun sayHello(): R|kotlin/Unit| {
+                println#(String(Privet!!!))
+            }

        }

@@ -26,7 +26,9 @@ FILE: enums2.kt
                super<R|SomeEnum|>(O1#)
            }

-            public? open? override fun check(y: Some): Boolean { LAZY_BLOCK }
+            public? open? override fun check(y: Some): Boolean {
+                ^check Boolean(true)
+            }

        }

@@ -35,7 +37,9 @@ FILE: enums2.kt
                super<R|SomeEnum|>(O2#)
            }

-            public? open? override fun check(y: Some): Boolean { LAZY_BLOCK }
+            public? open? override fun check(y: Some): Boolean {
+                ^check ==(y#, O2#)
+            }

        }

@@ -0,0 +1,22 @@
+
+class X {
+  init {
+    class classInInit {
+        fun funInClassInInit() {
+        }
+    }
+    fun funInInit() {
+    }
+  }
+}
+
+object {
+  init {
+    class classInInit {
+        fun funInClassInInit() {
+        }
+    }
+    fun funInInit() {
+    }
+  }
+}
@@ -0,0 +1,45 @@
+FILE: initBlockWithDeclarations.kt
+    public? final? class X : R|kotlin/Any| {
+        public? constructor(): R|X| {
+            super<R|kotlin/Any|>()
+        }
+
+        init {
+            local final? class classInInit : R|kotlin/Any| {
+                public? constructor(): R|X.classInInit| {
+                    super<R|kotlin/Any|>()
+                }
+
+                public? final? fun funInClassInInit(): R|kotlin/Unit| {
+                }
+
+            }
+
+            local final? fun funInInit(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }
+    public? final? object <no name provided> : R|kotlin/Any| {
+        private constructor(): R|<no name provided>| {
+            super<R|kotlin/Any|>()
+        }
+
+        init {
+            local final? class classInInit : R|kotlin/Any| {
+                public? constructor(): R|<no name provided>.classInInit| {
+                    super<R|kotlin/Any|>()
+                }
+
+                public? final? fun funInClassInInit(): R|kotlin/Unit| {
+                }
+
+            }
+
+            local final? fun funInInit(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }
@@ -0,0 +1,45 @@
+FILE: initBlockWithDeclarations.kt
+    public? final? class X : R|kotlin/Any| {
+        public? constructor(): R|X| {
+            super<R|kotlin/Any|>()
+        }
+
+        init {
+            local final? class classInInit : R|kotlin/Any| {
+                public? constructor(): R|X.classInInit| {
+                    super<R|kotlin/Any|>()
+                }
+
+                public? final? fun funInClassInInit(): R|kotlin/Unit| {
+                }
+
+            }
+
+            local final? fun funInInit(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }
+    public? final? object <no name provided> : R|kotlin/Any| {
+        private constructor(): R|<no name provided>| {
+            super<R|kotlin/Any|>()
+        }
+
+        init {
+            local final? class classInInit : R|kotlin/Any| {
+                public? constructor(): R|<no name provided>.classInInit| {
+                    super<R|kotlin/Any|>()
+                }
+
+                public? final? fun funInClassInInit(): R|kotlin/Unit| {
+                }
+
+            }
+
+            local final? fun funInInit(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }
@@ -3,7 +3,10 @@ FILE: noPrimaryConstructor.kt
        public? final? val x: String
            public? get(): String

-        public? constructor(x: String): R|NoPrimary| { LAZY_BLOCK }
+        public? constructor(x: String): R|NoPrimary| {
+            super<R|kotlin/Any|>()
+            this#.x# = x#
+        }

        public? constructor(): R|NoPrimary| {
            this<R|NoPrimary|>(String())
@@ -44,6 +44,11 @@ public class PartialRawFirBuilderTestCaseGenerated extends AbstractPartialRawFir
        runTest("compiler/fir/raw-fir/psi2fir/testData/partialRawBuilder/memberProperty.kt");
    }

+    @TestMetadata("paramemtersCatching.kt")
+    public void testParamemtersCatching() throws Exception {
+        runTest("compiler/fir/raw-fir/psi2fir/testData/partialRawBuilder/paramemtersCatching.kt");
+    }
+
    @TestMetadata("simpleFunction.kt")
    public void testSimpleFunction() throws Exception {
        runTest("compiler/fir/raw-fir/psi2fir/testData/partialRawBuilder/simpleFunction.kt");
@@ -116,6 +116,11 @@ public class RawFirBuilderLazyBodiesTestCaseGenerated extends AbstractRawFirBuil
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericProperty.kt");
        }

+        @TestMetadata("initBlockWithDeclarations.kt")
+        public void testInitBlockWithDeclarations() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.kt");
+        }
+
        @TestMetadata("nestedClass.kt")
        public void testNestedClass() throws Exception {
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/nestedClass.kt");
@@ -116,6 +116,11 @@ public class RawFirBuilderTestCaseGenerated extends AbstractRawFirBuilderTestCas
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericProperty.kt");
        }

+        @TestMetadata("initBlockWithDeclarations.kt")
+        public void testInitBlockWithDeclarations() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.kt");
+        }
+
        @TestMetadata("nestedClass.kt")
        public void testNestedClass() throws Exception {
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/nestedClass.kt");
@@ -5,14 +5,15 @@

 package org.jetbrains.kotlin.fir.builder

+import junit.framework.TestCase
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirRenderer
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
-import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
@@ -34,29 +35,67 @@ abstract class AbstractPartialRawFirBuilderTestCase : AbstractRawFirBuilderTestC

    private fun testFunctionPartialBuilding(filePath: String, nameToFind: String) {
        testPartialBuilding(
-            filePath,
-            { file -> file.findDescendantOfType<KtNamedFunction> { it.name == nameToFind }!! }
-        ) { rawFirBuilder, function -> rawFirBuilder.buildFunctionWithBody(function, original = null) }
+            filePath
+        ) { file -> file.findDescendantOfType<KtNamedFunction> { it.name == nameToFind }!! }
    }

    private fun testPropertyPartialBuilding(filePath: String, nameToFind: String) {
        testPartialBuilding(
-            filePath,
-            { file -> file.findDescendantOfType<KtProperty> { it.name == nameToFind }!! }
-        ) { rawFirBuilder, property -> rawFirBuilder.buildPropertyWithBody(property, original = null) }
+            filePath
+        ) { file -> file.findDescendantOfType<KtProperty> { it.name == nameToFind }!! }
+    }
+
+    private class DesignationBuilder(private val elementToBuild: KtDeclaration) : FirVisitorVoid() {
+        val designation = mutableListOf<FirDeclaration>()
+        var originalDeclaration: FirDeclaration? = null
+        var built = false
+
+        override fun visitElement(element: FirElement) {
+            if (built) return
+            when (element) {
+                is FirSimpleFunction, is FirProperty -> {
+                    if (element.psi == elementToBuild) {
+                        originalDeclaration = element as FirDeclaration
+                        built = true
+                    } else {
+                        element.acceptChildren(this)
+                    }
+                }
+                is FirRegularClass -> {
+                    designation.add(element)
+                    element.acceptChildren(this)
+                    if (!built) {
+                        designation.removeLast()
+                    }
+                }
+                else -> {
+                    element.acceptChildren(this)
+                }
+            }
+        }
    }

    private fun <T : KtElement> testPartialBuilding(
        filePath: String,
-        findPsiElement: (KtFile) -> T,
-        buildFirElement: (RawFirBuilder, T) -> FirElement
+        findPsiElement: (KtFile) -> T
    ) {
        val file = createKtFile(filePath)
-        val elementToBuild = findPsiElement(file)
+        val elementToBuild = findPsiElement(file) as KtDeclaration

        val session = FirSessionFactory.createEmptySession()
+        val firBuilder = RawFirBuilder(session, StubFirScopeProvider)
+        val original = firBuilder.buildFirFile(file)

-        val firElement = buildFirElement(RawFirBuilder(session, StubFirScopeProvider), elementToBuild)
+        val designationBuilder = DesignationBuilder(elementToBuild)
+        original.accept(designationBuilder)
+        TestCase.assertTrue(designationBuilder.built)
+
+        val firElement = RawFirFragmentForLazyBodiesBuilder.build(
+            session,
+            StubFirScopeProvider,
+            designationBuilder.designation,
+            elementToBuild
+        )

        val firDump = firElement.render(FirRenderer.RenderMode.WithFqNames)
        val expectedPath = filePath.replace(".kt", ".txt")
@@ -92,9 +92,11 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte

    inline fun <T> withCapturedTypeParameters(block: () -> T): T {
        val previous = context.capturedTypeParameters
-        val result = block()
-        context.capturedTypeParameters = previous
-        return result
+        return try {
+            block()
+        } finally {
+            context.capturedTypeParameters = previous
+        }
    }

    fun addCapturedTypeParameters(typeParameters: List<FirTypeParameterRef>) {
@@ -191,17 +193,17 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
    }

    fun T?.toDelegatedSelfType(firClass: FirRegularClassBuilder): FirResolvedTypeRef =
-        toDelegatedSelfType(firClass, firClass.symbol)
+        toDelegatedSelfType(firClass.typeParameters, firClass.symbol)

    fun T?.toDelegatedSelfType(firObject: FirAnonymousObjectBuilder): FirResolvedTypeRef =
-        toDelegatedSelfType(firObject, firObject.symbol)
+        toDelegatedSelfType(firObject.typeParameters, firObject.symbol)

-    private fun T?.toDelegatedSelfType(firClass: FirClassBuilder, symbol: FirClassLikeSymbol<*>): FirResolvedTypeRef {
+    protected fun T?.toDelegatedSelfType(typeParameters: List<FirTypeParameterRef>, symbol: FirClassLikeSymbol<*>): FirResolvedTypeRef {
        return buildResolvedTypeRef {
            source = this@toDelegatedSelfType?.toFirSourceElement(FirFakeSourceElementKind.ClassSelfTypeRef)
            type = ConeClassLikeTypeImpl(
                symbol.toLookupTag(),
-                firClass.typeParameters.map { ConeTypeParameterTypeImpl(it.symbol.toLookupTag(), false) }.toTypedArray(),
+                typeParameters.map { ConeTypeParameterTypeImpl(it.symbol.toLookupTag(), false) }.toTypedArray(),
                false
            )
        }
@@ -123,6 +123,12 @@ public class FirVisualizerForRawFirDataGenerated extends AbstractFirVisualizerTe
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericProperty.kt");
        }

+        @Test
+        @TestMetadata("initBlockWithDeclarations.kt")
+        public void testInitBlockWithDeclarations() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.kt");
+        }
+
        @Test
        @TestMetadata("nestedClass.kt")
        public void testNestedClass() throws Exception {
@@ -123,6 +123,12 @@ public class PsiVisualizerForRawFirDataGenerated extends AbstractPsiVisualizerTe
            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericProperty.kt");
        }

+        @Test
+        @TestMetadata("initBlockWithDeclarations.kt")
+        public void testInitBlockWithDeclarations() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.kt");
+        }
+
        @Test
        @TestMetadata("nestedClass.kt")
        public void testNestedClass() throws Exception {
@@ -1069,6 +1069,12 @@ fun main(args: Array<String>) {
            }
        }

+        testGroup("idea/idea-frontend-fir/idea-fir-low-level-api/tests", "compiler/fir/raw-fir/psi2fir/testData") {
+            testClass<AbstractFirLazyBodiesCalculatorTest> {
+                model("rawBuilder", testMethod = "doTest")
+            }
+        }
+
        testGroup("idea/idea-frontend-fir/idea-fir-low-level-api/tests", "idea/idea-frontend-fir/idea-fir-low-level-api/testdata") {
            testClass<AbstractFirMultiModuleLazyResolveTest> {
                model("multiModuleLazyResolve", recursive = false, extension = null)
@@ -1,28 +1,22 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 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.api

-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyAccessorCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildPropertyCopy
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunctionCopy
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
-import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateForCompletion
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveStateImpl
-import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.util.getElementTextInContext
-import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNamedFunction
 import org.jetbrains.kotlin.psi.KtProperty

@@ -66,7 +60,7 @@ object LowLevelFirApiFacadeForCompletion {
        firIdeProvider: FirIdeProvider,
        element: KtNamedFunction,
        originalFunction: FirSimpleFunction,
-        state: FirModuleResolveState
+        state: FirModuleResolveState,
    ): FirSimpleFunction {
        val builtFunction = firIdeProvider.buildFunctionWithBody(element, originalFunction)

@@ -79,7 +73,7 @@ object LowLevelFirApiFacadeForCompletion {
            resolvePhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
            source = builtFunction.source
            session = state.rootModuleSession
-        }.apply { reassignAllReturnTargets (builtFunction) }
+        }.apply { reassignAllReturnTargets(builtFunction) }
    }

    private fun buildPropertyCopyForCompletion(
@@ -17,36 +17,62 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyRe
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator

-class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator(
+fun FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculator(
    designation: Iterator<FirElement>,
    session: FirSession,
    scopeSession: ScopeSession,
    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
    returnTypeCalculator: ReturnTypeCalculator,
    outerBodyResolveContext: BodyResolveContext?,
+): FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl {
+
+    val designationList = mutableListOf<FirElement>()
+    for (element in designation) {
+        designationList.add(element)
+    }
+
+    return FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
+        designationList,
+        session,
+        scopeSession,
+        implicitBodyResolveComputationSession,
+        returnTypeCalculator,
+        outerBodyResolveContext
+    )
+}
+
+class FirIdeDesignatedBodyResolveTransformerForReturnTypeCalculatorImpl(
+    private val designation: List<FirElement>,
+    session: FirSession,
+    scopeSession: ScopeSession,
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+    returnTypeCalculator: ReturnTypeCalculator,
+    outerBodyResolveContext: BodyResolveContext?,
 ) : FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
-    designation,
+    designation.iterator(),
    session,
    scopeSession,
    implicitBodyResolveComputationSession,
    returnTypeCalculator,
    outerBodyResolveContext
 ) {
+    private val declarationDesignation = designation.filterIsInstance<FirDeclaration>()
+
    override fun transformSimpleFunction(
        simpleFunction: FirSimpleFunction,
        data: ResolutionMode
    ): CompositeTransformResult<FirSimpleFunction> {
-        FirLazyBodiesCalculator.calculateLazyBodiesForFunction(simpleFunction)
+        FirLazyBodiesCalculator.calculateLazyBodiesForFunction(simpleFunction, declarationDesignation)
        return super.transformSimpleFunction(simpleFunction, data)
    }

    override fun transformConstructor(constructor: FirConstructor, data: ResolutionMode): CompositeTransformResult<FirDeclaration> {
-        FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(constructor)
+        FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(constructor, declarationDesignation)
        return super.transformConstructor(constructor, data)
    }

    override fun transformProperty(property: FirProperty, data: ResolutionMode): CompositeTransformResult<FirProperty> {
-        FirLazyBodiesCalculator.calculateLazyBodyForProperty(property)
+        FirLazyBodiesCalculator.calculateLazyBodyForProperty(property, declarationDesignation)
        return super.transformProperty(property, data)
    }
 }
@@ -1,12 +1,11 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 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.file.structure

 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
@@ -20,10 +19,8 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import java.util.concurrent.ConcurrentHashMap

-
 internal class FileStructure(
    private val ktFile: KtFile,
    private val firFile: FirFile,
@@ -48,11 +45,11 @@ internal class FileStructure(
                structureElement == null -> createStructureElement(declaration)
                structureElement is ReanalyzableStructureElement<KtDeclaration> && !structureElement.isUpToDate() -> {
                    structureElement.reanalyze(
-                        declaration as KtDeclaration,
-                        moduleFileCache,
-                        firLazyDeclarationResolver,
-                        firIdeProvider,
-                        collector
+                        newKtDeclaration = declaration as KtDeclaration,
+                        cache = moduleFileCache,
+                        firLazyDeclarationResolver = firLazyDeclarationResolver,
+                        firIdeProvider = firIdeProvider,
+                        towerDataContextCollector = collector
                    )
                }
                else -> structureElement
@@ -5,10 +5,8 @@

 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure

-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorDeclarationAction
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
@@ -6,54 +6,66 @@
 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.builder.RawFirFragmentForLazyBodiesBuilder
 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
-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.fir.visitors.*
 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 calculateLazyBodiesInside(element: FirElement, designation: List<FirDeclaration>) {
+        element.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, designation.toMutableList())
    }

    fun calculateLazyBodiesIfPhaseRequires(firFile: FirFile, phase: FirResolvePhase) {
        if (phase == FIRST_PHASE_WHICH_NEEDS_BODIES) {
-            calculateLazyBodiesInside(firFile)
+            firFile.transform<FirElement, MutableList<FirDeclaration>>(FirLazyBodiesCalculatorTransformer, mutableListOf())
        }
    }

-    fun calculateLazyBodiesForFunction(simpleFunction: FirSimpleFunction) {
+    fun calculateLazyBodiesForFunction(simpleFunction: FirSimpleFunction, designation: List<FirDeclaration>) {
        if (simpleFunction.body !is FirLazyBlock) return
-        val rawFirBuilder = createRawFirBuilder(simpleFunction)
-        val newFunction = rawFirBuilder.buildFunctionWithBody(simpleFunction.psi as KtNamedFunction, simpleFunction) as FirSimpleFunction
+        val newFunction = RawFirFragmentForLazyBodiesBuilder.build(
+            session = simpleFunction.session,
+            baseScopeProvider = simpleFunction.session.firIdeProvider.kotlinScopeProvider,
+            designation = designation,
+            declaration = simpleFunction.psi as KtNamedFunction
+        ) as FirSimpleFunction
        simpleFunction.apply {
            replaceBody(newFunction.body)
            replaceContractDescription(newFunction.contractDescription)
        }
    }

-    fun calculateLazyBodyForSecondaryConstructor(secondaryConstructor: FirConstructor) {
+    fun calculateLazyBodyForSecondaryConstructor(secondaryConstructor: FirConstructor, designation: List<FirDeclaration>) {
        require(!secondaryConstructor.isPrimary)
        if (secondaryConstructor.body !is FirLazyBlock) return
-        val rawFirBuilder = createRawFirBuilder(secondaryConstructor)
-        val newFunction = rawFirBuilder.buildSecondaryConstructor(secondaryConstructor.psi as KtSecondaryConstructor, secondaryConstructor)
+
+        val newFunction = RawFirFragmentForLazyBodiesBuilder.build(
+            session = secondaryConstructor.session,
+            baseScopeProvider = secondaryConstructor.session.firIdeProvider.kotlinScopeProvider,
+            designation = designation,
+            declaration = secondaryConstructor.psi as KtSecondaryConstructor
+        ) as FirSimpleFunction
+
        secondaryConstructor.apply {
            replaceBody(newFunction.body)
        }
    }

-    fun calculateLazyBodyForProperty(firProperty: FirProperty) {
+    fun calculateLazyBodyForProperty(firProperty: FirProperty, designation: List<FirDeclaration>) {
        if (!needCalculatingLazyBodyForProperty(firProperty)) return

-        val rawFirBuilder = createRawFirBuilder(firProperty)
-        val newProperty = rawFirBuilder.buildPropertyWithBody(firProperty.psi as KtProperty, firProperty)
+        val newProperty = RawFirFragmentForLazyBodiesBuilder.build(
+            session = firProperty.session,
+            baseScopeProvider = firProperty.session.firIdeProvider.kotlinScopeProvider,
+            designation = designation,
+            declaration = firProperty.psi as KtProperty
+        ) as FirProperty

        firProperty.getter?.takeIf { it.body is FirLazyBlock }?.let { getter ->
            val newGetter = newProperty.getter!!
@@ -84,42 +96,54 @@ internal object FirLazyBodiesCalculator {
                || firProperty.initializer is FirLazyExpression
                || (firProperty.delegate as? FirWrappedDelegateExpression)?.expression is FirLazyExpression

-
-    private fun createRawFirBuilder(firDeclaration: FirDeclaration): RawFirBuilder {
-        val scopeProvider = firDeclaration.session.firIdeProvider.kotlinScopeProvider
-        return RawFirBuilder(firDeclaration.session, scopeProvider)
-    }
-
    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()
+private object FirLazyBodiesCalculatorTransformer : FirTransformer<MutableList<FirDeclaration>>() {
+
+    override fun transformFile(file: FirFile, data: MutableList<FirDeclaration>): CompositeTransformResult<FirDeclaration> {
+        file.declarations.forEach {
+            it.transformSingle(this, data)
+        }
+        return file.compose()
    }

-    override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+    override fun <E : FirElement> transformElement(element: E, data: MutableList<FirDeclaration>): CompositeTransformResult<E> {
+        if (element is FirRegularClass) {
+            data.add(element)
+            element.declarations.forEach {
+                it.transformSingle(this, data)
+            }
+            element.transformChildren(this, data)
+            data.removeLast()
+        }
+        return element.compose()
+    }
+
+    override fun transformSimpleFunction(
+        simpleFunction: FirSimpleFunction,
+        data: MutableList<FirDeclaration>
+    ): CompositeTransformResult<FirDeclaration> {
        if (simpleFunction.body is FirLazyBlock) {
-            FirLazyBodiesCalculator.calculateLazyBodiesForFunction(simpleFunction)
-            return simpleFunction.compose()
+            FirLazyBodiesCalculator.calculateLazyBodiesForFunction(simpleFunction, data)
        }
-        return (simpleFunction.transformChildren(this, data) as FirDeclaration).compose()
+        return simpleFunction.compose()
    }

-    override fun transformConstructor(constructor: FirConstructor, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+    override fun transformConstructor(
+        constructor: FirConstructor,
+        data: MutableList<FirDeclaration>
+    ): CompositeTransformResult<FirDeclaration> {
        if (constructor.body is FirLazyBlock) {
-            FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(constructor)
-            return constructor.compose()
+            FirLazyBodiesCalculator.calculateLazyBodyForSecondaryConstructor(constructor, data)
        }
-        return (constructor.transformChildren(this, data) as FirDeclaration).compose()
+        return constructor.compose()
    }

-    override fun transformProperty(property: FirProperty, data: Nothing?): CompositeTransformResult<FirDeclaration> {
+    override fun transformProperty(property: FirProperty, data: MutableList<FirDeclaration>): CompositeTransformResult<FirDeclaration> {
        if (FirLazyBodiesCalculator.needCalculatingLazyBodyForProperty(property)) {
-            FirLazyBodiesCalculator.calculateLazyBodyForProperty(property)
-            return property.compose()
+            FirLazyBodiesCalculator.calculateLazyBodyForProperty(property, data)
        }
-        return super.transformProperty(property, data)
+        return property.compose()
    }
 }
@@ -91,8 +91,8 @@ internal class FirLazyDeclarationResolver(
        }
    }

-    private fun calculateLazyBodies(firDeclaration: FirDeclaration) {
-        FirLazyBodiesCalculator.calculateLazyBodiesInside(firDeclaration)
+    private fun calculateLazyBodies(firDeclaration: FirDeclaration, designation: List<FirDeclaration>) {
+        FirLazyBodiesCalculator.calculateLazyBodiesInside(firDeclaration, designation)
    }

    fun runLazyResolveWithoutLock(
@@ -117,8 +117,12 @@ internal class FirLazyDeclarationResolver(
        }
        if (toPhase <= nonLazyPhase) return

+
+        val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
+        val designation = nonLocalDeclarationToResolve.getDesignation(containerFirFile, provider, moduleFileCache)
+
        executeWithoutPCE {
-            calculateLazyBodies(firDeclarationToResolve)
+            calculateLazyBodies(firDeclarationToResolve, designation)
        }

        var currentPhase = nonLazyPhase
@@ -129,30 +133,22 @@ internal class FirLazyDeclarationResolver(
            if (currentPhase.pluginPhase) continue
            if (checkPCE) checkCanceled()
            runLazyResolvePhase(
-                firDeclarationToResolve,
                containerFirFile,
-                moduleFileCache,
-                provider,
                currentPhase,
                scopeSession,
-                towerDataContextCollector
+                towerDataContextCollector,
+                designation
            )
        }
    }

    private fun runLazyResolvePhase(
-        firDeclarationToResolve: FirDeclaration,
        containerFirFile: FirFile,
-        moduleFileCache: ModuleFileCache,
-        provider: FirProvider,
        phase: FirResolvePhase,
        scopeSession: ScopeSession,
        towerDataContextCollector: FirTowerDataContextCollector?,
+        designation: List<FirDeclaration>
    ) {
-        val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider, moduleFileCache)
-
-        val designation = nonLocalDeclarationToResolve.getDesignation(containerFirFile, provider, moduleFileCache)
-
        if (designation.all { it.resolvePhase >= phase }) {
            return
        }
@@ -10,7 +10,7 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.NoMutableState
 import org.jetbrains.kotlin.fir.ThreadSafeMutableState
-import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.builder.RawFirFragmentForLazyBodiesBuilder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.originalForSubstitutionOverride
@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
 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.collectDesignation
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -95,15 +96,30 @@ internal class FirIdeProvider(

    // TODO move out of here
    // used only for completion
-    fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction, original: FirFunction<*>): FirFunction<*> {
-        return RawFirBuilder(session, kotlinScopeProvider).buildFunctionWithBody(ktNamedFunction, original)
+    fun buildFunctionWithBody(
+        ktNamedFunction: KtNamedFunction,
+        original: FirFunction<*>,
+    ): FirFunction<*> {
+        return RawFirFragmentForLazyBodiesBuilder.build(
+            session = original.session,
+            baseScopeProvider = original.session.firIdeProvider.kotlinScopeProvider,
+            designation = original.collectDesignation(),
+            declaration = ktNamedFunction
+        ) as FirFunction<*>
    }

-    fun buildPropertyWithBody(ktNamedFunction: KtProperty, original: FirProperty): FirProperty {
-        return RawFirBuilder(session, kotlinScopeProvider).buildPropertyWithBody(ktNamedFunction, original)
+    fun buildPropertyWithBody(
+        ktNamedFunction: KtProperty,
+        original: FirProperty,
+    ): FirProperty {
+        return RawFirFragmentForLazyBodiesBuilder.build(
+            session = original.session,
+            baseScopeProvider = original.session.firIdeProvider.kotlinScopeProvider,
+            designation = original.collectDesignation(),
+            declaration = ktNamedFunction
+        ) as FirProperty
    }

-
    @FirProviderInternals
    override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {
        TODO()
@@ -5,10 +5,8 @@

 package org.jetbrains.kotlin.idea.fir.low.level.api.util

-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
+import org.jetbrains.kotlin.fir.containingClass
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
@@ -16,6 +14,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.api.InvalidFirElementTypeExce
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 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.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.name.ClassId
@@ -134,4 +133,20 @@ private fun KtTypeAlias.findFir(firSymbolProvider: FirSymbolProvider): FirTypeAl
 }

 val FirDeclaration.isGeneratedDeclaration
-    get() = realPsi == null
+    get() = realPsi == null
+
+internal fun FirDeclaration.collectDesignation(): List<FirDeclaration> {
+    require(this is FirCallableDeclaration<*>)
+    val designation = mutableListOf<FirDeclaration>()
+    val firProvider = session.firIdeProvider
+    var containingClassId = containingClass()?.classId
+    while (containingClassId != null) {
+        val klass = firProvider.getFirClassifierByFqName(containingClassId)
+        if (klass != null) {
+            designation.add(klass)
+        }
+        containingClassId = containingClassId.outerClassId
+    }
+    designation.reverse()
+    return designation
+}
@@ -4,6 +4,9 @@ FILE: secondaryConstructor.kt
    }
    public final [STATUS] fun receive([STATUS] value: R|A|): R|kotlin/Unit| { LAZY_BLOCK }
    public final [STATUS] class A : R|kotlin/Any| {
-        public [STATUS] constructor([STATUS] x: R|kotlin/Int|): R|A| { LAZY_BLOCK }
+        public [STATUS] constructor([STATUS] x: R|kotlin/Int|): R|A| {
+            super<R|kotlin/Any|>()
+            [RAW_FIR] lval a: <implicit> = x#
+        }

    }
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2021 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
+
+import junit.framework.TestCase
+import org.jetbrains.kotlin.fir.FirRenderer
+import org.jetbrains.kotlin.fir.builder.RawFirBuilder
+import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getResolveState
+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.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.psi.KtFile
+
+abstract class AbstractFirLazyBodiesCalculatorTest : KotlinLightCodeInsightFixtureTestCase() {
+
+    override fun isFirPlugin(): Boolean = true
+
+    protected fun doTest(filePath: String) {
+
+        val file = myFixture.configureByFile(fileName()) as KtFile
+        val resolveState = file.getResolveState()
+        val session = resolveState.rootModuleSession
+        val provider = session.firIdeProvider.kotlinScopeProvider
+
+        val laziedFirFile = RawFirBuilder(session, provider, RawFirBuilderMode.LAZY_BODIES).buildFirFile(file)
+        FirLazyBodiesCalculator.calculateLazyBodiesIfPhaseRequires(laziedFirFile, FirResolvePhase.CONTRACTS)
+        val fullFirFile = RawFirBuilder(session, provider, RawFirBuilderMode.NORMAL).buildFirFile(file)
+
+        val laziedFirFileDump = StringBuilder().also { FirRenderer(it).visitFile(laziedFirFile) }.toString()
+        val fullFirFileDump = StringBuilder().also { FirRenderer(it).visitFile(fullFirFile) }.toString()
+
+        TestCase.assertEquals(laziedFirFileDump, fullFirFileDump)
+    }
+}
@@ -0,0 +1,401 @@
+/*
+ * Copyright 2010-2021 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;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("compiler/fir/raw-fir/psi2fir/testData/rawBuilder")
+@TestDataPath("$PROJECT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+public class FirLazyBodiesCalculatorTestGenerated extends AbstractFirLazyBodiesCalculatorTest {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    public void testAllFilesPresentInRawBuilder() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/raw-fir/psi2fir/testData/rawBuilder"), Pattern.compile("^(.+)\\.kt$"), null, true);
+    }
+
+    @TestMetadata("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Declarations extends AbstractFirLazyBodiesCalculatorTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInDeclarations() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations"), Pattern.compile("^(.+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("annotation.kt")
+        public void testAnnotation() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/annotation.kt");
+        }
+
+        @TestMetadata("complexTypes.kt")
+        public void testComplexTypes() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/complexTypes.kt");
+        }
+
+        @TestMetadata("constructorInObject.kt")
+        public void testConstructorInObject() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorInObject.kt");
+        }
+
+        @TestMetadata("constructorOfAnonymousObject.kt")
+        public void testConstructorOfAnonymousObject() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/constructorOfAnonymousObject.kt");
+        }
+
+        @TestMetadata("delegates.kt")
+        public void testDelegates() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/delegates.kt");
+        }
+
+        @TestMetadata("derivedClass.kt")
+        public void testDerivedClass() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/derivedClass.kt");
+        }
+
+        @TestMetadata("emptyAnonymousObject.kt")
+        public void testEmptyAnonymousObject() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/emptyAnonymousObject.kt");
+        }
+
+        @TestMetadata("enums.kt")
+        public void testEnums() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/enums.kt");
+        }
+
+        @TestMetadata("enums2.kt")
+        public void testEnums2() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/enums2.kt");
+        }
+
+        @TestMetadata("expectActual.kt")
+        public void testExpectActual() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/expectActual.kt");
+        }
+
+        @TestMetadata("external.kt")
+        public void testExternal() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/external.kt");
+        }
+
+        @TestMetadata("F.kt")
+        public void testF() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/F.kt");
+        }
+
+        @TestMetadata("functionTypes.kt")
+        public void testFunctionTypes() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/functionTypes.kt");
+        }
+
+        @TestMetadata("genericFunctions.kt")
+        public void testGenericFunctions() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericFunctions.kt");
+        }
+
+        @TestMetadata("genericProperty.kt")
+        public void testGenericProperty() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/genericProperty.kt");
+        }
+
+        @TestMetadata("initBlockWithDeclarations.kt")
+        public void testInitBlockWithDeclarations() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/initBlockWithDeclarations.kt");
+        }
+
+        @TestMetadata("nestedClass.kt")
+        public void testNestedClass() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/nestedClass.kt");
+        }
+
+        @TestMetadata("NestedOfAliasedType.kt")
+        public void testNestedOfAliasedType() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/NestedOfAliasedType.kt");
+        }
+
+        @TestMetadata("NestedSuperType.kt")
+        public void testNestedSuperType() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/NestedSuperType.kt");
+        }
+
+        @TestMetadata("noPrimaryConstructor.kt")
+        public void testNoPrimaryConstructor() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/noPrimaryConstructor.kt");
+        }
+
+        @TestMetadata("simpleClass.kt")
+        public void testSimpleClass() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/simpleClass.kt");
+        }
+
+        @TestMetadata("simpleFun.kt")
+        public void testSimpleFun() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/simpleFun.kt");
+        }
+
+        @TestMetadata("simpleTypeAlias.kt")
+        public void testSimpleTypeAlias() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/simpleTypeAlias.kt");
+        }
+
+        @TestMetadata("typeAliasWithGeneric.kt")
+        public void testTypeAliasWithGeneric() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/typeAliasWithGeneric.kt");
+        }
+
+        @TestMetadata("typeParameterVsNested.kt")
+        public void testTypeParameterVsNested() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/typeParameterVsNested.kt");
+        }
+
+        @TestMetadata("typeParameters.kt")
+        public void testTypeParameters() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/typeParameters.kt");
+        }
+
+        @TestMetadata("where.kt")
+        public void testWhere() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/where.kt");
+        }
+
+        @TestMetadata("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class Contracts extends AbstractFirLazyBodiesCalculatorTest {
+            private void runTest(String testDataFilePath) throws Exception {
+                KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+            }
+
+            public void testAllFilesPresentInContracts() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts"), Pattern.compile("^(.+)\\.kt$"), null, true);
+            }
+
+            @TestMetadata("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/newSyntax")
+            @TestDataPath("$PROJECT_ROOT")
+            @RunWith(JUnit3RunnerWithInners.class)
+            public static class NewSyntax extends AbstractFirLazyBodiesCalculatorTest {
+                private void runTest(String testDataFilePath) throws Exception {
+                    KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                }
+
+                public void testAllFilesPresentInNewSyntax() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/newSyntax"), Pattern.compile("^(.+)\\.kt$"), null, true);
+                }
+
+                @TestMetadata("functionWithBothOldAndNewSyntaxContractDescription.kt")
+                public void testFunctionWithBothOldAndNewSyntaxContractDescription() throws Exception {
+                    runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/newSyntax/functionWithBothOldAndNewSyntaxContractDescription.kt");
+                }
+
+                @TestMetadata("propertyAccessorsContractDescription.kt")
+                public void testPropertyAccessorsContractDescription() throws Exception {
+                    runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/newSyntax/propertyAccessorsContractDescription.kt");
+                }
+
+                @TestMetadata("simpleFunctionsContractDescription.kt")
+                public void testSimpleFunctionsContractDescription() throws Exception {
+                    runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/newSyntax/simpleFunctionsContractDescription.kt");
+                }
+            }
+
+            @TestMetadata("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/oldSyntax")
+            @TestDataPath("$PROJECT_ROOT")
+            @RunWith(JUnit3RunnerWithInners.class)
+            public static class OldSyntax extends AbstractFirLazyBodiesCalculatorTest {
+                private void runTest(String testDataFilePath) throws Exception {
+                    KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+                }
+
+                public void testAllFilesPresentInOldSyntax() throws Exception {
+                    KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/oldSyntax"), Pattern.compile("^(.+)\\.kt$"), null, true);
+                }
+
+                @TestMetadata("contractDescription.kt")
+                public void testContractDescription() throws Exception {
+                    runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/contracts/oldSyntax/contractDescription.kt");
+                }
+            }
+        }
+    }
+
+    @TestMetadata("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class Expressions extends AbstractFirLazyBodiesCalculatorTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInExpressions() throws Exception {
+            KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions"), Pattern.compile("^(.+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("annotated.kt")
+        public void testAnnotated() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/annotated.kt");
+        }
+
+        @TestMetadata("arrayAccess.kt")
+        public void testArrayAccess() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/arrayAccess.kt");
+        }
+
+        @TestMetadata("arrayAssignment.kt")
+        public void testArrayAssignment() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/arrayAssignment.kt");
+        }
+
+        @TestMetadata("branches.kt")
+        public void testBranches() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/branches.kt");
+        }
+
+        @TestMetadata("callableReferences.kt")
+        public void testCallableReferences() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/callableReferences.kt");
+        }
+
+        @TestMetadata("calls.kt")
+        public void testCalls() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/calls.kt");
+        }
+
+        @TestMetadata("classReference.kt")
+        public void testClassReference() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/classReference.kt");
+        }
+
+        @TestMetadata("collectionLiterals.kt")
+        public void testCollectionLiterals() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/collectionLiterals.kt");
+        }
+
+        @TestMetadata("destructuring.kt")
+        public void testDestructuring() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/destructuring.kt");
+        }
+
+        @TestMetadata("for.kt")
+        public void testFor() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/for.kt");
+        }
+
+        @TestMetadata("genericCalls.kt")
+        public void testGenericCalls() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/genericCalls.kt");
+        }
+
+        @TestMetadata("in.kt")
+        public void testIn() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/in.kt");
+        }
+
+        @TestMetadata("inBrackets.kt")
+        public void testInBrackets() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/inBrackets.kt");
+        }
+
+        @TestMetadata("init.kt")
+        public void testInit() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/init.kt");
+        }
+
+        @TestMetadata("labelForInfix.kt")
+        public void testLabelForInfix() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/labelForInfix.kt");
+        }
+
+        @TestMetadata("lambda.kt")
+        public void testLambda() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambda.kt");
+        }
+
+        @TestMetadata("lambdaAndAnonymousFunction.kt")
+        public void testLambdaAndAnonymousFunction() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/lambdaAndAnonymousFunction.kt");
+        }
+
+        @TestMetadata("locals.kt")
+        public void testLocals() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/locals.kt");
+        }
+
+        @TestMetadata("modifications.kt")
+        public void testModifications() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/modifications.kt");
+        }
+
+        @TestMetadata("namedArgument.kt")
+        public void testNamedArgument() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/namedArgument.kt");
+        }
+
+        @TestMetadata("nullability.kt")
+        public void testNullability() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/nullability.kt");
+        }
+
+        @TestMetadata("qualifierWithTypeArguments.kt")
+        public void testQualifierWithTypeArguments() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/qualifierWithTypeArguments.kt");
+        }
+
+        @TestMetadata("simpleReturns.kt")
+        public void testSimpleReturns() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/simpleReturns.kt");
+        }
+
+        @TestMetadata("super.kt")
+        public void testSuper() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/super.kt");
+        }
+
+        @TestMetadata("these.kt")
+        public void testThese() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/these.kt");
+        }
+
+        @TestMetadata("try.kt")
+        public void testTry() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/try.kt");
+        }
+
+        @TestMetadata("typeOperators.kt")
+        public void testTypeOperators() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/typeOperators.kt");
+        }
+
+        @TestMetadata("unary.kt")
+        public void testUnary() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/unary.kt");
+        }
+
+        @TestMetadata("variables.kt")
+        public void testVariables() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/variables.kt");
+        }
+
+        @TestMetadata("while.kt")
+        public void testWhile() throws Exception {
+            runTest("compiler/fir/raw-fir/psi2fir/testData/rawBuilder/expressions/while.kt");
+        }
+    }
+}