FIR: use enum in RawFirBuilder to indicate its mode: normals, stubs, or lazy bodies

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

@@ -152,7 +152,7 @@ abstract class AbstractSimpleFileBenchmark {
             .uniteWith(TopDownAnalyzerFacadeForJVM.AllJavaSourcesInProjectScope(env.project))
         val session = createSession(env, scope)
         val firProvider = session.firProvider as FirProviderImpl
-        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)
+        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
 
         val totalTransformer = FirTotalResolveProcessor(session)
         val firFile = builder.buildFirFile(file).also(firProvider::recordFile)

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

@@ -33,7 +33,7 @@ class FirAnalyzerFacade(val session: FirSession, val languageVersionSettings: La
     private fun buildRawFir() {
         if (firFiles != null) return
         val firProvider = (session.firProvider as FirProviderImpl)
-        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)
+        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
         firFiles = ktFiles.map {
             val firFile = builder.buildFirFile(it)
             firProvider.recordFile(firFile)

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

@@ -121,7 +121,7 @@ class FirResolveModularizedTotalKotlinTest : AbstractModularizedTest() {
             }
             bench.buildFiles(lightTree2Fir, allSourceFiles)
         } else {
-            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)
+            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
             bench.buildFiles(builder, ktFiles)
         }
 

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

