From 673ac2227abb47dbb62d7b5ef2a180bfa09fa3a6 Mon Sep 17 00:00:00 2001 From: Egor Kulikov Date: Wed, 2 Aug 2023 18:24:57 +0000 Subject: [PATCH] [FIR] Refactor getFir to return only sane option or null Merge-request: KT-MR-11138 Merged-by: Egor Kulikov --- .../api/fir/LLFirResolveSessionDepended.kt | 2 +- .../fir/element/builder/FirElementBuilder.kt | 20 +++--- .../file/structure/FileStructureElement.kt | 67 +++++++++++++++++-- 3 files changed, 72 insertions(+), 17 deletions(-) diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveSessionDepended.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveSessionDepended.kt index c8a0ccde08a..2500809450b 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveSessionDepended.kt +++ b/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) } diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt index 335d3c824d8..03f6af76514 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt +++ b/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 } - 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 getOrBuildFirForNonBodyElement( + private inline fun 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(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(withSelf = true).lastOrNull() }, diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt index 2d2b3d269f3..1a6bd018dd1 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt +++ b/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) { - 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 + } } }