RawFirBuilder: optimize package name calculation in compiler mode

benchmarks/src/org/jetbrains/kotlin/benchmarks/AbstractSimpleFileBenchmark.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.context.withModule
 import org.jetbrains.kotlin.context.withProject
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.Severity
+import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.createSessionForTests
 import org.jetbrains.kotlin.fir.java.FirJavaElementFinder
@@ -154,7 +155,7 @@ abstract class AbstractSimpleFileBenchmark {
             .uniteWith(TopDownAnalyzerFacadeForJVM.AllJavaSourcesInProjectScope(env.project))
         val session = createSessionForTests(env, scope)
         val firProvider = session.firProvider as FirProviderImpl
-        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
+        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
 
         val totalTransformer = FirTotalResolveProcessor(session)
         val firFile = builder.buildFirFile(file).also(firProvider::recordFile)

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/analysis/FirAnalyzerFacade.kt

@@ -16,6 +16,7 @@ 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
@@ -63,7 +64,7 @@ class FirAnalyzerFacade(
                 }
             }
         } else {
-            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
+            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
             ktFiles.map {
                 builder.buildFirFile(it).also { firFile ->
                     firProvider.recordFile(firFile)

compiler/fir/modularized-tests/tests/org/jetbrains/kotlin/fir/FirResolveModularizedTotalKotlinTest.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.fir.analysis.FirCheckersResolveProcessor
+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
@@ -158,7 +159,7 @@ class FirResolveModularizedTotalKotlinTest : AbstractModularizedTest() {
             }
             bench.buildFiles(lightTree2Fir, allSourceFiles)
         } else {
-            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
+            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, PsiHandlingMode.COMPILER)
             bench.buildFiles(builder, ktFiles)
         }
 

compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/SimpleTestCase.kt

@@ -13,7 +13,7 @@ import com.intellij.testFramework.LightVirtualFile
 import com.intellij.testFramework.TestDataPath
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.builder.AbstractRawFirBuilderTestCase
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.StubFirScopeProvider
 import org.jetbrains.kotlin.fir.lightTree.converter.DeclarationsConverter
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
@@ -73,7 +73,7 @@ class SimpleTestCase : AbstractRawFirBuilderTestCase() {
         println(StringBuilder().also { FirRenderer(it).visitFile(firFromLightTreeFile) }.toString())
 
         val psiFile = createPsiFile("foo", code) as KtFile
-        val firFromPsiFile = psiFile.toFirFile(RawFirBuilderMode.STUBS)
+        val firFromPsiFile = psiFile.toFirFile(BodyBuildingMode.STUBS)
         println("Fir from PSI")
         println(StringBuilder().also { FirRenderer(it).visitFile(firFromPsiFile) }.toString())
 

compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/TotalKotlinTest.kt

@@ -12,7 +12,7 @@ import com.intellij.testFramework.TestDataPath
 import com.intellij.util.PathUtil
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.builder.AbstractRawFirBuilderTestCase
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.StubFirScopeProvider
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
 import org.jetbrains.kotlin.psi.KtFile
@@ -29,7 +29,7 @@ class TotalKotlinTest : AbstractRawFirBuilderTestCase() {
         if (onlyPsi) {
             DebugUtil.psiTreeToString(ktFile, false)
         } else {
-            val firFile = ktFile.toFirFile(RawFirBuilderMode.STUBS)
+            val firFile = ktFile.toFirFile(BodyBuildingMode.STUBS)
             StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
         }
     }

compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/benchmark/generators/Psi2FirGenerator.kt

@@ -10,7 +10,7 @@ import com.intellij.psi.impl.DebugUtil
 import com.intellij.util.PathUtil
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.builder.AbstractRawFirBuilderTestCase
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.psi.KtFile
 import org.openjdk.jmh.annotations.Scope
 import org.openjdk.jmh.annotations.State
@@ -25,7 +25,7 @@ open class Psi2FirGenerator : TreeGenerator, AbstractRawFirBuilderTestCase() {
 
     override fun generateFir(text: String, file: File, stubMode: Boolean) {
         val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(file.path)), text) as KtFile
-        val firFile = ktFile.toFirFile(RawFirBuilderMode.stubs(stubMode))
+        val firFile = ktFile.toFirFile(BodyBuildingMode.stubs(stubMode))
         StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
     }
 

compiler/fir/raw-fir/light-tree2fir/tests/org/jetbrains/kotlin/fir/lightTree/compare/TreesCompareTest.kt

@@ -13,7 +13,7 @@ import junit.framework.TestCase
 import org.jetbrains.kotlin.checkers.BaseDiagnosticsTest.Companion.DIAGNOSTIC_IN_TESTDATA_PATTERN
 import org.jetbrains.kotlin.fir.FirRenderer
 import org.jetbrains.kotlin.fir.builder.AbstractRawFirBuilderTestCase
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.StubFirScopeProvider
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
 import org.jetbrains.kotlin.fir.lightTree.walkTopDown
@@ -66,7 +66,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
 
             //psi
             val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(file.path)), text) as KtFile
