diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt index 1040dd90080..88349e176a8 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt @@ -8,11 +8,13 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder import org.jetbrains.kotlin.fir.FirElement import org.jetbrains.kotlin.fir.declarations.FirDeclaration import org.jetbrains.kotlin.fir.declarations.FirResolvePhase +import org.jetbrains.kotlin.fir.resolve.firProvider import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder +import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration import org.jetbrains.kotlin.idea.util.getElementTextInContext import org.jetbrains.kotlin.psi.* import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject @@ -36,10 +38,7 @@ internal class FirElementBuilder( val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache) val containerFir = when (val container = element.getNonLocalContainingDeclarationWithFqName()) { - is KtDeclaration -> { - FirElementFinder.findElementByPsiIn(firFile, container) - ?: error("Declaration was not found in FIR file which was build by declaration KtFile, declaration is\n${container.getElementTextInContext()}") - } + is KtDeclaration -> container.findNonLocalFirDeclaration(firFileBuilder, firFile.session.firProvider, moduleFileCache) null -> firFile else -> error("Unsupported: ${container.text}") } diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt index c16ea1232ec..f267493158b 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt @@ -19,11 +19,9 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE +import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile -import org.jetbrains.kotlin.idea.util.classIdIfNonLocal -import org.jetbrains.kotlin.idea.util.getElementTextInContext import org.jetbrains.kotlin.psi.* -import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject internal class FirLazyDeclarationResolver( private val firFileBuilder: FirFileBuilder @@ -139,31 +137,9 @@ internal class FirLazyDeclarationResolver( if (this is FirFile) return this val ktDeclaration = psi as? KtDeclaration ?: error("FirDeclaration should have a PSI of type KtDeclaration") if (!KtPsiUtil.isLocal(ktDeclaration)) return this - return when (val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName()) { - is KtClassOrObject -> provider.getFirClassifierByFqName( - nonLocalPsi.classIdIfNonLocal() - ?: error("Container classId should not be null for non-local declaration") - ) ?: error("Could not find class ${nonLocalPsi.classIdIfNonLocal()}") - is KtProperty, is KtNamedFunction -> { - val containerClass = nonLocalPsi.containingClassOrObject - if (containerClass != null) { - val containerClassId = containerClass.classIdIfNonLocal() - ?: error("Container classId should not be null for non-local declaration") - val containingFir = provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass - ?: error("Could not find class $containerClassId") - containingFir.declarations.first { it.psi === nonLocalPsi } - } else { - val packageFqName = ktDeclaration.containingKtFile.packageFqName - provider.symbolProvider.getTopLevelCallableSymbols(packageFqName, nonLocalPsi.nameAsSafeName) - .firstOrNull { it.fir.psi === nonLocalPsi } - val ktFile = ktDeclaration.containingKtFile - val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache) - firFile.declarations.firstOrNull { it.psi === nonLocalPsi } - ?: error("Cannot find FIR for\n${nonLocalPsi.getElementTextInContext()}") - } - } - else -> error("Invalid container ${nonLocalPsi?.let { it::class } ?: "null"}") - } + val nonLocalPsi = ktDeclaration.getNonLocalContainingDeclarationWithFqName() + ?: error("Container for local declaration cannot be null") + return nonLocalPsi.findNonLocalFirDeclaration(firFileBuilder, provider, moduleFileCache) } } diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt new file mode 100644 index 00000000000..49ebb6286ca --- /dev/null +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt @@ -0,0 +1,51 @@ +/* + * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.kotlin.idea.fir.low.level.api.util + +import org.jetbrains.kotlin.fir.declarations.FirDeclaration +import org.jetbrains.kotlin.fir.declarations.FirRegularClass +import org.jetbrains.kotlin.fir.psi +import org.jetbrains.kotlin.fir.resolve.providers.FirProvider +import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder +import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache +import org.jetbrains.kotlin.idea.util.classIdIfNonLocal +import org.jetbrains.kotlin.psi.* +import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject + +internal fun KtDeclaration.findNonLocalFirDeclaration( + firFileBuilder: FirFileBuilder, + provider: FirProvider, + moduleFileCache: ModuleFileCache +): FirDeclaration { + require(!KtPsiUtil.isLocal(this)) + return when { + this is KtClassOrObject -> findFir(provider) + this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> { + val containerClass = containingClassOrObject + if (containerClass != null) { + val containerClassFir = containerClass.findFir(provider) + containerClassFir.declarations.first { it.psi === this } + } else { + val ktFile = containingKtFile + val firFile = firFileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache) + firFile.declarations.first { it.psi === this } + } + } + this is KtConstructor<*> -> { + val containerClassFir = containingClassOrObject?.findFir(provider) + ?: error("Container class should be not null for KtConstructor") + containerClassFir.declarations.first { it.psi === this } + } + else -> error("Invalid container $this") + } +} + +private fun KtClassOrObject.findFir(provider: FirProvider): FirRegularClass { + val containerClassId = classIdIfNonLocal() + ?: error("Container classId should not be null for non-local declaration") + return provider.getFirClassifierByFqName(containerClassId) as? FirRegularClass + ?: error("Could not find class $containerClassId") +} \ No newline at end of file