FIR: Get rid of replacements map during calls completion

Otherwise, it's complicated to work with nested lambdas:
outer/inner ones may be replaced independently and still refer
to the old one instances.
So in the changed test, if we don't apply the commit, will remain
implicit return and receiver types for the nested lambda

Anyway, once we decided to leave immutable semantics,
replacements are not necessary anymore

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/PostponedArgumentsAnalyzer.kt

@@ -40,7 +40,6 @@ class PostponedArgumentsAnalyzer(
     private val lambdaAnalyzer: LambdaAnalyzer,
     private val components: InferenceComponents,
     private val candidate: Candidate,
-    private val replacements: MutableMap<FirExpression, FirExpression>,
     private val callResolver: FirCallResolver
 ) {
 
@@ -94,8 +93,6 @@ class PostponedArgumentsAnalyzer(
                 replaceTypeRef(FirResolvedTypeRefImpl(null, candidate.resultingTypeForCallableReference!!))
             }
         }
-
-        replacements[callableReferenceAccess] = transformedCalleeReference
     }
 
     private fun analyzeLambda(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/FirCallCompleter.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.inference
 
-import org.jetbrains.kotlin.fir.copy
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
 import org.jetbrains.kotlin.fir.expressions.FirExpression
@@ -16,8 +15,8 @@ import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.calls.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
-import org.jetbrains.kotlin.fir.resolve.transformers.*
-import org.jetbrains.kotlin.fir.resolve.transformers.MapArguments
+import org.jetbrains.kotlin.fir.resolve.transformers.FirCallCompletionResultsWriterTransformer
+import org.jetbrains.kotlin.fir.resolve.transformers.InvocationKindTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
@@ -75,11 +74,10 @@ class FirCallCompleter(
         }
 
         val completionMode = candidate.computeCompletionMode(inferenceComponents, expectedTypeRef, initialType)
-        val replacements = mutableMapOf<FirExpression, FirExpression>()
 
         val analyzer =
             PostponedArgumentsAnalyzer(
-                LambdaAnalyzerImpl(replacements), inferenceComponents, candidate, replacements,
+                LambdaAnalyzerImpl(), inferenceComponents, candidate,
                 transformer.components.callResolver
             )
 
@@ -88,8 +86,6 @@ class FirCallCompleter(
             analyzer.analyze(candidate.system.asPostponedArgumentsAnalyzerContext(), it)
         }
 
-        call.transformChildren(MapArguments, replacements.toMap())
-
         if (completionMode == KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL) {
             val finalSubstitutor =
                 candidate.system.asReadOnlyStorage().buildAbstractResultingSubstitutor(inferenceComponents.ctx) as ConeSubstitutor
@@ -106,9 +102,7 @@ class FirCallCompleter(
         return call.transformSingle(integerOperatorsTypeUpdater, null)
     }
 
-    private inner class LambdaAnalyzerImpl(
-        val replacements: MutableMap<FirExpression, FirExpression>
-    ) : LambdaAnalyzer {
+    private inner class LambdaAnalyzerImpl : LambdaAnalyzer {
         override fun analyzeAndGetLambdaReturnArguments(
             lambdaArgument: FirAnonymousFunction,
             isSuspend: Boolean,
@@ -142,24 +136,24 @@ class FirCallCompleter(
 
             val expectedReturnTypeRef = expectedReturnType?.let { lambdaArgument.returnTypeRef.resolvedTypeFromPrototype(it) }
 
-            val newLambdaExpression = lambdaArgument.copy(
-                receiverTypeRef = receiverType?.let {
-                    lambdaArgument.receiverTypeRef?.resolvedTypeFromPrototype(it.approximateLambdaInputType())
-                },
-                valueParameters = lambdaArgument.valueParameters.mapIndexed { index, parameter ->
-                    parameter.transformReturnTypeRef(
-                        StoreType,
-                        parameter.returnTypeRef.resolvedTypeFromPrototype(parameters[index].approximateLambdaInputType())
-                    )
-                    parameter
-                } + listOfNotNull(itParam),
-                returnTypeRef = expectedReturnTypeRef ?: noExpectedType
+            lambdaArgument.replaceReceiverTypeRef(
+                receiverType?.approximateLambdaInputType()?.let {
+                    lambdaArgument.receiverTypeRef?.resolvedTypeFromPrototype(it)
+                }
             )
 
-            replacements[lambdaArgument] =
-                newLambdaExpression.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+            lambdaArgument.valueParameters.forEachIndexed { index, parameter ->
+                parameter.replaceReturnTypeRef(
+                    parameter.returnTypeRef.resolvedTypeFromPrototype(parameters[index].approximateLambdaInputType())
+                )
+            }
 
-            val returnArguments = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(newLambdaExpression)
+            lambdaArgument.replaceValueParameters(lambdaArgument.valueParameters + listOfNotNull(itParam))
+            lambdaArgument.replaceReturnTypeRef(expectedReturnTypeRef ?: noExpectedType)
+
+            lambdaArgument.transformSingle(transformer, ResolutionMode.LambdaResolution(expectedReturnTypeRef))
+
+            val returnArguments = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambdaArgument)
             return returnArguments to InferenceSession.default
         }
     }

compiler/fir/resolve/testData/resolve/cfg/postponedLambdas.txt

@@ -2,8 +2,8 @@ FILE: postponedLambdas.kt
     public final inline fun foo(vararg x: R|kotlin/Array<kotlin/Any>|): R|kotlin/Unit| {
     }
     public final fun test(a: R|kotlin/Any|, b: R|kotlin/Any|, c: R|kotlin/Any|): R|kotlin/Unit| {
-        R|/foo|(vararg(R|<local>/a|, R|<local>/b|), foo@fun <anonymous>(): R|kotlin/String| <kind=UNKNOWN>  {
+        R|/foo|(vararg(R|<local>/a|, foo@fun <anonymous>(): R|kotlin/String| <kind=UNKNOWN>  {
             String()
         }
-        )
+        , R|<local>/b|))
     }

compiler/fir/resolve/testData/resolve/inference/nestedLambdas.kt

@@ -8,7 +8,7 @@ interface MyMap<K, V> {
 
 val w: Inv<String> = TODO()
 
-public fun <X, K> Inv<X>.associateBy1(keySelector: (X) -> K): MyMap<K, String> = TODO()
+public fun <X, K> Inv<X>.associateBy1(keySelector: (X) -> K): MyMap<K, X> = TODO()
 
 val x = myRun {
     w.associateBy1 { f -> f.length }

compiler/fir/resolve/testData/resolve/inference/nestedLambdas.txt

@@ -14,7 +14,7 @@ FILE: nestedLambdas.kt
     }
     public final val w: R|Inv<kotlin/String>| = R|kotlin/TODO|()
         public get(): R|Inv<kotlin/String>|
-    public final fun <X, K> R|Inv<X>|.associateBy1(keySelector: R|(X) -> K|): R|MyMap<K, kotlin/String>| {
+    public final fun <X, K> R|Inv<X>|.associateBy1(keySelector: R|(X) -> K|): R|MyMap<K, X>| {
         ^associateBy1 R|kotlin/TODO|()
     }
     public final val x: R|MyMap<kotlin/Int, kotlin/String>| = R|/myRun|<R|MyMap<kotlin/Int, kotlin/String>|>(<L> = myRun@fun <anonymous>(): R|MyMap<kotlin/Int, kotlin/String>| {