LL API: allow getOrBuildFir to return null

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

@@ -41,7 +41,7 @@ internal class FirModuleResolveStateDepended(
     override fun getSessionFor(moduleInfo: ModuleInfo): FirSession =
         originalState.getSessionFor(moduleInfo)
 
-    override fun getOrBuildFirFor(element: KtElement): FirElement {
+    override fun getOrBuildFirFor(element: KtElement): FirElement? {
         val psi = originalState.elementBuilder.getPsiAsFirElementSource(element)
 
         ktToFirMapping?.getFirOfClosestParent(psi, this)?.let { return it }

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

@@ -61,7 +61,7 @@ internal class FirModuleResolveStateImpl(
     override fun getSessionFor(moduleInfo: ModuleInfo): FirSession =
         sessionProvider.getSession(moduleInfo)!!
 
-    override fun getOrBuildFirFor(element: KtElement): FirElement =
+    override fun getOrBuildFirFor(element: KtElement): FirElement? =
         elementBuilder.getOrBuildFirFor(
             element = element,
             firFileBuilder = firFileBuilder,

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

@@ -34,7 +34,7 @@ abstract class FirModuleResolveState {
      * This operation could be performance affective because it create FIleStructureElement and resolve non-local declaration into BODY phase, use
      * @see tryGetCachedFirFile to get [FirFile] in undefined phase
      */
-    internal abstract fun getOrBuildFirFor(element: KtElement): FirElement
+    internal abstract fun getOrBuildFirFor(element: KtElement): FirElement?
 
     /**
      * Get or build or get cached [FirFile] for requested file in undefined phase

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

@@ -208,11 +208,13 @@ fun <D : FirDeclaration> D.resolvedFirToType(
 /**
  * Get a [FirElement] which was created by [KtElement]
  * Returned [FirElement] is guaranteed to be resolved to [FirResolvePhase.BODY_RESOLVE] phase
- * This operation could be performance affective because it create FIleStructureElement and resolve non-local declaration into BODY phase
+ * This operation could be performance affective because it create FIleStructureElement and resolve non-local declaration into BODY phase.
+ *
+ * The `null` value is returned iff FIR tree does not have corresponding element
  */
 fun KtElement.getOrBuildFir(
     resolveState: FirModuleResolveState,
-): FirElement = resolveState.getOrBuildFirFor(this)
+): FirElement? = resolveState.getOrBuildFirFor(this)
 
 /**
  * Get a [FirElement] which was created by [KtElement], but only if it is subtype of [E], `null` otherwise

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

@@ -233,8 +233,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
             }
         }
 
-        val originalDeclaration = nonLocalDeclaration.getOrBuildFir(state)
+        val originalDeclaration = nonLocalDeclaration.getOrBuildFirOfType<FirDeclaration>(state)
-        check(originalDeclaration is FirDeclaration) { "Invalid original declaration type ${originalDeclaration::class.simpleName}" }
 
         val originalDesignation = originalDeclaration.collectDesignation()
 

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

@@ -55,7 +55,7 @@ internal class FirElementBuilder {
         fileStructureCache: FileStructureCache,
         firLazyDeclarationResolver: FirLazyDeclarationResolver,
         state: FirModuleResolveState,
-    ): FirElement = when (element) {
+    ): FirElement? = when (element) {
         is KtFile -> getOrBuildFirForKtFile(element, firFileBuilder, moduleFileCache, firLazyDeclarationResolver)
         else -> getOrBuildFirForNonKtFileElement(element, fileStructureCache, moduleFileCache, state)
     }
@@ -82,7 +82,7 @@ internal class FirElementBuilder {
         fileStructureCache: FileStructureCache,
         moduleFileCache: ModuleFileCache,
         state: FirModuleResolveState,
-    ): FirElement {
+    ): FirElement? {
         require(element !is KtFile)
         val firFile = element.containingKtFile
         val fileStructure = fileStructureCache.getFileStructure(firFile, moduleFileCache)

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

@@ -36,7 +36,7 @@ abstract class AbstractGetOrBuildFirTest : AbstractLowLevelApiSingleFileTest() {
         val actual = resolveWithClearCaches(ktFile) { state ->
             val fir = selectedElement.getOrBuildFir(state)
             """|KT element: ${selectedElement::class.simpleName}
-               |FIR element: ${fir::class.simpleName}
+               |FIR element: ${fir?.let { it::class.simpleName }}
                |
                |FIR element rendered:
                |${render(fir)}""".trimMargin()
@@ -55,7 +55,8 @@ abstract class AbstractGetOrBuildFirTest : AbstractLowLevelApiSingleFileTest() {
         ).single { it.textRange == selectedElement.textRange }
     }
 
-    private fun render(firElement: FirElement): String = when (firElement) {
+    private fun render(firElement: FirElement?): String = when (firElement) {
+        null -> "null"
         is FirImport -> "import ${firElement.importedFqName}"
         else -> firElement.render(renderingMode)
     }

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.fir.FirRealSourceElementKind
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorVisitor
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.fir.types.FirTypeRef
@@ -17,6 +18,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.collectDiagnosticsForFile
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.BeforeElementDiagnosticCollectionHandler
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.fir.PersistentCheckerContextFactory
@@ -44,7 +46,7 @@ abstract class AbstractDiagnosticTraversalCounterTest  : AbstractLowLevelApiSing
             // we should get diagnostics before we resolve the whole file by  ktFile.getOrBuildFir
             ktFile.collectDiagnosticsForFile(resolveState, DiagnosticCheckerFilter.ONLY_COMMON_CHECKERS)
 
-            val firFile = ktFile.getOrBuildFir(resolveState)
+            val firFile = ktFile.getOrBuildFirOfType<FirFile>(resolveState)
 
             val errorElements = collectErrorElements(firFile, handler)
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt

@@ -11,6 +11,7 @@ import com.intellij.psi.SmartPsiElementPointer
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClassSymbol
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
@@ -41,6 +42,7 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForResolveOnAir
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFir
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.parentsOfType
 import org.jetbrains.kotlin.idea.frontend.api.components.KtReferenceShortener
@@ -74,7 +76,7 @@ internal class KtFirReferenceShortener(
         val declarationToVisit = file.findSmallestDeclarationContainingSelection(selection)
             ?: file.withDeclarationsResolvedToBodyResolve()
 
-        val firDeclaration = declarationToVisit.getOrBuildFir(firResolveState)
+        val firDeclaration = declarationToVisit.getOrBuildFirOfType<FirDeclaration>(firResolveState)
 
         val towerContext =
             LowLevelFirApiFacadeForResolveOnAir.onAirGetTowerContextProvider(firResolveState, declarationToVisit)