-            val firFileFromPsi = ktFile.toFirFile(RawFirBuilderMode.stubs(stubMode))
+            val firFileFromPsi = ktFile.toFirFile(BodyBuildingMode.stubs(stubMode))
             val treeFromPsi = StringBuilder().also { FirRenderer(it).visitFile(firFileFromPsi) }.toString()
 
             //light tree

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

@@ -48,23 +48,30 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 open class RawFirBuilder(
-    session: FirSession, val baseScopeProvider: FirScopeProvider, builderMode: RawFirBuilderMode = RawFirBuilderMode.NORMAL
+    session: FirSession,
+    val baseScopeProvider: FirScopeProvider,
+    private val psiMode: PsiHandlingMode,
+    bodyBuildingMode: BodyBuildingMode = BodyBuildingMode.NORMAL
 ) : BaseFirBuilder<PsiElement>(session) {
 
-    private val stubMode get() = mode == RawFirBuilderMode.STUBS
+    @Deprecated("Please replace with primary constructor call")
+    constructor(session: FirSession, baseScopeProvider: FirScopeProvider, mode: BodyBuildingMode = BodyBuildingMode.NORMAL) :
+            this(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE, bodyBuildingMode = mode)
+
+    private val stubMode get() = mode == BodyBuildingMode.STUBS
 
     protected open fun bindFunctionTarget(target: FirFunctionTarget, function: FirFunction) = target.bind(function)
 
-    var mode: RawFirBuilderMode = builderMode
+    var mode: BodyBuildingMode = bodyBuildingMode
         private set
 
     private inline fun <T> disabledLazyMode(body: () -> T): T {
-        if (mode != RawFirBuilderMode.LAZY_BODIES) return body()
+        if (mode != BodyBuildingMode.LAZY_BODIES) return body()
         return try {
-            mode = RawFirBuilderMode.NORMAL
+            mode = BodyBuildingMode.NORMAL
             body()
         } finally {
-            mode = RawFirBuilderMode.LAZY_BODIES
+            mode = BodyBuildingMode.LAZY_BODIES
         }
     }
 
@@ -289,7 +296,7 @@ open class RawFirBuilder(
             when {
                 !hasBody() ->
                     null to null
-                mode == RawFirBuilderMode.LAZY_BODIES -> {
+                mode == BodyBuildingMode.LAZY_BODIES -> {
                     val block = buildLazyBlock {
                         source = bodyExpression?.toFirSourceElement()
                             ?: error("hasBody() == true but body is null")
@@ -734,7 +741,7 @@ open class RawFirBuilder(
         }
 
         override fun visitKtFile(file: KtFile, data: Unit): FirElement {
-            context.packageFqName = file.packageFqName
+            context.packageFqName = if (psiMode == PsiHandlingMode.COMPILER) file.packageFqNameByTree else file.packageFqName
             return buildFile {
                 source = file.toFirSourceElement()
                 moduleData = baseModuleData
@@ -1317,10 +1324,10 @@ open class RawFirBuilder(
             val isVar = isVar
             val propertyInitializer = when {
                 !hasInitializer() -> null
-                mode == RawFirBuilderMode.LAZY_BODIES -> buildLazyExpression {
+                mode == BodyBuildingMode.LAZY_BODIES -> buildLazyExpression {
                     source = initializer?.toFirSourceElement()
                 }
-                mode == RawFirBuilderMode.STUBS -> buildExpressionStub()
+                mode == BodyBuildingMode.STUBS -> buildExpressionStub()
                 else -> initializer.toFirExpression("Should have initializer")
 
             }
@@ -1373,9 +1380,9 @@ open class RawFirBuilder(
                                 source =
                                     if (stubMode) null else delegateExpression?.toFirSourceElement(FirFakeSourceElementKind.WrappedDelegate)
                                 expression = when (mode) {
-                                    RawFirBuilderMode.NORMAL -> delegateExpression.toFirExpression("Should have delegate")
-                                    RawFirBuilderMode.STUBS -> buildExpressionStub()
-                                    RawFirBuilderMode.LAZY_BODIES -> buildLazyExpression {
+                                    BodyBuildingMode.NORMAL -> delegateExpression.toFirExpression("Should have delegate")
+                                    BodyBuildingMode.STUBS -> buildExpressionStub()
+                                    BodyBuildingMode.LAZY_BODIES -> buildLazyExpression {
                                         source = delegateExpression!!.toFirSourceElement()
                                     }
                                 }
@@ -2284,7 +2291,7 @@ open class RawFirBuilder(
     }
 }
 
-enum class RawFirBuilderMode {
+enum class BodyBuildingMode {
     /**
      * Build every expression and every body
      */
@@ -2302,10 +2309,25 @@ enum class RawFirBuilderMode {
     LAZY_BODIES;
 
     companion object {
-        fun lazyBodies(lazyBodies: Boolean): RawFirBuilderMode =
+        fun lazyBodies(lazyBodies: Boolean): BodyBuildingMode =
             if (lazyBodies) LAZY_BODIES else NORMAL
 
-        fun stubs(stubs: Boolean): RawFirBuilderMode =
+        fun stubs(stubs: Boolean): BodyBuildingMode =
             if (stubs) STUBS else NORMAL
     }
 }
+
+@Deprecated("Please replace with BodyBuildingMode")
+typealias RawFirBuilderMode = BodyBuildingMode
+
+enum class PsiHandlingMode {
+    /**
+     * Do not build any stubs while handling PSI
+     */
+    COMPILER,
+
+    /**
+     * Build stubs if possible while handling PSI
+     */
+    IDE;
+}

compiler/fir/raw-fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderLazyBodiesTestCase.kt

@@ -12,7 +12,7 @@ import java.io.File
 abstract class AbstractRawFirBuilderLazyBodiesTestCase : AbstractRawFirBuilderTestCase() {
     override fun doRawFirTest(filePath: String) {
         val file = createKtFile(filePath)
-        val firFile = file.toFirFile(RawFirBuilderMode.LAZY_BODIES)
+        val firFile = file.toFirFile(BodyBuildingMode.LAZY_BODIES)
         val firFileDump = StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
         val expectedPath = filePath.replace(".kt", ".lazyBodies.txt")
         KotlinTestUtils.assertEqualsToFile(File(expectedPath), firFileDump)

compiler/fir/raw-fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilderTestCase.kt

@@ -62,7 +62,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
 
     protected open fun doRawFirTest(filePath: String) {
         val file = createKtFile(filePath)
-        val firFile = file.toFirFile(RawFirBuilderMode.NORMAL)
+        val firFile = file.toFirFile(BodyBuildingMode.NORMAL)
         val firFileDump = StringBuilder().also { FirRenderer(it, mode = FirRenderer.RenderMode.WithDeclarationAttributes).visitFile(firFile) }.toString()
         val expectedPath = filePath.replace(".kt", ".txt")
         KotlinTestUtils.assertEqualsToFile(File(expectedPath), firFileDump)
@@ -75,9 +75,14 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         }
     }
 
-    protected fun KtFile.toFirFile(mode: RawFirBuilderMode = RawFirBuilderMode.NORMAL): FirFile {
+    protected fun KtFile.toFirFile(bodyBuildingMode: BodyBuildingMode = BodyBuildingMode.NORMAL): FirFile {
         val session = FirSessionFactory.createEmptySession()
-        return RawFirBuilder(session, StubFirScopeProvider, mode).buildFirFile(this)
+        return RawFirBuilder(
+            session,
+            StubFirScopeProvider,
+            psiMode = PsiHandlingMode.COMPILER,
+            bodyBuildingMode = bodyBuildingMode
+        ).buildFirFile(this)
     }
 
     private fun FirElement.traverseChildren(result: MutableSet<FirElement> = hashSetOf()): MutableSet<FirElement> {

compiler/fir/raw-fir/psi2fir/tests/org/jetbrains/kotlin/fir/builder/RawFirBuilderTotalKotlinTestCase.kt

@@ -63,7 +63,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
                 val ktFile = createKtFile(file.toRelativeString(root))
                 val firFile: FirFile
                 time += measureNanoTime {
-                    firFile = ktFile.toFirFile(RawFirBuilderMode.stubs(stubMode))
+                    firFile = ktFile.toFirFile(BodyBuildingMode.stubs(stubMode))
                 }
                 totalLength += StringBuilder().also { FirRenderer(it).visitFile(firFile) }.length
                 counter++

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt

@@ -25,7 +25,8 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.PsiDiagnosticUtils
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+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.extensions.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
@@ -132,7 +133,8 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
             val firBuilder = RawFirBuilder(
                 session,
                 firProvider.kotlinScopeProvider,
-                RawFirBuilderMode.lazyBodies(useLazyBodiesModeForRawFir)
+                psiMode = if (useLazyBodiesModeForRawFir) PsiHandlingMode.IDE else PsiHandlingMode.COMPILER,
+                bodyBuildingMode = BodyBuildingMode.lazyBodies(useLazyBodiesModeForRawFir)
             )
             ktFiles.mapTo(firFiles) {
                 val firFile = firBuilder.buildFirFile(it)

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

@@ -9,7 +9,8 @@ import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
+import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -135,7 +136,8 @@ object KtDeclarationAndFirDeclarationEqualityChecker {
         return RawFirBuilder(
             createEmptySession(),
             DummyScopeProvider,
-            RawFirBuilderMode.STUBS
+            psiMode = PsiHandlingMode.IDE,
+            bodyBuildingMode = BodyBuildingMode.STUBS
         ).buildTypeReference(this)
     }
 

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

@@ -6,7 +6,8 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.builder
 
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+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.idea.fir.low.level.api.FirPhaseRunner
@@ -31,7 +32,12 @@ internal class FirFileBuilder(
         cache: ModuleFileCache,
         preferLazyBodies: Boolean
     ): FirFile = cache.fileCached(ktFile) {
-        RawFirBuilder(cache.session, scopeProvider, RawFirBuilderMode.lazyBodies(preferLazyBodies)).buildFirFile(ktFile)
+        RawFirBuilder(
+            cache.session,
+            scopeProvider,
+            psiMode = PsiHandlingMode.IDE,
+            bodyBuildingMode = BodyBuildingMode.lazyBodies(preferLazyBodies)
+        ).buildFirFile(ktFile)
     }
 }
 

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.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.FirAnnotationCall
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -22,7 +23,7 @@ internal fun buildFileFirAnnotation(
 
     val replacementApplier = replacement?.Applier()
 
-    val builder = object : RawFirBuilder(session, baseScopeProvider) {
+    val builder = object : RawFirBuilder(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE) {
         inner class VisitorWithReplacement : Visitor() {
             override fun convertElement(element: KtElement): FirElement? =
                 super.convertElement(replacementApplier?.tryReplace(element) ?: element)

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

@@ -8,7 +8,8 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.builder.RawFirBuilderMode
+import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
+import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isInner
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -24,7 +25,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, RawFirBuilderMode.NORMAL) {
+) : RawFirBuilder(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE, bodyBuildingMode = BodyBuildingMode.NORMAL) {
 
     companion object {
         fun buildWithReplacement(

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.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
@@ -16,7 +17,7 @@ internal fun buildFirUserTypeRef(
     session: FirSession,
     baseScopeProvider: FirScopeProvider
 ): FirUserTypeRef {
-    val builder = object : RawFirBuilder(session, baseScopeProvider) {
+    val builder = object : RawFirBuilder(session, baseScopeProvider, psiMode = PsiHandlingMode.IDE) {
         fun build(): FirUserTypeRef = Visitor().visitTypeReference(typeReference, Unit) as FirUserTypeRef
     }
     builder.context.packageFqName = typeReference.containingKtFile.packageFqName

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

@@ -8,7 +8,8 @@ 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.builder.BodyBuildingMode
+import org.jetbrains.kotlin.fir.builder.PsiHandlingMode
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.AbstractLowLevelApiSingleFileTest
@@ -22,9 +23,21 @@ abstract class AbstractFirLazyBodiesCalculatorTest : AbstractLowLevelApiSingleFi
             val session = resolveState.rootModuleSession
             val provider = session.firIdeProvider.kotlinScopeProvider
 
-            val laziedFirFile = RawFirBuilder(session, provider, RawFirBuilderMode.LAZY_BODIES).buildFirFile(ktFile)
+            val laziedFirFile = RawFirBuilder(
+                session,
+                provider,
+                psiMode = PsiHandlingMode.IDE,
+                bodyBuildingMode = BodyBuildingMode.LAZY_BODIES
+            ).buildFirFile(ktFile)
+
             FirLazyBodiesCalculator.calculateLazyBodies(laziedFirFile)
-            val fullFirFile = RawFirBuilder(session, provider, RawFirBuilderMode.NORMAL).buildFirFile(ktFile)
+
+            val fullFirFile = RawFirBuilder(
+                session,
+                provider,
+                psiMode = PsiHandlingMode.IDE,
+                bodyBuildingMode = BodyBuildingMode.NORMAL
+            ).buildFirFile(ktFile)
 
             val laziedFirFileDump = StringBuilder().also { FirRenderer(it).visitFile(laziedFirFile) }.toString()
             val fullFirFileDump = StringBuilder().also { FirRenderer(it).visitFile(fullFirFile) }.toString()

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import junit.framework.TestCase
 import org.jetbrains.kotlin.fir.*
+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.resolve.ScopeSession
@@ -106,7 +107,7 @@ abstract class AbstractPartialRawFirBuilderTestCase : AbstractLowLevelApiSingleF
         }
 
         val session = FirSessionFactory.createEmptySession()
-        val firBuilder = RawFirBuilder(session, scopeProvider)
+        val firBuilder = RawFirBuilder(session, scopeProvider, PsiHandlingMode.IDE)
         val original = firBuilder.buildFirFile(file)
 
         val designationBuilder = DesignationBuilder(elementToBuild)