From b8030ec1de53546e07b8ce8385347042df1d706c Mon Sep 17 00:00:00 2001
From: Victor Petukhov <victor.petukhov@jetbrains.com>
Date: Sat, 7 May 2022 22:42:06 +0200
Subject: [PATCH] [FIR] Substitute lambda implicit receivers with builder
 inference stub types

^KT-51988 Fixed
---
 .../resolve/inference/FirBuilderInferenceSession.kt   | 11 +++++++++++
 .../transformers/body/resolve/BodyResolveContext.kt   |  1 +
 .../codegen/box/inference/builderInference/kt51988.kt |  1 -
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirBuilderInferenceSession.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirBuilderInferenceSession.kt
index f15d2f95bd8..e69a0552eee 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirBuilderInferenceSession.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirBuilderInferenceSession.kt
@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirResolvable
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.calls.Candidate
+import org.jetbrains.kotlin.fir.resolve.calls.ImplicitExtensionReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.substitution.ChainedSubstitutor
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
@@ -35,6 +36,7 @@ class FirBuilderInferenceSession(
     private val stubsForPostponedVariables: Map<ConeTypeVariable, ConeStubType>,
 ) : FirInferenceSessionForChainedResolve(resolutionContext) {
     private val commonCalls: MutableList<Pair<FirStatement, Candidate>> = mutableListOf()
+    private var lambdaImplicitReceivers: MutableList<ImplicitExtensionReceiverValue> = mutableListOf()
 
     override val currentConstraintStorage: ConstraintStorage
         get() = ConstraintStorage.Empty
@@ -93,6 +95,10 @@ class FirBuilderInferenceSession(
         return postponedAtoms.any { !it.analyzed }
     }
 
+    fun addLambdaImplicitReceiver(receiver: ImplicitExtensionReceiverValue) {
+        lambdaImplicitReceivers += receiver
+    }
+
     override fun <T> addCompletedCall(call: T, candidate: Candidate) where T : FirResolvable, T : FirStatement {
         if (skipCall(call)) return
         commonCalls += call to candidate
@@ -233,6 +239,11 @@ class FirBuilderInferenceSession(
     private fun updateCalls(substitutor: ConeSubstitutor) {
         val stubTypeSubstitutor = FirStubTypeTransformer(substitutor)
         lambda.transformSingle(stubTypeSubstitutor, null)
+
+        for (receiver in lambdaImplicitReceivers) {
+            receiver.replaceType(substitutor.substituteOrSelf(receiver.type))
+        }
+
         // TODO: support diagnostics, see [CoroutineInferenceSession#updateCalls]
 
         val completionResultsWriter = components.callCompleter.createCompletionResultsWriter(substitutor)
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
index 8a03082c040..08dd6f12c80 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
@@ -235,6 +235,7 @@ class BodyResolveContext(
                 holder.scopeSession
             )
             addReceiver(labelName, receiver, additionalLabelName)
+            (inferenceSession as? FirBuilderInferenceSession)?.addLambdaImplicitReceiver(receiver)
         }
 
         f()
diff --git a/compiler/testData/codegen/box/inference/builderInference/kt51988.kt b/compiler/testData/codegen/box/inference/builderInference/kt51988.kt
index 8573e35c0a6..1d37a53f45e 100644
--- a/compiler/testData/codegen/box/inference/builderInference/kt51988.kt
+++ b/compiler/testData/codegen/box/inference/builderInference/kt51988.kt
@@ -1,5 +1,4 @@
 // WITH_STDLIB
-// IGNORE_BACKEND_FIR: JVM_IR
 
 import kotlin.experimental.ExperimentalTypeInference