@@ -13,6 +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.StubFirScopeProvider
 import org.jetbrains.kotlin.fir.lightTree.converter.DeclarationsConverter
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
@@ -72,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(stubMode = true)
+        val firFromPsiFile = psiFile.toFirFile(RawFirBuilderMode.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,6 +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.StubFirScopeProvider
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
 import org.jetbrains.kotlin.psi.KtFile
@@ -28,7 +29,7 @@ class TotalKotlinTest : AbstractRawFirBuilderTestCase() {
         if (onlyPsi) {
             DebugUtil.psiTreeToString(ktFile, false)
         } else {
-            val firFile = ktFile.toFirFile(stubMode = true)
+            val firFile = ktFile.toFirFile(RawFirBuilderMode.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,6 +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.psi.KtFile
 import org.openjdk.jmh.annotations.Scope
 import org.openjdk.jmh.annotations.State
@@ -24,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(stubMode)
+        val firFile = ktFile.toFirFile(RawFirBuilderMode.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,6 +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.StubFirScopeProvider
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
 import org.jetbrains.kotlin.fir.lightTree.walkTopDown
@@ -65,7 +66,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
 
             //psi
             val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(file.path)), text) as KtFile
-            val firFileFromPsi = ktFile.toFirFile(stubMode)
+            val firFileFromPsi = ktFile.toFirFile(RawFirBuilderMode.stubs(stubMode))
             val treeFromPsi = StringBuilder().also { FirRenderer(it).visitFile(firFileFromPsi) }.toString()
 
             //light tree
@@ -91,7 +92,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
 
             //psi
             val ktFile = createPsiFile(FileUtil.getNameWithoutExtension(PathUtil.getFileName(file.path)), text) as KtFile
-            val firFileFromPsi = ktFile.toFirFile(stubMode = false)
+            val firFileFromPsi = ktFile.toFirFile()
             val treeFromPsi = StringBuilder().also { FirRenderer(it).visitFile(firFileFromPsi) }.toString()
                 .replace("<Unsupported LValue.*?>".toRegex(), "<Unsupported LValue>")
 

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

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.expressions.impl.FirExpressionStub
 import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
 import org.jetbrains.kotlin.fir.references.builder.*
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -49,9 +50,11 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class RawFirBuilder(
-    session: FirSession, val baseScopeProvider: FirScopeProvider, val stubMode: Boolean, val lazyBodiesMode: Boolean = false
+    session: FirSession, val baseScopeProvider: FirScopeProvider, val mode: RawFirBuilderMode = RawFirBuilderMode.NORMAL
 ) : BaseFirBuilder<PsiElement>(session) {
 
+    private val stubMode get() = mode == RawFirBuilderMode.STUBS
+
     fun buildFirFile(file: KtFile): FirFile {
         return file.accept(Visitor(), Unit) as FirFile
     }
@@ -74,7 +77,7 @@ class RawFirBuilder(
     }
 
     private fun buildDeclaration(declaration: KtDeclaration): FirDeclaration {
-        assert(!stubMode) { "Building FIR declarations isn't supported in stub mode" }
+        assert(mode ==  RawFirBuilderMode.NORMAL) { "Building FIR declarations isn't supported in stub or lazy mode mode" }
         setupContextForPosition(declaration)
         return declaration.accept(Visitor(), Unit) as FirDeclaration
     }
@@ -258,7 +261,7 @@ class RawFirBuilder(
             when {
                 !hasBody() ->
                     null to null
-                lazyBodiesMode -> {
+                mode == RawFirBuilderMode.LAZY_BODIES -> {
                     val block = buildLazyBlock {
                         source = bodyExpression?.toFirSourceElement()
                             ?: error("hasBody() == true but body is null")
@@ -1196,12 +1199,12 @@ class RawFirBuilder(
             val isVar = isVar
             val propertyInitializer = when {
                 !hasInitializer() -> null
-                lazyBodiesMode -> buildLazyExpression {
+                mode == RawFirBuilderMode.LAZY_BODIES -> buildLazyExpression {
                     source = initializer?.toFirSourceElement()
                 }
-                else -> {
+                mode == RawFirBuilderMode.STUBS -> buildExpressionStub()
-                    { initializer }.toFirExpression("Should have initializer")
+                else -> initializer.toFirExpression("Should have initializer")
-                }
+
             }
             val delegateExpression by lazy { delegate?.expression }
             val propertySource = toFirSourceElement()
@@ -2080,3 +2083,29 @@ class RawFirBuilder(
         }
     }
 }
+
+enum class RawFirBuilderMode {
+    /**
+     * Build every expression and every body
+     */
+    NORMAL,
+
+    /**
+     * Build [org.jetbrains.kotlin.fir.expressions.impl.FirExpressionStub] for expressions
+     */
+    STUBS,
+
+    /**
+     * Build [org.jetbrains.kotlin.fir.expressions.impl.FirLazyBlock] for function bodies, constructors & getters/setters
+     * Build [org.jetbrains.kotlin.fir.expressions.impl.FirLazyExpression] for property initializers
+     */
+    LAZY_BODIES;
+
+    companion object {
+        fun lazyBodies(lazyBodies: Boolean): RawFirBuilderMode =
+            if (lazyBodies) LAZY_BODIES else NORMAL
+
+        fun stubs(stubs: Boolean): RawFirBuilderMode =
+            if (stubs) STUBS else NORMAL
+    }
+}

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

@@ -59,7 +59,7 @@ abstract class AbstractPartialRawFirBuilderTestCase : AbstractRawFirBuilderTestC
 
         val session = FirSessionFactory.createEmptySession()
 
-        val firElement = buildFirElement(RawFirBuilder(session, StubFirScopeProvider, false), elementToBuild)
+        val firElement = buildFirElement(RawFirBuilder(session, StubFirScopeProvider), elementToBuild)
 
         val firDump = firElement.render(FirRenderer.RenderMode.WithFqNames)
         val expectedPath = filePath.replace(".kt", ".txt")

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(stubMode = false, bodyStubMode = true)
+        val firFile = file.toFirFile(RawFirBuilderMode.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/AbstractRawFirBuilderSourceElementMappingTestCase.kt

@@ -39,7 +39,7 @@ abstract class AbstractRawFirBuilderSourceElementMappingTestCase : AbstractRawFi
             }
             elements.single() as KtElement
         }
-        val firFile = ktFile.toFirFile(stubMode = false)
+        val firFile = ktFile.toFirFile()
         val foundElement = run {
             val foundElements = FindElementVisitor.find(firFile, selectedExpression)
             if (foundElements.size != 1) {

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

@@ -61,7 +61,7 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
 
     protected open fun doRawFirTest(filePath: String) {
         val file = createKtFile(filePath)
-        val firFile = file.toFirFile(stubMode = false)
+        val firFile = file.toFirFile(RawFirBuilderMode.NORMAL)
         val firFileDump = StringBuilder().also { FirRenderer(it).visitFile(firFile) }.toString()
         val expectedPath = filePath.replace(".kt", ".txt")
         KotlinTestUtils.assertEqualsToFile(File(expectedPath), firFileDump)
@@ -74,9 +74,9 @@ abstract class AbstractRawFirBuilderTestCase : KtParsingTestCase(
         }
     }
 
-    protected fun KtFile.toFirFile(stubMode: Boolean, bodyStubMode: Boolean = false): FirFile {
+    protected fun KtFile.toFirFile(mode: RawFirBuilderMode = RawFirBuilderMode.NORMAL): FirFile {
         val session = FirSessionFactory.createEmptySession()
-        return RawFirBuilder(session, StubFirScopeProvider, stubMode, bodyStubMode).buildFirFile(this)
+        return RawFirBuilder(session, StubFirScopeProvider, mode).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

@@ -59,7 +59,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
                 val ktFile = createKtFile(file.toRelativeString(root))
                 val firFile: FirFile
                 time += measureNanoTime {
-                    firFile = ktFile.toFirFile(stubMode)
+                    firFile = ktFile.toFirFile(RawFirBuilderMode.stubs(stubMode))
                 }
                 totalLength += StringBuilder().also { FirRenderer(it).visitFile(firFile) }.length
                 counter++
@@ -180,7 +180,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
             if ("testData" in path || "resources" in path || "api/js" in path.replace('\\', '/')) continue
             if (file.extension != "kt") continue
             val ktFile = createKtFile(file.toRelativeString(root))
-            val firFile = ktFile.toFirFile(stubMode = false)
+            val firFile = ktFile.toFirFile()
             try {
                 firFile.checkConsistency()
             } catch (e: Throwable) {
@@ -207,7 +207,7 @@ class RawFirBuilderTotalKotlinTestCase : AbstractRawFirBuilderTestCase() {
             if ("testData" in path || "resources" in path || "api/js" in path.replace('\\', '/')) continue
             if (file.extension != "kt") continue
             val ktFile = createKtFile(file.toRelativeString(root))
-            val firFile: FirFile = ktFile.toFirFile(stubMode = false)
+            val firFile: FirFile = ktFile.toFirFile()
             val psiSetViaFir = mutableSetOf<KtElement>()
             val psiSetDirect = mutableSetOf<KtElement>()
             firFile.accept(object : FirVisitorVoid() {

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

@@ -118,7 +118,7 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
                 firFile
             }
         } else {
-            val firBuilder = RawFirBuilder(session, firProvider.kotlinScopeProvider, false)
+            val firBuilder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
             ktFiles.mapTo(firFiles) {
                 val firFile = firBuilder.buildFirFile(it)
                 firProvider.recordFile(firFile)

compiler/visualizer/tests/org/jetbrains/kotlin/visualizer/fir/AbstractFirVisualizer.kt

@@ -33,7 +33,7 @@ abstract class AbstractFirVisualizer : AbstractVisualizer() {
         val session = createSession(environment, scope)
 
         val firProvider = (session.firProvider as FirProviderImpl)
-        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)
+        val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
 
         val transformer = FirTotalResolveProcessor(session)
         val firFiles = ktFiles.map {

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

@@ -6,6 +6,7 @@
 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.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -33,7 +34,7 @@ internal class FirFileBuilder(
         cache: ModuleFileCache,
         lazyBodiesMode: Boolean
     ): FirFile = cache.fileCached(ktFile) {
-        RawFirBuilder(cache.session, scopeProvider, stubMode = false, lazyBodiesMode = lazyBodiesMode).buildFirFile(ktFile)
+        RawFirBuilder(cache.session, scopeProvider, RawFirBuilderMode.lazyBodies(lazyBodiesMode)).buildFirFile(ktFile)
     }
 
     fun getBuiltFirFileOrNull(ktFile: KtFile, cache: ModuleFileCache): FirFile? =

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

@@ -78,7 +78,7 @@ internal object FirLazyBodiesCalculator {
 
     private fun createRawFirBuilder(firDeclaration: FirDeclaration): RawFirBuilder {
         val scopeProvider = firDeclaration.session.firIdeProvider.kotlinScopeProvider
-        return RawFirBuilder(firDeclaration.session, scopeProvider, stubMode = false)
+        return RawFirBuilder(firDeclaration.session, scopeProvider)
     }
 
     private val FIRST_PHASE_WHICH_NEEDS_BODIES = FirResolvePhase.CONTRACTS

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

@@ -97,11 +97,11 @@ internal class FirIdeProvider(
     // TODO move out of here
     // used only for completion
     fun buildFunctionWithBody(ktNamedFunction: KtNamedFunction): FirFunction<*> {
-        return RawFirBuilder(session, kotlinScopeProvider, stubMode = false).buildFunctionWithBody(ktNamedFunction)
+        return RawFirBuilder(session, kotlinScopeProvider).buildFunctionWithBody(ktNamedFunction)
     }
 
     fun buildPropertyWithBody(ktNamedFunction: KtProperty): FirProperty {
-        return RawFirBuilder(session, kotlinScopeProvider, stubMode = false).buildPropertyWithBody(ktNamedFunction)
+        return RawFirBuilder(session, kotlinScopeProvider).buildPropertyWithBody(ktNamedFunction)
     }
 
 

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

@@ -75,7 +75,7 @@ abstract class AbstractFirMultiModuleResolveTest : AbstractMultiModuleTest() {
             sessions += session
 
             val firProvider = (session.firProvider as FirProviderImpl)
-            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)
+            val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider)
             val psiManager = PsiManager.getInstance(project)
 
             val ideaModuleInfo = session.moduleInfo.cast<IdeaModuleInfo>()

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

@@ -9,6 +9,7 @@ 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.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -116,7 +117,7 @@ object KtDeclarationAndFirDeclarationEqualityChecker {
 
     private fun KtTypeReference.toKotlinTypReference(session: FirSession): FirTypeRef {
         // Maybe resolve all types here to not to work with FirTypeRef directly
-        return RawFirBuilder(session, DummyScopeProvider, stubMode = true).buildTypeReference(this)
+        return RawFirBuilder(session, DummyScopeProvider, RawFirBuilderMode.STUBS).buildTypeReference(this)
     }
 
     private fun ConeKotlinType.renderTypeAsKotlinType(): String {