[FIR] Refactor getFir to return only sane option or null

Merge-request: KT-MR-11138
Merged-by: Egor Kulikov <Egor.Kulikov@jetbrains.com>

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveSessionDepended.kt

@@ -61,7 +61,7 @@ internal class LLFirResolveSessionDepended(
 
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
         val psi = FirElementBuilder.getPsiAsFirElementSource(element) ?: return null
-        ktToFirMapping?.getFirOfClosestParent(psi)?.let { return it }
+        ktToFirMapping?.getFir(psi)?.let { return it }
         return originalFirResolveSession.getOrBuildFirFor(element = element)
     }
 

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

@@ -77,7 +77,7 @@ internal class FirElementBuilder(
         firResolveSession: LLFirResolveSession,
     ): FirElement? = when (element) {
         is KtFile -> getOrBuildFirForKtFile(element)
-        else -> getOrBuildFirForNonKtFileElement(element, firResolveSession)
+        else -> getFirForNonKtFileElement(element, firResolveSession)
     }
 
     private fun getOrBuildFirForKtFile(ktFile: KtFile): FirFile {
@@ -86,7 +86,7 @@ internal class FirElementBuilder(
         return firFile
     }
 
-    private fun getOrBuildFirForNonKtFileElement(
+    private fun getFirForNonKtFileElement(
         element: KtElement,
         firResolveSession: LLFirResolveSession,
     ): FirElement? {
@@ -96,8 +96,8 @@ internal class FirElementBuilder(
             return null
         }
 
-        getOrBuildFirForElementInsideAnnotations(element, firResolveSession)?.let { return it }
+        getFirForElementInsideAnnotations(element, firResolveSession)?.let { return it }
-        getOrBuildFirForElementInsideTypes(element, firResolveSession)?.let { return it }
+        getFirForElementInsideTypes(element, firResolveSession)?.let { return it }
 
         val psi = getPsiAsFirElementSource(element) ?: return null
         val firFile = element.containingKtFile
@@ -105,10 +105,10 @@ internal class FirElementBuilder(
 
         val structureElement = fileStructure.getStructureElementFor(element)
         val mappings = structureElement.mappings
-        return mappings.getFirOfClosestParent(psi) ?: firResolveSession.getOrBuildFirFile(firFile)
+        return mappings.getFir(psi)
     }
 
-    private inline fun <T : KtElement> getOrBuildFirForNonBodyElement(
+    private inline fun <T : KtElement> getFirForNonBodyElement(
         element: KtElement,
         firResolveSession: LLFirResolveSession,
         anchorElementProvider: (KtElement) -> T?,
@@ -138,10 +138,10 @@ internal class FirElementBuilder(
         return modifierList.owner as? KtDeclaration
     }
 
-    private fun getOrBuildFirForElementInsideAnnotations(
+    private fun getFirForElementInsideAnnotations(
         element: KtElement,
         firResolveSession: LLFirResolveSession,
-    ): FirElement? = getOrBuildFirForNonBodyElement(
+    ): FirElement? = getFirForNonBodyElement(
         element = element,
         firResolveSession = firResolveSession,
         anchorElementProvider = { it.parentOfType<KtAnnotationEntry>(withSelf = true) },
@@ -149,10 +149,10 @@ internal class FirElementBuilder(
         resolveAndFindFirForAnchor = { declaration, anchor -> declaration.resolveAndFindAnnotation(anchor, goDeep = true) },
     )
 
-    private fun getOrBuildFirForElementInsideTypes(
+    private fun getFirForElementInsideTypes(
         element: KtElement,
         firResolveSession: LLFirResolveSession,
-    ): FirElement? = getOrBuildFirForNonBodyElement(
+    ): FirElement? = getFirForNonBodyElement(
         element = element,
         firResolveSession = firResolveSession,
         anchorElementProvider = { it.parentsOfType<KtTypeReference>(withSelf = true).lastOrNull() },

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiErrorElement
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.ClassDiagnosticRetriever
@@ -25,6 +26,8 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 internal sealed class FileStructureElement(val firFile: FirFile, protected val moduleComponents: LLFirModuleResolveComponents) {
     abstract val psi: KtAnnotated
@@ -39,15 +42,67 @@ internal class KtToFirMapping(firElement: FirElement, recorder: FirElementsRecor
         return mapping[ktElement]
     }
 
-    fun getFirOfClosestParent(element: KtElement): FirElement? {
+    fun getFir(element: KtElement): FirElement? {
         var current: PsiElement? = element
-        while (current != null && current !is KtFile) {
+        while (
-            if (current is KtElement) {
+            current == element ||
-                getElement(current)?.let { return it }
+            current is KtUserType ||
-            }
+            current is KtTypeReference ||
+            current is KtDotQualifiedExpression ||
+            current is KtNullableType
+        ) {
+            // We are still referring to the same element with possible type parameter/name qualification/nullability, hence it is always
+            // sane to return corresponding element if present
+            if (current is KtElement) getElement(current)?.let { return it }
             current = current.parent
         }
-        return null
+
+        // Here current is the lowest ancestor that has different corresponding text
+        return when (current) {
+            // Constants with unary operation (i.e. +1 or -1) are saved as leaf element of FIR tree
+            is KtPrefixExpression,
+            // There is no separate element for annotation construction call
+            is KtAnnotationEntry,
+            // We replace source for selector for that of whole expression
+            is KtSafeQualifiedExpression,
+            // Top level destructuring declarations does not have FIR for r-value at the moment, would probably be changed later
+            is KtDestructuringDeclaration,
+            // There is no separate FIR node for this in this@foo expressions, same for super@Foo
+            is KtThisExpression,
+            is KtSuperExpression,
+            // Part of path in import/package directives has no FIR node
+            is KtImportDirective,
+            is KtPackageDirective,
+            // Super type refs are not recorded
+            is KtSuperTypeCallEntry,
+            // this/super in delegation calls are not part of FIR tree, this(args) is
+            is KtConstructorDelegationCall,
+            // In case of type projection we are not recording corresponding type reference
+            is KtTypeProjection,
+            // If we have, say, A(), reference A is not recorded, while call A() is recorded
+            is KtCallExpression ->
+                getElement(current as KtElement)
+            is KtBinaryExpression ->
+                // Here there is no separate FIR node for partial operator calls (like for a[i] = 1, there is no separate node for a[i])
+                if (element is KtArrayAccessExpression || element is KtOperationReferenceExpression) getElement(current) else null
+            is KtBlockExpression ->
+                // For script initializers we need to return FIR element for script itself
+                if (element is KtScriptInitializer) getElement(current.parent as KtScript) else null
+            is PsiErrorElement -> {
+                val parent = current.parent
+                if (parent is KtDestructuringDeclaration) getElement(parent) else null
+            }
+            // Value argument names and corresponding references are not part of FIR tree
+            is KtValueArgumentName -> getElement(current.parent as KtValueArgument)
+            is KtContainerNode -> {
+                val parent = current.parent
+                // Labels in labeled expression (i.e. return@foo) has no FIR node
+                if (parent is KtExpressionWithLabel) getElement(parent) else null
+            }
+            // Enum entries/annotation entries constructor calls
+            is KtConstructorCalleeExpression -> getElement(current.parent as KtCallElement)
+            else -> null
+        }
     }
 }