From 90c89bddc5478490ec77a06f53afa9001aab1179 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Mon, 31 Aug 2020 16:59:06 +0300
Subject: [PATCH] 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)
---
 .../transformers/body/resolve/FirBodyResolveTransformer.kt   | 2 ++
 .../body/resolve/FirDeclarationsResolveTransformer.kt        | 1 +
 .../builder/FirDesignatedBodyResolveTransformerForIDE.kt     | 4 ++++
 .../api/element/builder/FirTowerDataContextCollector.kt      | 5 +++++
 4 files changed, 12 insertions(+)

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<FirStatement> {
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<FirDeclaration> {
+        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<FirDeclaration> {
         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) {