[FIR] Substitute lambda implicit receivers with builder inference stub types

^KT-51988 Fixed

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)

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

compiler/testData/codegen/box/inference/builderInference/kt51988.kt

@@ -1,5 +1,4 @@
 // WITH_STDLIB
-// IGNORE_BACKEND_FIR: JVM_IR
 
 import kotlin.experimental.ExperimentalTypeInference