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

Merge-request: KT-MR-11138
Merged-by: Egor Kulikov <Egor.Kulikov@jetbrains.com>
This commit is contained in:
Egor Kulikov
2023-08-02 18:24:57 +00:00
committed by Space Team
parent 64c8ce18a0
commit 673ac2227a
3 changed files with 72 additions and 17 deletions
@@ -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)
}
@@ -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 }
getOrBuildFirForElementInsideTypes(element, firResolveSession)?.let { return it }
getFirForElementInsideAnnotations(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() },
@@ -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) {
if (current is KtElement) {
getElement(current)?.let { return it }
}
while (
current == element ||
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
}
}
}