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 6b1c2895437..266e07b2d78 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 @@ -85,6 +85,8 @@ open class FirBodyResolveTransformer( open fun onBeforeStatementResolution(statement: FirStatement) {} + open fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) {} + // ------------------------------------- Expressions ------------------------------------- override fun transformExpression(expression: FirExpression, data: ResolutionMode): CompositeTransformResult { diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt index e14ada4a702..30579ab60cc 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt +++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt @@ -53,6 +53,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor private fun transformDeclarationContent( declaration: FirDeclaration, data: ResolutionMode ): CompositeTransformResult { + transformer.onBeforeDeclarationContentResolve(declaration) return context.withContainer(declaration) { declaration.replaceResolvePhase(transformerPhase) transformer.transformDeclarationContent(declaration, data) 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 749714dd152..a7e67b8f24f 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 @@ -31,6 +31,10 @@ internal class FirDesignatedBodyResolveTransformerForIDE( returnTypeCalculator = createReturnTypeCalculatorForIDE(session, scopeSession) ) { + override fun onBeforeDeclarationContentResolve(declaration: FirDeclaration) { + towerDataContextCollector?.addDeclarationContext(declaration, context.towerDataContext) + } + override fun transformDeclarationContent(declaration: FirDeclaration, data: ResolutionMode): CompositeTransformResult { if (designation.hasNext()) { designation.next().visitNoTransform(this, data) 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 index af37a3aedd2..5cb66359b75 100644 --- 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 @@ -6,6 +6,7 @@ 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.expressions.FirStatement import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext @@ -17,6 +18,10 @@ internal class FirTowerDataContextCollector { state[statement] = context } + fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext) { + state[declaration] = context + } + fun getContext(statement: FirElement): FirTowerDataContext? { this.state[statement]?.let { return it } for ((key, value) in this.state) {