diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/PersistentMultimap.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/PersistentMultimap.kt index c33ff9ebcae..2d2961c3314 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/PersistentMultimap.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/PersistentMultimap.kt @@ -34,6 +34,8 @@ class PersistentMultimap private constructor(private val map: PersistentMa operator fun get(key: K): List { return map[key] ?: emptyList() } + + val keys: ImmutableSet get() = map.keys } class PersistentSetMultimap private constructor(private val map: PersistentMap>) { @@ -63,4 +65,4 @@ class PersistentSetMultimap private constructor(private val map: Persisten operator fun get(key: K): Set { return map[key] ?: emptySet() } -} \ No newline at end of file +} diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt index 027608eba7d..6b1c2895437 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt @@ -83,6 +83,8 @@ open class FirBodyResolveTransformer( return data.expectedTypeRef.compose() } + open fun onBeforeStatementResolution(statement: FirStatement) {} + // ------------------------------------- Expressions ------------------------------------- override fun transformExpression(expression: FirExpression, data: ResolutionMode): CompositeTransformResult { diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt index 3d4f780eab3..9960d35d77f 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt @@ -291,6 +291,9 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform block.transformStatementsIndexed(transformer) { index -> val value = if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent + + transformer.onBeforeStatementResolution(block.statements[index]) + TransformData.Data(value) } if (data == ResolutionMode.ContextIndependent) { diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirLocalScope.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirLocalScope.kt index 9e018f6263e..6a10fa07d08 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirLocalScope.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirLocalScope.kt @@ -70,4 +70,6 @@ class FirLocalScope private constructor( } override fun mayContainName(name: Name) = properties.containsKey(name) || functions[name].isNotEmpty() || classes.containsKey(name) + + override fun getCallableNames(): Set = properties.keys + functions.keys + classes.keys } diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirResolutionApi.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirResolutionApi.kt index 8613f5a15b2..9efae9301b2 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirResolutionApi.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirResolutionApi.kt @@ -9,8 +9,12 @@ import com.intellij.openapi.progress.ProgressIndicatorProvider import com.intellij.psi.PsiElement import org.jetbrains.kotlin.fir.* import org.jetbrains.kotlin.fir.declarations.* +import org.jetbrains.kotlin.fir.expressions.FirStatement +import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.references.* -import org.jetbrains.kotlin.fir.resolve.* +import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext +import org.jetbrains.kotlin.fir.resolve.ResolutionMode +import org.jetbrains.kotlin.fir.resolve.ScopeSession import org.jetbrains.kotlin.fir.resolve.providers.FirProvider import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredCallableSymbols import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer @@ -162,7 +166,8 @@ internal fun FirDeclaration.runResolve( file: FirFile, firProvider: FirIdeProvider, toPhase: FirResolvePhase, - state: FirModuleResolveState + state: FirModuleResolveState, + towerDataContextForStatement: MutableMap? = null, ) { val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS) file.runResolve(toPhase = nonLazyPhase, fromPhase = this.resolvePhase) @@ -193,7 +198,8 @@ internal fun FirDeclaration.runResolve( val transformer = FirDesignatedBodyResolveTransformerForIDE( designation.iterator(), state.getSession(psi as KtElement), scopeSession, - implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE + implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE, + towerDataContextForStatement ) file.transform(transformer, ResolutionMode.ContextDependent) } @@ -202,7 +208,8 @@ private class FirDesignatedBodyResolveTransformerForIDE( private val designation: Iterator, session: FirSession, scopeSession: ScopeSession, - implicitTypeOnly: Boolean + implicitTypeOnly: Boolean, + private val towerDataContextForStatement: MutableMap? = null ) : FirBodyResolveTransformer( session, phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE, @@ -210,6 +217,7 @@ private class FirDesignatedBodyResolveTransformerForIDE( scopeSession = scopeSession, returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession) ) { + override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult { if (designation.hasNext()) { designation.next().visitNoTransform(this, data) @@ -218,6 +226,11 @@ private class FirDesignatedBodyResolveTransformerForIDE( return super.transformDeclarationContent(declaration, data) } + + override fun onBeforeStatementResolution(statement: FirStatement) { + if (towerDataContextForStatement == null) return + towerDataContextForStatement[statement] = context.towerDataContext + } } diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt index 5fcdf5e96ae..bc32a7dc54a 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/LowLevelFirApiFacade.kt @@ -5,12 +5,15 @@ package org.jetbrains.kotlin.idea.fir.low.level.api +import com.intellij.psi.PsiElement import org.jetbrains.kotlin.diagnostics.Diagnostic import org.jetbrains.kotlin.fir.FirElement import org.jetbrains.kotlin.fir.declarations.FirFile import org.jetbrains.kotlin.fir.declarations.FirFunction import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.declarations.FirResolvePhase +import org.jetbrains.kotlin.fir.expressions.FirStatement +import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext import org.jetbrains.kotlin.psi.KtElement import org.jetbrains.kotlin.psi.KtNamedFunction @@ -31,21 +34,48 @@ object LowLevelFirApiFacade { fun getFirOfClosestParent(element: KtElement): FirElement? = element.getFirOfClosestParent(element.firResolveState())?.second - fun buildFunctionWithResolvedBody( + class FirCompletionContext internal constructor( + val session: FirSession, + private val towerDataContextForStatement: MutableMap, + private val state: FirModuleResolveState, + ) { + fun getTowerDataContext(element: KtElement): FirTowerDataContext { + var current: PsiElement? = element + while (current is KtElement) { + val mappedFir = state.getCachedMapping(current) + + if (mappedFir is FirStatement) { + towerDataContextForStatement[mappedFir]?.let { return it } + } + current = current.parent + } + + error("No context for $element") + } + } + + fun buildCompletionContextForFunction( firFile: FirFile, element: KtNamedFunction, phase: FirResolvePhase = FirResolvePhase.BODY_RESOLVE - ): FirFunction<*> { + ): FirCompletionContext { val state = element.firResolveState() val firIdeProvider = firFile.session.firIdeProvider val builtFunction = firIdeProvider.buildFunctionWithBody(element) + val towerDataContextForStatement = mutableMapOf() - return builtFunction.apply { - runResolve(firFile, firIdeProvider, phase, state) + val function = builtFunction.apply { + runResolve(firFile, firIdeProvider, phase, state, towerDataContextForStatement) // TODO this PSI caching should be somewhere else state.recordElementsFrom(this) } + + return FirCompletionContext( + function.session, + towerDataContextForStatement, + state + ) } fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection { @@ -53,4 +83,4 @@ object LowLevelFirApiFacade { file.getOrBuildFirWithDiagnostics(resolveState) return resolveState.getDiagnostics(element) } -} \ No newline at end of file +} diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt index 4726959e486..412105a59a3 100644 --- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt +++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt @@ -12,18 +12,19 @@ import com.intellij.util.ProcessingContext import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.declarations.FirFile import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression -import org.jetbrains.kotlin.fir.resolve.* -import org.jetbrains.kotlin.fir.resolve.providers.getSymbolByTypeRef -import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol +import org.jetbrains.kotlin.fir.resolve.ScopeSession +import org.jetbrains.kotlin.fir.resolve.scope import org.jetbrains.kotlin.fir.types.ConeKotlinType import org.jetbrains.kotlin.fir.types.coneTypeUnsafe import org.jetbrains.kotlin.idea.fir.getFirOfClosestParent -import org.jetbrains.kotlin.idea.fir.getOrBuildFir import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe import org.jetbrains.kotlin.idea.fir.low.level.api.LowLevelFirApiFacade import org.jetbrains.kotlin.idea.references.mainReference import org.jetbrains.kotlin.name.Name -import org.jetbrains.kotlin.psi.* +import org.jetbrains.kotlin.psi.KtFile +import org.jetbrains.kotlin.psi.KtLabelReferenceExpression +import org.jetbrains.kotlin.psi.KtNamedFunction +import org.jetbrains.kotlin.psi.KtSimpleNameExpression import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType class KotlinFirCompletionContributor : CompletionContributor() { @@ -46,16 +47,25 @@ private object KotlinFirCompletionProvider : CompletionProvider() ?: return - val firFunction = LowLevelFirApiFacade.buildFunctionWithResolvedBody(originalFileFir, parentFunction) + val completionContext = LowLevelFirApiFacade.buildCompletionContextForFunction(originalFileFir, parentFunction) val element = nameExpression.getFirOfClosestParent() as? FirQualifiedAccessExpression ?: return - for (scope in getScopes(element, firFunction.session)) { + for (scope in getScopes(element, completionContext.session)) { for (name in scope.availableNames()) { result.addElement(LookupElementBuilder.create(name.asString())) } } + val towerDataContext = completionContext.getTowerDataContext(nameExpression) + + for (localScope in towerDataContext.localScopes) { + for (name in localScope.getCallableNames()) { + result.addElement(LookupElementBuilder.create(name.asString())) + } + } + + // val receiver = element.explicitReceiver // if (receiver != null) { //