FIR IDE: Add collecting scopes for declarations

- They are required to perform completion in default arguments,
return types, and other declaration parts (except for body)

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<FirStatement> {

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<FirDeclaration> {
+        transformer.onBeforeDeclarationContentResolve(declaration)
         return context.withContainer(declaration) {
             declaration.replaceResolvePhase(transformerPhase)
             transformer.transformDeclarationContent(declaration, data)

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<FirDeclaration> {
         if (designation.hasNext()) {
             designation.next().visitNoTransform(this, data)

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) {