diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt index 868e12de9f0..58c610c2b4f 100644 --- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt +++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt @@ -9,6 +9,7 @@ import com.intellij.codeInsight.completion.* import com.intellij.patterns.PlatformPatterns import com.intellij.patterns.PsiJavaPatterns import com.intellij.util.ProcessingContext +import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession import org.jetbrains.kotlin.idea.frontend.api.getAnalysisSessionFor @@ -71,9 +72,16 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch lookupElementFactory.createLookupElement(symbol)?.let(::addElement) } + private fun recordOriginalFile(completionParameters: CompletionParameters) { + val originalFile = completionParameters.originalFile as? KtFile ?: return + val fakeFile = completionParameters.position.containingFile as? KtFile ?: return + fakeFile.originalKtFile = originalFile + } + @OptIn(InvalidWayOfUsingAnalysisSession::class) fun addCompletions(parameters: CompletionParameters, result: CompletionResultSet) { val originalFile = parameters.originalFile as? KtFile ?: return + recordOriginalFile(parameters) val reference = (parameters.position.parent as? KtSimpleNameExpression)?.mainReference ?: return val nameExpression = reference.expression.takeIf { it !is KtLabelReferenceExpression } ?: return diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt index ccf109be9b7..5fe56061d2e 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt @@ -12,12 +12,15 @@ import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.declarations.FirDeclaration import org.jetbrains.kotlin.fir.declarations.FirFile import org.jetbrains.kotlin.fir.declarations.FirResolvePhase +import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.resolve.providers.FirProvider import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder +import org.jetbrains.kotlin.idea.fir.low.level.api.util.containingKtFileIfAny +import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile import org.jetbrains.kotlin.psi.KtDeclaration import org.jetbrains.kotlin.psi.KtElement import org.jetbrains.kotlin.psi.KtFile @@ -74,9 +77,13 @@ internal class FirModuleResolveStateForCompletion( } override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? { - return cache.getContainerFirFile(declaration) + val ktFile = declaration.containingKtFileIfAny ?: return null + cache.getCachedFirFile(ktFile)?.let { return it } + ktFile.originalKtFile?.let(cache::getCachedFirFile)?.let { return it } + return null } + override fun getDiagnostics(element: KtElement): List { error("Diagnostics should not be retrieved in completion") } 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 index aa678a9e372..ca1b4284318 100644 --- 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 @@ -42,8 +42,9 @@ internal fun KtDeclaration.findFirDeclarationForAnyFirSourceDeclaration( val nonLocalDeclaration = getNonLocalContainingOrThisDeclaration() ?.findSourceNonLocalFirDeclaration(firFileBuilder, firSymbolProvider, moduleFileCache) ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile, moduleFileCache, lazyBodiesMode = true) + val originalDeclaration = originalDeclaration val fir = FirElementFinder.findElementIn(nonLocalDeclaration) { firDeclaration -> - firDeclaration.psi == this + firDeclaration.psi == this || firDeclaration.psi == originalDeclaration } return fir ?: error("FirDeclaration was not found for\n${getElementTextInContext()}") @@ -105,9 +106,11 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider( } val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key("ORIGINAL_DECLARATION_KEY") - var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY) +private val ORIGINAL_KT_FILE_KEY = com.intellij.openapi.util.Key("ORIGINAL_KT_FILE_KEY") +var KtFile.originalKtFile by UserDataProperty(ORIGINAL_KT_FILE_KEY) + private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass? { val classId = classIdIfNonLocal() ?: return null diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/utils.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/utils.kt index 885aedd2c31..066e0e23cda 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/utils.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/utils.kt @@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.render import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo import org.jetbrains.kotlin.psi.KtDeclaration +import org.jetbrains.kotlin.psi.KtFile import java.util.concurrent.TimeUnit import java.util.concurrent.locks.Lock @@ -63,6 +64,9 @@ internal val FirDeclaration.ktDeclaration: KtDeclaration return psi as KtDeclaration } +internal val FirDeclaration.containingKtFileIfAny: KtFile? + get() = psi?.containingFile as? KtFile + internal fun IdeaModuleInfo.collectTransitiveDependenciesWithSelf(): List { val result = mutableSetOf() diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/EnclosingDeclarationContext.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/EnclosingDeclarationContext.kt index 2db0bab2634..4d190bfe2ea 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/EnclosingDeclarationContext.kt +++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/EnclosingDeclarationContext.kt @@ -8,7 +8,6 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.utils import com.intellij.psi.util.parentsOfType import org.jetbrains.kotlin.fir.declarations.FirFile import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState -import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacade import org.jetbrains.kotlin.idea.fir.low.level.api.api.LowLevelFirApiFacadeForCompletion import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalDeclaration import org.jetbrains.kotlin.idea.util.getElementTextInContext @@ -22,7 +21,7 @@ internal sealed class EnclosingDeclarationContext { if (fakeFunction != null) { val originalFunction = originalFile.findDeclarationOfTypeAt(fakeFunction.textOffset) ?: error("Cannot find original function matching to ${fakeFunction.getElementTextInContext()} in $originalFile") - fakeFunction.originalDeclaration = originalFunction + recordOriginalDeclaration(originalFunction, fakeFunction) return FunctionContext(fakeFunction, originalFunction) } @@ -30,13 +29,23 @@ internal sealed class EnclosingDeclarationContext { if (fakeProperty != null) { val originalProperty = originalFile.findDeclarationOfTypeAt(fakeProperty.textOffset) ?: error("Cannot find original property matching to ${fakeProperty.getElementTextInContext()} in $originalFile") - fakeProperty.originalDeclaration = originalProperty - + recordOriginalDeclaration(originalProperty, fakeProperty) return PropertyContext(fakeProperty, originalProperty) } error("Cannot find enclosing declaration for ${positionInFakeFile.getElementTextInContext()}") } + + private fun recordOriginalDeclaration(originalDeclaration: KtNamedDeclaration, fakeDeclaration: KtNamedDeclaration) { + require(!fakeDeclaration.isPhysical) + require(originalDeclaration.containingKtFile !== fakeDeclaration.containingKtFile) + val originalDeclrationParents = originalDeclaration.parentsOfType().toList() + val fakeDeclarationParents = fakeDeclaration.parentsOfType().toList() + + originalDeclrationParents.zip(fakeDeclarationParents) { original, fake -> + fake.originalDeclaration = original + } + } } }