From fcefdcc3e9c49b7c4c34534f937046383d3c0320 Mon Sep 17 00:00:00 2001 From: Yan Zhulanow Date: Tue, 1 Aug 2023 14:57:50 +0900 Subject: [PATCH] [LL API] Move PSI element context collection to 'ContextCollector' The new API is going to be reused in 'KtFirScopeProvider'. --- .../fir/transformers/LLFirBodyLazyResolver.kt | 34 ++----------- .../level/api/fir/util/ContextCollector.kt | 50 +++++++++++++++++-- 2 files changed, 51 insertions(+), 33 deletions(-) diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt index e1130ce8f09..03baed7e271 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt +++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt @@ -5,15 +5,11 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers -import com.intellij.psi.PsiElement import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget -import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget -import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.asResolveTarget import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError import org.jetbrains.kotlin.analysis.low.level.api.fir.compile.codeFragmentScopeProvider -import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData @@ -53,9 +49,8 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol import org.jetbrains.kotlin.fir.types.ConeKotlinType import org.jetbrains.kotlin.psi.KtCodeFragment import org.jetbrains.kotlin.psi.KtFile -import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf -import org.jetbrains.kotlin.utils.exceptions.buildErrorWithAttachment import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment +import org.jetbrains.kotlin.utils.exceptions.withPsiEntry import org.jetbrains.kotlin.utils.findIsInstanceAnd internal object LLFirBodyLazyResolver : LLFirLazyResolver(FirResolvePhase.BODY_RESOLVE) { @@ -180,20 +175,11 @@ private class LLFirBodyTargetResolver( firCodeFragment.codeFragmentContext = if (contextKtFile != null) { val contextFirFile = resolveSession.getOrBuildFirFile(contextKtFile) - val target = computeCodeFragmentContextResolveTarget(contextFirFile, contextPsiElement) - val contextParentKtElements = contextPsiElement.parentsWithSelf.toSet() - - val contextProvider = ContextCollector.process(target, transformer.components) { candidate -> - when (candidate) { - contextPsiElement -> ContextCollector.FilterResponse.STOP - in contextParentKtElements -> ContextCollector.FilterResponse.CONTINUE - else -> ContextCollector.FilterResponse.SKIP + val elementContext = ContextCollector.process(contextFirFile, transformer.components, contextPsiElement, preferBody = true) + ?: errorWithAttachment("Cannot find enclosing context for ${contextPsiElement::class}") { + withPsiEntry("contextPsiElement", contextPsiElement) } - } - - val elementContext = contextProvider[contextPsiElement, ContextCollector.ContextKind.BODY] - ?: contextParentKtElements.firstNotNullOf { contextProvider[it, ContextCollector.ContextKind.SELF] } LLFirCodeFragmentContext(elementContext.towerDataContext.withExtraScopes(), elementContext.smartCasts) } else { @@ -202,18 +188,6 @@ private class LLFirBodyTargetResolver( } } - private fun computeCodeFragmentContextResolveTarget(contextFirFile: FirFile, contextPsiElement: PsiElement): LLFirResolveTarget { - val contextKtDeclaration = contextPsiElement.getNonLocalContainingOrThisDeclaration() - if (contextKtDeclaration != null) { - val designationPath = FirElementFinder.collectDesignationPath(contextFirFile, contextKtDeclaration) - if (designationPath != null) { - return FirDesignationWithFile(designationPath.path, designationPath.target, contextFirFile).asResolveTarget() - } - } - - return LLFirWholeFileResolveTarget(contextFirFile) - } - override fun doLazyResolveUnderLock(target: FirElementWithResolveState) { when (target) { is FirRegularClass, is FirCodeFragment -> error("Should have been resolved in ${::doResolveWithoutLock.name}") diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt index 2aa336d4dc4..312a9e1b4c8 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt +++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/ContextCollector.kt @@ -9,7 +9,11 @@ import com.intellij.openapi.progress.ProgressManager import com.intellij.psi.PsiElement import kotlinx.collections.immutable.PersistentMap import kotlinx.collections.immutable.persistentMapOf +import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget +import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget +import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.asResolveTarget +import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.ContextKind import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.Context import org.jetbrains.kotlin.analysis.low.level.api.fir.util.ContextCollector.FilterResponse @@ -29,6 +33,7 @@ import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase import org.jetbrains.kotlin.fir.types.ConeKotlinType import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid +import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf internal object ContextCollector { enum class ContextKind { @@ -55,6 +60,45 @@ internal object ContextCollector { SKIP } + /** + * Process the [file], collecting contexts for the [targetElement] and all its PSI tree parents. + */ + fun process(file: FirFile, holder: SessionHolder, targetElement: PsiElement, preferBody: Boolean): Context? { + val acceptedElements = targetElement.parentsWithSelf.toSet() + + val contextProvider = process(computeResolveTarget(file, targetElement), holder) { candidate -> + when (candidate) { + targetElement -> FilterResponse.STOP + in acceptedElements -> FilterResponse.CONTINUE + else -> FilterResponse.SKIP + } + } + + if (preferBody) { + val bodyContext = contextProvider[targetElement, ContextKind.BODY] + if (bodyContext != null) { + return bodyContext + } + } + + return acceptedElements.firstNotNullOfOrNull { contextProvider[it, ContextKind.SELF] } + } + + private fun computeResolveTarget(file: FirFile, targetElement: PsiElement): LLFirResolveTarget { + val contextKtDeclaration = targetElement.getNonLocalContainingOrThisDeclaration() + if (contextKtDeclaration != null) { + val designationPath = FirElementFinder.collectDesignationPath(file, contextKtDeclaration) + if (designationPath != null) { + return FirDesignationWithFile(designationPath.path, designationPath.target, file).asResolveTarget() + } + } + + return LLFirWholeFileResolveTarget(file) + } + + /** + * Processes the [FirFile] that owns the [target], collecting contexts for elements matching the [filter]. + */ fun process(target: LLFirResolveTarget, holder: SessionHolder, filter: (PsiElement) -> FilterResponse): ContextProvider { val pathIterator = target.path.iterator() @@ -67,10 +111,10 @@ internal object ContextCollector { return ContextProvider { element, kind -> visitor[element, kind] } } -} -internal fun interface ContextProvider { - operator fun get(element: PsiElement, kind: ContextKind): Context? + fun interface ContextProvider { + operator fun get(element: PsiElement, kind: ContextKind): Context? + } } private class ContextCollectorVisitor(