diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveState.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveState.kt index 1c0b89c4cad..dec3de26d0e 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveState.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveState.kt @@ -7,17 +7,16 @@ package org.jetbrains.kotlin.idea.fir.low.level.api import org.jetbrains.kotlin.diagnostics.Diagnostic import org.jetbrains.kotlin.fir.FirElement -import org.jetbrains.kotlin.fir.declarations.FirDeclaration -import org.jetbrains.kotlin.fir.declarations.FirFunction -import org.jetbrains.kotlin.fir.declarations.FirFile 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.FirFunction 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.idea.caches.project.IdeaModuleInfo import org.jetbrains.kotlin.idea.caches.project.getModuleInfo import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.DiagnosticsCollector import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder +import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache @@ -52,7 +51,7 @@ abstract class FirModuleResolveState { containerFirFile: FirFile, firIdeProvider: FirIdeProvider, toPhase: FirResolvePhase, - towerDataContextForStatement: MutableMap + towerDataContextCollector: FirTowerDataContextCollector ) } @@ -97,14 +96,14 @@ internal open class FirModuleResolveStateImpl( containerFirFile: FirFile, firIdeProvider: FirIdeProvider, toPhase: FirResolvePhase, - towerDataContextForStatement: MutableMap + towerDataContextCollector: FirTowerDataContextCollector ) { firLazyDeclarationResolver.runLazyResolveWithoutLock( firFunction, containerFirFile, firIdeProvider, toPhase, - towerDataContextForStatement, + towerDataContextCollector, checkPCE = false ) } 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 f98db435234..44a9ad50805 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,9 +12,8 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration import org.jetbrains.kotlin.fir.declarations.FirFile import org.jetbrains.kotlin.fir.declarations.FirFunction 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.idea.caches.project.IdeaModuleInfo +import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.PsiToFirCache import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider import org.jetbrains.kotlin.psi.KtElement @@ -61,9 +60,9 @@ internal class FirModuleResolveStateForCompletion( containerFirFile: FirFile, firIdeProvider: FirIdeProvider, toPhase: FirResolvePhase, - towerDataContextForStatement: MutableMap + towerDataContextCollector: FirTowerDataContextCollector ) { - originalState.lazyResolveFunctionForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextForStatement) + originalState.lazyResolveFunctionForCompletion(firFunction, containerFirFile, firIdeProvider, toPhase, towerDataContextCollector) } override fun getDiagnostics(element: KtElement): List { 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 2180b548999..4529d6aa092 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 @@ -19,8 +19,8 @@ import org.jetbrains.kotlin.idea.caches.project.getModuleInfo import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider import org.jetbrains.kotlin.psi.KtElement import org.jetbrains.kotlin.psi.KtNamedFunction -import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol +import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector import org.jetbrains.kotlin.idea.util.getElementTextInContext object LowLevelFirApiFacade { @@ -40,7 +40,7 @@ object LowLevelFirApiFacade { class FirCompletionContext internal constructor( val session: FirSession, - private val towerDataContextForStatement: Map, + private val towerDataContextCollector: FirTowerDataContextCollector, private val state: FirModuleResolveState, ) { fun getTowerDataContext(element: KtElement): FirTowerDataContext { @@ -49,13 +49,7 @@ object LowLevelFirApiFacade { val mappedFir = state.getCachedMappingForCompletion(current) if (mappedFir is FirStatement) { - towerDataContextForStatement[mappedFir]?.let { return it } - for ((key, value) in towerDataContextForStatement) { - // TODO Rework that - if (key.psi == mappedFir.psi) { - return value - } - } + towerDataContextCollector.getContext(mappedFir)?.let { return it } } current = current.parent } @@ -74,7 +68,7 @@ object LowLevelFirApiFacade { val firIdeProvider = firFile.session.firIdeProvider val originalFunction = state.getOrBuildFirFor(originalElement, phase) as FirSimpleFunction val builtFunction = firIdeProvider.buildFunctionWithBody(element) - val towerDataContextForStatement = mutableMapOf() + val contextCollector = FirTowerDataContextCollector() // right now we can't resolve builtFunction header properly, as it built right in air, // without file, which is now required for running stages other then body resolve, so we @@ -87,13 +81,13 @@ object LowLevelFirApiFacade { } val function = frankensteinFunction.apply { - state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, towerDataContextForStatement) + state.lazyResolveFunctionForCompletion(this, firFile, firIdeProvider, phase, contextCollector) state.recordPsiToFirMappingsForCompletionFrom(this, firFile, element.containingKtFile) } return FirCompletionContext( function.session, - towerDataContextForStatement, + contextCollector, state ) } diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirDesignatedBodyResolveTransformerForIDE.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirDesignatedBodyResolveTransformerForIDE.kt index 4a07f08380d..749714dd152 100644 --- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirDesignatedBodyResolveTransformerForIDE.kt +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirDesignatedBodyResolveTransformerForIDE.kt @@ -10,7 +10,6 @@ import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.declarations.FirDeclaration 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.fir.resolve.ResolutionMode import org.jetbrains.kotlin.fir.resolve.ScopeSession import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer @@ -23,7 +22,7 @@ internal class FirDesignatedBodyResolveTransformerForIDE( session: FirSession, scopeSession: ScopeSession, implicitTypeOnly: Boolean, - private val towerDataContextForStatement: MutableMap? = null + private val towerDataContextCollector: FirTowerDataContextCollector? = null ) : FirBodyResolveTransformer( session, phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE, @@ -42,7 +41,6 @@ internal class FirDesignatedBodyResolveTransformerForIDE( } override fun onBeforeStatementResolution(statement: FirStatement) { - if (towerDataContextForStatement == null) return - towerDataContextForStatement[statement] = context.towerDataContext + towerDataContextCollector?.addStatementContext(statement, context.towerDataContext) } } diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt new file mode 100644 index 00000000000..af37a3aedd2 --- /dev/null +++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt @@ -0,0 +1,30 @@ +/* + * 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.element.builder + +import org.jetbrains.kotlin.fir.FirElement +import org.jetbrains.kotlin.fir.expressions.FirStatement +import org.jetbrains.kotlin.fir.psi +import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext + +internal class FirTowerDataContextCollector { + private val state: MutableMap = mutableMapOf() + + fun addStatementContext(statement: FirStatement, context: FirTowerDataContext) { + state[statement] = context + } + + fun getContext(statement: FirElement): FirTowerDataContext? { + this.state[statement]?.let { return it } + for ((key, value) in this.state) { + // TODO Rework that + if (key.psi == statement.psi) { + return value + } + } + return null + } +} \ No newline at end of file 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 8645a85fff3..46b5ad33a6e 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 @@ -6,13 +6,12 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve import org.jetbrains.kotlin.fir.declarations.* -import org.jetbrains.kotlin.fir.expressions.FirStatement import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.render -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.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingDeclarationWithFqName import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder @@ -32,7 +31,7 @@ internal class FirLazyDeclarationResolver( declaration: FirDeclaration, moduleFileCache: ModuleFileCache, toPhase: FirResolvePhase, - towerDataContextForStatement: MutableMap? = null, + towerDataContextCollector: FirTowerDataContextCollector? = null, checkPCE: Boolean = false ) { if (declaration.resolvePhase >= toPhase) return @@ -41,12 +40,12 @@ internal class FirLazyDeclarationResolver( val provider = firFile.session.firIdeProvider if (checkPCE) { firFileBuilder.runCustomResolveWithPCECheck(firFile, moduleFileCache) { - runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextForStatement, checkPCE = true) + runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextCollector, checkPCE = true) } } else { firFileBuilder.runCustomResolveUnderLock(firFile, moduleFileCache) { executeWithoutPCE { - runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextForStatement, checkPCE = false) + runLazyResolveWithoutLock(declaration, firFile, provider, toPhase, towerDataContextCollector, checkPCE = false) } } } @@ -57,7 +56,7 @@ internal class FirLazyDeclarationResolver( containerFirFile: FirFile, provider: FirProvider, toPhase: FirResolvePhase, - towerDataContextForStatement: MutableMap? = null, + towerDataContextCollector: FirTowerDataContextCollector? = null, checkPCE: Boolean ) { val nonLazyPhase = minOf(toPhase, FirResolvePhase.DECLARATIONS) @@ -71,7 +70,7 @@ internal class FirLazyDeclarationResolver( } if (toPhase <= nonLazyPhase) return if (checkPCE) checkCanceled() - runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextForStatement) + runLazyResolvePhase(firDeclarationToResolve, containerFirFile, provider, toPhase, towerDataContextCollector) } private fun runLazyResolvePhase( @@ -79,7 +78,7 @@ internal class FirLazyDeclarationResolver( containerFirFile: FirFile, provider: FirProvider, toPhase: FirResolvePhase, - towerDataContextForStatement: MutableMap?, + towerDataContextCollector: FirTowerDataContextCollector?, ) { val nonLocalDeclarationToResolve = firDeclarationToResolve.getNonLocalDeclarationToResolve(provider) @@ -110,7 +109,7 @@ internal class FirLazyDeclarationResolver( designation.iterator(), containerFirFile.session, scopeSession, implicitTypeOnly = toPhase == FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE, - towerDataContextForStatement + towerDataContextCollector ) containerFirFile.transform(transformer, ResolutionMode.ContextDependent) }