RawFirBuilder refactoring

Second step for KT-52615 Get rid of PsiHandlingMode Get rid of source in FirLazyBlock Refactor lazy creation Merge-request: KT-MR-7753 Merged-by: Egor Kulikov <Egor.Kulikov@jetbrains.com>
2022-11-22 14:43:53 +00:00
parent 27949bccf6
commit 067036c162
18 changed files with 96 additions and 126 deletions
@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecific
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClass
@@ -180,7 +179,6 @@ object KtDeclarationAndFirDeclarationEqualityChecker {
        return RawFirBuilder(
            createEmptySession(),
            DummyScopeProvider,
-            psiMode = PsiHandlingMode.IDE,
            bodyBuildingMode = BodyBuildingMode.NORMAL
        ).buildTypeReference(this)
    }
@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.analysis.api.impl.barebone.annotations.ThreadSafe
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirPhaseRunner
@@ -26,7 +25,6 @@ internal class LLFirFileBuilder(
        RawFirBuilder(
            moduleComponents.session,
            moduleComponents.scopeProvider,
-            psiMode = PsiHandlingMode.IDE,
            bodyBuildingMode = BodyBuildingMode.LAZY_BODIES
        ).buildFirFile(ktFile)
    }
@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve

 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -23,7 +22,7 @@ internal fun buildFileFirAnnotation(

    val replacementApplier = replacement?.Applier()

-    val builder = object : RawFirBuilder(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE) {
+    val builder = object : RawFirBuilder(session, baseScopeProvider) {
        inner class VisitorWithReplacement : Visitor() {
            override fun convertElement(element: KtElement): FirElement? =
                super.convertElement(replacementApplier?.tryReplace(element) ?: element)
@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDeclarationDesignation
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
@@ -30,7 +29,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
    private val declarationToBuild: KtDeclaration,
    private val functionsToRebind: Set<FirFunction>? = null,
    private val replacementApplier: RawFirReplacement.Applier? = null
-) : RawFirBuilder(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE, bodyBuildingMode = BodyBuildingMode.NORMAL) {
+) : RawFirBuilder(session, baseScopeProvider, bodyBuildingMode = BodyBuildingMode.NORMAL) {

    companion object {
        fun buildWithReplacement(
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve

 import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
 import org.jetbrains.kotlin.fir.types.FirUserTypeRef
@@ -17,7 +16,7 @@ internal fun buildFirUserTypeRef(
    session: FirSession,
    baseScopeProvider: FirScopeProvider
 ): FirUserTypeRef {
-    val builder = object : RawFirBuilder(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE) {
+    val builder = object : RawFirBuilder(session, baseScopeProvider) {
        fun build(): FirUserTypeRef = Visitor().visitTypeReference(typeReference, Unit) as FirUserTypeRef
    }
    builder.context.packageFqName = typeReference.containingKtFile.packageFqName
@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodie
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression
@@ -38,7 +37,6 @@ abstract class AbstractFirLazyBodiesCalculatorTest : AbstractLowLevelApiSingleFi
            val laziedFirFile = RawFirBuilder(
                session,
                provider,
-                psiMode = PsiHandlingMode.IDE,
                bodyBuildingMode = BodyBuildingMode.LAZY_BODIES
            ).buildFirFile(ktFile)

@@ -48,7 +46,6 @@ abstract class AbstractFirLazyBodiesCalculatorTest : AbstractLowLevelApiSingleFi
            val fullFirFile = RawFirBuilder(
                session,
                provider,
-                psiMode = PsiHandlingMode.IDE,
                bodyBuildingMode = BodyBuildingMode.NORMAL
            ).buildFirFile(ktFile)

@@ -40,10 +40,12 @@ abstract class AbstractFirLazyDeclarationResolveTest : AbstractLowLevelApiSingle
            var result: FirDeclaration? = null
            override fun visitElement(element: FirElement) {
                if (result != null) return
-                val declaration = element.realPsi as? KtDeclaration
-                if (element is FirDeclaration && declaration != null && declaration.name?.decapitalizeAsciiOnly() == "resolveMe") {
-                    result = element
-                    return
+                if (element is FirDeclaration) {
+                    val declaration = element.realPsi as? KtDeclaration
+                    if (declaration != null && declaration.name?.decapitalizeAsciiOnly() == "resolveMe") {
+                        result = element
+                        return
+                    }
                }
                element.acceptChildren(this)
            }
@@ -11,7 +11,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.RawFirNonLoc
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
@@ -111,7 +110,7 @@ abstract class AbstractPartialRawFirBuilderTestCase : AbstractLowLevelApiSingleF
        }

        val session = FirSessionFactoryHelper.createEmptySession()
-        val firBuilder = RawFirBuilder(session, scopeProvider, PsiHandlingMode.IDE)
+        val firBuilder = RawFirBuilder(session, scopeProvider)
        val original = firBuilder.buildFirFile(file)

        val designationBuilder = DesignationBuilder(elementToBuild)
@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.context.withProject
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.FirTestSessionFactoryHelper
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.java.FirJavaElementFinder
 import org.jetbrains.kotlin.fir.resolve.providers.firProvider
@@ -157,7 +156,7 @@ abstract class AbstractSimpleFileBenchmark {
            .uniteWith(TopDownAnalyzerFacadeForJVM.AllJavaSourcesInProjectScope(env.project))
        val session = FirTestSessionFactoryHelper.createSessionForTests(env.toAbstractProjectEnvironment(), scope.toAbstractProjectFileSearchScope())
        val firProvider = session.firProvider as FirProviderImpl
-        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
+        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)

        val totalTransformer = FirTotalResolveProcessor(session)
        val firFile = builder.buildFirFile(file).also(firProvider::recordFile)
@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.pipeline
 import org.jetbrains.kotlin.KtSourceFile
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
@@ -46,7 +45,7 @@ fun FirSession.buildFirViaLightTree(

 fun FirSession.buildFirFromKtFiles(ktFiles: Collection<KtFile>): List<FirFile> {
    val firProvider = (firProvider as FirProviderImpl)
-    val builder = RawFirBuilder(this, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
+    val builder = RawFirBuilder(this, firProvider.kotlinScopeProvider)
    return ktFiles.map {
        builder.buildFirFile(it).also { firFile ->
            firProvider.recordFile(firFile)
@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollector
 import org.jetbrains.kotlin.fir.analysis.collectors.FirDiagnosticsCollector
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.dump.MultiModuleHtmlFirDump
@@ -162,7 +161,7 @@ class FirResolveModularizedTotalKotlinTest : AbstractModularizedTest() {
            }
            bench.buildFiles(lightTree2Fir, allSourceFiles)
        } else {
-            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
+            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
            bench.buildFiles(builder, ktFiles)
        }

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.builder

 import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.IElementType
+import com.intellij.util.AstLoadingFilter
 import org.jetbrains.kotlin.*
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.StandardNames.BACKING_FIELD
@@ -51,14 +52,8 @@ import org.jetbrains.kotlin.utils.addToStdlib.runIf
 open class RawFirBuilder(
    session: FirSession,
    val baseScopeProvider: FirScopeProvider,
-    private val psiMode: PsiHandlingMode,
    bodyBuildingMode: BodyBuildingMode = BodyBuildingMode.NORMAL
 ) : BaseFirBuilder<PsiElement>(session) {
-
-    @Deprecated("Please replace with primary constructor call")
-    constructor(session: FirSession, baseScopeProvider: FirScopeProvider, mode: BodyBuildingMode = BodyBuildingMode.NORMAL) :
-            this(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE, bodyBuildingMode = mode)
-
    protected open fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction) = target.bind(function)

    var mode: BodyBuildingMode = bodyBuildingMode
@@ -74,6 +69,15 @@ open class RawFirBuilder(
        }
    }

+    private inline fun <T> runOnStubs(crossinline body: () -> T): T {
+        return when (mode) {
+            BodyBuildingMode.NORMAL -> body()
+            BodyBuildingMode.LAZY_BODIES -> {
+                AstLoadingFilter.disallowTreeLoading<T, Nothing> { body() }
+            }
+        }
+    }
+
    fun buildFirFile(file: KtFile): FirFile {
        return file.accept(Visitor(), Unit) as FirFile
    }
@@ -165,6 +169,28 @@ open class RawFirBuilder(
        private inline fun <reified R : FirElement> KtElement.convert(): R =
            convertElement(this) as R

+        private inline fun <T> buildOrLazy(build: () -> T, noinline lazy: () -> T): T {
+            return when (mode) {
+                BodyBuildingMode.NORMAL -> build()
+                BodyBuildingMode.LAZY_BODIES -> runOnStubs(lazy)
+            }
+        }
+
+        private inline fun buildOrLazyExpression(
+            sourceElement: KtSourceElement?,
+            buildExpression: () -> FirExpression,
+        ): FirExpression {
+            return buildOrLazy(buildExpression, {
+                buildLazyExpression {
+                    source = sourceElement
+                }
+            })
+        }
+
+        private inline fun buildOrLazyBlock(buildBlock: () -> FirBlock): FirBlock {
+            return buildOrLazy(buildBlock, ::buildLazyBlock)
+        }
+
        open fun convertElement(element: KtElement): FirElement? =
            element.accept(this@Visitor, Unit)

@@ -317,28 +343,22 @@ open class RawFirBuilder(
            }

        private fun KtDeclarationWithBody.buildFirBody(): Pair<FirBlock?, FirContractDescription?> =
-            when {
-                !hasBody() ->
-                    null to null
-                mode == BodyBuildingMode.LAZY_BODIES -> {
-                    val block = buildLazyBlock {
-                        source = bodyExpression?.toFirSourceElement()
-                            ?: error("hasBody() == true but body is null")
-                    }
-                    block to null
-                }
-                hasBlockBody() -> {
-                    val block = bodyBlockExpression?.accept(this@Visitor, Unit) as? FirBlock
-                    if (hasContractEffectList()) {
-                        block to null
+            if (hasBody()) {
+                buildOrLazyBlock {
+                    if (hasBlockBody()) {
+                        val block = bodyBlockExpression?.accept(this@Visitor, Unit) as? FirBlock
+                        return@buildFirBody if (hasContractEffectList()) {
+                            block to null
+                        } else {
+                            block.extractContractDescriptionIfPossible()
+                        }
                    } else {
-                        block.extractContractDescriptionIfPossible()
+                        val result = { bodyExpression }.toFirExpression("Function has no body (but should)")
+                        FirSingleExpressionBlock(result.toReturn(baseSource = result.source))
                    }
-                }
-                else -> {
-                    val result = { bodyExpression }.toFirExpression("Function has no body (but should)")
-                    FirSingleExpressionBlock(result.toReturn(baseSource = result.source)) to null
-                }
+                } to null
+            } else {
+                null to null
            }

        private fun ValueArgument?.toFirExpression(): FirExpression {
@@ -497,13 +517,7 @@ open class RawFirBuilder(
                Visibilities.Private
            }
            val status = obtainPropertyComponentStatus(componentVisibility, this, property)
-            val backingFieldInitializer = when {
-                mode == BodyBuildingMode.LAZY_BODIES -> buildLazyExpression {
-                    source = this@toFirBackingField?.initializer?.toFirSourceElement()
-                }
-                this?.hasInitializer() != true -> null
-                else -> this@toFirBackingField.initializer?.toFirExpression("Should have initializer")
-            }
+            val backingFieldInitializer = this?.toInitializerExpression()
            val returnType = this?.returnTypeReference.toFirOrImplicitType()
            val source = this?.toFirSourceElement()
            return if (this != null) {
@@ -950,7 +964,10 @@ open class RawFirBuilder(
            if (hasModifier(INNER_KEYWORD)) dispatchReceiverForInnerClassConstructor() else null

        override fun visitKtFile(file: KtFile, data: Unit): FirElement {
-            context.packageFqName = if (psiMode == PsiHandlingMode.COMPILER) file.packageFqNameByTree else file.packageFqName
+            context.packageFqName = when (mode) {
+                BodyBuildingMode.NORMAL -> file.packageFqNameByTree
+                BodyBuildingMode.LAZY_BODIES -> file.packageFqName
+            }
            return buildFile {
                source = file.toFirSourceElement()
                moduleData = baseModuleData
@@ -1573,6 +1590,13 @@ open class RawFirBuilder(
            }
        }

+        private fun KtDeclarationWithInitializer.toInitializerExpression() =
+            runIf (hasInitializer()) {
+                buildOrLazyExpression(initializer?.toFirSourceElement()) {
+                    initializer.toFirExpression("Should have initializer")
+                }
+            }
+
        private fun <T> KtProperty.toFirProperty(
            ownerRegularOrAnonymousObjectSymbol: FirClassSymbol<*>?,
            ownerRegularClassTypeParametersCount: Int?,
@@ -1581,14 +1605,7 @@ open class RawFirBuilder(
            val propertyType = typeReference.toFirOrImplicitType()
            val propertyName = nameAsSafeName
            val isVar = isVar
-            val propertyInitializer = when {
-                !hasInitializer() -> null
-                mode == BodyBuildingMode.LAZY_BODIES -> buildLazyExpression {
-                    source = initializer?.toFirSourceElement()
-                }
-                else -> initializer.toFirExpression("Should have initializer")
-
-            }
+            val propertyInitializer = toInitializerExpression()

            val propertySource = toFirSourceElement()

@@ -1680,14 +1697,11 @@ open class RawFirBuilder(
                        }

                        if (hasDelegate()) {
-                            fun extractDelegateExpression() = when (mode) {
-                                BodyBuildingMode.NORMAL -> this@toFirProperty.delegate?.expression.toFirExpression("Should have delegate")
-                                BodyBuildingMode.LAZY_BODIES -> this@toFirProperty.delegate?.expression?.let {
-                                    buildLazyExpression { source = it.toFirSourceElement() }
-                                } ?: buildErrorExpression {
-                                    ConeSimpleDiagnostic("Should have delegate", DiagnosticKind.ExpressionExpected)
+                            fun extractDelegateExpression() = this@toFirProperty.delegate?.expression?.let { expression ->
+                                buildOrLazyExpression(expression.toFirSourceElement()) {
+                                    expression.toFirExpression("Should have delegate")
                                }
-                            }
+                            } ?: buildErrorExpression { ConeSimpleDiagnostic("Should have delegate", DiagnosticKind.ExpressionExpected) }

                            val delegateBuilder = FirWrappedDelegateExpressionBuilder().apply {
                                val delegateExpression = extractDelegateExpression()
@@ -1905,7 +1919,8 @@ open class RawFirBuilder(
                source = expression.toFirSourceElement()
                for (statement in expression.statements) {
                    val firStatement = statement.toFirStatement { "Statement expected: ${statement.text}" }
-                    val isForLoopBlock = firStatement is FirBlock && firStatement.source?.kind == KtFakeSourceElementKind.DesugaredForLoop
+                    val isForLoopBlock =
+                        firStatement is FirBlock && firStatement.source?.kind == KtFakeSourceElementKind.DesugaredForLoop
                    if (firStatement !is FirBlock || isForLoopBlock || firStatement.annotations.isNotEmpty()) {
                        statements += firStatement
                    } else {
@@ -2500,7 +2515,10 @@ open class RawFirBuilder(
        override fun visitParenthesizedExpression(expression: KtParenthesizedExpression, data: Unit): FirElement {
            context.forwardLabelUsagePermission(expression, expression.expression)
            return expression.expression?.accept(this, data)
-                ?: buildErrorExpression(expression.toFirSourceElement(), ConeSimpleDiagnostic("Empty parentheses", DiagnosticKind.Syntax))
+                ?: buildErrorExpression(
+                    expression.toFirSourceElement(),
+                    ConeSimpleDiagnostic("Empty parentheses", DiagnosticKind.Syntax)
+                )
        }

        override fun visitLabeledExpression(expression: KtLabeledExpression, data: Unit): FirElement {
@@ -2546,7 +2564,10 @@ open class RawFirBuilder(
                baseModuleData,
                multiDeclaration.toFirSourceElement(),
                "destruct",
-                multiDeclaration.initializer.toFirExpression("Initializer required for destructuring declaration", DiagnosticKind.Syntax),
+                multiDeclaration.initializer.toFirExpression(
+                    "Initializer required for destructuring declaration",
+                    DiagnosticKind.Syntax
+                ),
                extractAnnotationsTo = { extractAnnotationsTo(it) }
            )
            return generateDestructuringBlock(
@@ -2621,15 +2642,3 @@ enum class BodyBuildingMode {
            if (lazyBodies) LAZY_BODIES else NORMAL
    }
 }
-
-enum class PsiHandlingMode {
-    /**
-     * Do not build any stubs while handling PSI
-     */
-    COMPILER,
-
-    /**
-     * Build stubs if possible while handling PSI
-     */
-    IDE;
-}
@@ -84,7 +84,6 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
        return RawFirBuilder(
            session,
            StubFirScopeProvider,
-            psiMode = PsiHandlingMode.COMPILER,
            bodyBuildingMode = bodyBuildingMode
        ).buildFirFile(this)
    }
@@ -25,33 +25,7 @@ import org.jetbrains.kotlin.fir.visitors.*
 * DO NOT MODIFY IT MANUALLY
 */

-@FirBuilderDsl
-class FirLazyBlockBuilder : FirAnnotationContainerBuilder, FirExpressionBuilder {
-    override var source: KtSourceElement? = null
-
-    @OptIn(FirImplementationDetail::class)
-    override fun build(): FirBlock {
-        return FirLazyBlock(
-            source,
-        )
-    }
-
-
-    @Deprecated("Modification of 'annotations' has no impact for FirLazyBlockBuilder", level = DeprecationLevel.HIDDEN)
-    override val annotations: MutableList<FirAnnotation> = mutableListOf()
-
-    @Deprecated("Modification of 'typeRef' has no impact for FirLazyBlockBuilder", level = DeprecationLevel.HIDDEN)
-    override var typeRef: FirTypeRef
-        get() = throw IllegalStateException()
-        set(_) {
-            throw IllegalStateException()
-        }
-}
-
-@OptIn(ExperimentalContracts::class)
-inline fun buildLazyBlock(init: FirLazyBlockBuilder.() -> Unit = {}): FirBlock {
-    contract {
-        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
-    }
-    return FirLazyBlockBuilder().apply(init).build()
+@OptIn(FirImplementationDetail::class)
+fun buildLazyBlock(): FirBlock {
+    return FirLazyBlock()
 }
@@ -20,9 +20,8 @@ import org.jetbrains.kotlin.fir.visitors.*
 * DO NOT MODIFY IT MANUALLY
 */

-class FirLazyBlock @FirImplementationDetail constructor(
-    override val source: KtSourceElement?,
-) : FirBlock() {
+class FirLazyBlock : FirBlock() {
+    override val source: KtSourceElement? get() = error("FirLazyBlock should be calculated before accessing")
    override val annotations: List<FirAnnotation> get() = error("FirLazyBlock should be calculated before accessing")
    override val statements: List<FirStatement> get() = error("FirLazyBlock should be calculated before accessing")
    override val typeRef: FirTypeRef get() = error("FirLazyBlock should be calculated before accessing")
@@ -142,6 +142,10 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()

        impl(block, "FirLazyBlock") {
            val error = """error("FirLazyBlock should be calculated before accessing")"""
+            default("source") {
+                value = error
+                withGetter = true
+            }
            default("statements") {
                value = error
                withGetter = true
@@ -30,7 +30,6 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.KtDiagnostic
 import org.jetbrains.kotlin.diagnostics.PsiDiagnosticUtils
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
@@ -142,7 +141,6 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
            val firBuilder = RawFirBuilder(
                session,
                firProvider.kotlinScopeProvider,
-                psiMode = if (useLazyBodiesModeForRawFir) PsiHandlingMode.IDE else PsiHandlingMode.COMPILER,
                bodyBuildingMode = BodyBuildingMode.lazyBodies(useLazyBodiesModeForRawFir)
            )
            ktFiles.mapTo(firFiles) {
@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.backend.Fir2IrResult
 import org.jetbrains.kotlin.fir.backend.jvm.Fir2IrJvmSpecialAnnotationSymbolProvider
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmKotlinMangler
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmVisibilityConverter
-import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
@@ -69,7 +68,7 @@ class FirAnalyzerFacade(
                }
            }
        } else {
-            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
+            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
            ktFiles.map {
                builder.buildFirFile(it).also { firFile ->
                    firProvider.recordFile(firFile)