From d21437e29dfc18ff476c6916ce71e4af027337d1 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Thu, 19 Mar 2020 12:02:34 +0300
Subject: [PATCH] FIR: Prepare body transformers for local classes resolution

- Add outerBodyResolveContext
- Rename FirImplicitBodyResolveTransformer and generalize because it will be used for full body
resolution on local classes
---
 .../body/resolve/FirBodyResolveTransformer.kt |  5 ++-
 .../body/resolve/FirImplicitBodyResolve.kt    | 38 +++++++++++++------
 2 files changed, 29 insertions(+), 14 deletions(-)

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 237cb3d1bd6..9669f06661a 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
@@ -26,11 +26,12 @@ open class FirBodyResolveTransformer(
     phase: FirResolvePhase,
     override var implicitTypeOnly: Boolean,
     scopeSession: ScopeSession,
-    val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve()
+    val returnTypeCalculator: ReturnTypeCalculator = ReturnTypeCalculatorForFullBodyResolve(),
+    outerBodyResolveContext: BodyResolveContext? = null
 ) : FirAbstractBodyResolveTransformer(phase) {
     private var packageFqName = FqName.ROOT
 
-    final override val context: BodyResolveContext = BodyResolveContext(returnTypeCalculator)
+    final override val context: BodyResolveContext = outerBodyResolveContext ?: BodyResolveContext(returnTypeCalculator)
     final override val components: BodyResolveTransformerComponents =
         BodyResolveTransformerComponents(session, scopeSession, this, context)
 
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt
index a8c6e92e887..08fcb0c5821 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirImplicitBodyResolve.kt
@@ -38,11 +38,14 @@ class FirImplicitTypeBodyResolveTransformerAdapter(private val scopeSession: Sco
 
     override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirFile> {
         val session = file.session
-        scopeSession.returnTypeCalculator = ReturnTypeCalculatorWithJump(session, scopeSession, implicitBodyResolveComputationSession)
-        val transformer = FirImplicitBodyResolveTransformer(
+        val returnTypeCalculator = ReturnTypeCalculatorWithJump(session, scopeSession, implicitBodyResolveComputationSession)
+        scopeSession.returnTypeCalculator = returnTypeCalculator
+        val transformer = FirImplicitAwareBodyResolveTransformer(
             session,
             scopeSession,
-            implicitBodyResolveComputationSession
+            implicitBodyResolveComputationSession,
+            FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE, implicitTypeOnly = true,
+            returnTypeCalculator
         )
         return file.transform(transformer, ResolutionMode.ContextIndependent)
     }
@@ -51,16 +54,21 @@ class FirImplicitTypeBodyResolveTransformerAdapter(private val scopeSession: Sco
 fun createReturnTypeCalculatorForIDE(session: FirSession, scopeSession: ScopeSession): ReturnTypeCalculator =
     ReturnTypeCalculatorWithJump(session, scopeSession, ImplicitBodyResolveComputationSession())
 
-private open class FirImplicitBodyResolveTransformer(
+private open class FirImplicitAwareBodyResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession,
-    private val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession
+    private val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+    phase: FirResolvePhase,
+    implicitTypeOnly: Boolean,
+    returnTypeCalculator: ReturnTypeCalculator,
+    outerBodyResolveContext: BodyResolveContext? = null
 ) : FirBodyResolveTransformer(
     session,
-    phase = FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
-    implicitTypeOnly = true,
-    scopeSession = scopeSession,
-    returnTypeCalculator = ReturnTypeCalculatorWithJump(session, scopeSession, implicitBodyResolveComputationSession)
+    phase,
+    implicitTypeOnly,
+    scopeSession,
+    returnTypeCalculator,
+    outerBodyResolveContext
 ) {
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,
@@ -179,11 +187,17 @@ private class FirDesignatedBodyResolveTransformerForReturnTypeCalculator(
     private val designation: Iterator<FirElement>,
     session: FirSession,
     scopeSession: ScopeSession,
-    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession
-) : FirImplicitBodyResolveTransformer(
+    implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
+    returnTypeCalculator: ReturnTypeCalculator,
+    outerBodyResolveContext: BodyResolveContext? = null
+) : FirImplicitAwareBodyResolveTransformer(
     session,
     scopeSession,
-    implicitBodyResolveComputationSession
+    implicitBodyResolveComputationSession,
+    FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE,
+    implicitTypeOnly = true,
+    returnTypeCalculator,
+    outerBodyResolveContext
 ) {
 
     var lastResult: FirElement? = null