From 6ad396f417bf8b3b84034be470ab316eca5a267a Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Fri, 18 Dec 2020 13:06:31 +0100
Subject: [PATCH] FIR: use transformer creator from existing return type
 calculator when creating a new one

---
 .../transformers/body/resolve/FirImplicitBodyResolve.kt   | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

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 725bf09fb9e..caccfb7499b 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
@@ -73,6 +73,7 @@ fun <F : FirClass<F>> F.runContractAndBodiesResolutionForLocalClass(
         components.scopeSession,
         implicitBodyResolveComputationSession,
         designationMap,
+        createTransformer = components.returnTypeCalculator.getTransformerCreator()
     )
 
     val newContext = components.context.createSnapshotForLocalClasses(returnTypeCalculator, targetedClasses)
@@ -98,6 +99,11 @@ fun <F : FirClass<F>> F.runContractAndBodiesResolutionForLocalClass(
     }
 }
 
+private fun ReturnTypeCalculator.getTransformerCreator() = when (this) {
+    is ReturnTypeCalculatorWithJump -> createTransformer
+    else -> ::FirDesignatedBodyResolveTransformerForReturnTypeCalculator
+}
+
 fun createReturnTypeCalculatorForIDE(
     session: FirSession,
     scopeSession: ScopeSession,
@@ -178,7 +184,7 @@ private class ReturnTypeCalculatorWithJump(
     private val scopeSession: ScopeSession,
     val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession,
     val designationMapForLocalClasses: Map<FirCallableMemberDeclaration<*>, List<FirClass<*>>> = mapOf(),
-    private val createTransformer: (
+    val createTransformer: (
         designation: Iterator<FirElement>,
         session: FirSession,
         scopeSession: ScopeSession,